Fix 10-bit handling in py_stats.
diff --git a/tools/py_stats/avm_stats/extract_proto.py b/tools/py_stats/avm_stats/extract_proto.py index 0a70284..3b57bf8 100644 --- a/tools/py_stats/avm_stats/extract_proto.py +++ b/tools/py_stats/avm_stats/extract_proto.py
@@ -42,7 +42,8 @@ output_path: pathlib.Path, skip_if_output_already_exists: bool = False, yuv_path: pathlib.Path | None = None, - frame_limit: int | None = None + frame_limit: int | None = None, + extra_args: list[str] | None = None ) -> ExtractProtoResult: if ( skip_if_output_already_exists @@ -75,6 +76,8 @@ "--limit", str(frame_limit), ]) + if extra_args is not None: + extract_proto_args.extend(extra_args) logging.info("Running:\n %s", " ".join(extract_proto_args)) p = subprocess.run(extract_proto_args, capture_output=True, check=True) @@ -107,7 +110,8 @@ output_path: pathlib.Path, skip_if_output_already_exists: bool = False, yuv_path: pathlib.Path | None = None, - frame_limit: int | None = None + frame_limit: int | None = None, + extra_args: list[str] | None = None ) -> Iterator[proto_helpers.Frame]: result = extract_proto( extract_proto_path=extract_proto_path, @@ -115,6 +119,7 @@ output_path=output_path, skip_if_output_already_exists=skip_if_output_already_exists, yuv_path=yuv_path, - frame_limit=frame_limit + frame_limit=frame_limit, + extra_args=extra_args, ) yield from load_protos(result.output_path)
diff --git a/tools/py_stats/avm_stats/frame_visualizations.py b/tools/py_stats/avm_stats/frame_visualizations.py index 2dfbfd1..ce42e23 100644 --- a/tools/py_stats/avm_stats/frame_visualizations.py +++ b/tools/py_stats/avm_stats/frame_visualizations.py
@@ -108,8 +108,9 @@ pixels = getattr(frame.pixels[self.plane], self.pixels_attribute) width = frame.width height = frame.height + vmax = 2 ** frame.bit_depth - 1 ax.imshow( - pixels, cmap="gray", vmin=0, vmax=255, extent=[0, width, height, 0] + pixels, cmap="gray", vmin=0, vmax=vmax, extent=[0, width, height, 0] ) @@ -160,7 +161,8 @@ pixels = self.get_pixels(frame) pixel_min = np.min(pixels) pixel_max = np.max(pixels) - vmin, vmax = -255, 255 + vmax = 2 ** frame.bit_depth - 1 + vmin = -vmax annotation = "Relative" if self.show_relative else "Absolute" if self.show_relative: vmin, vmax = pixel_min, pixel_max
diff --git a/tools/py_stats/avm_stats/proto_helpers.py b/tools/py_stats/avm_stats/proto_helpers.py index b6c861a..0bd57a5 100644 --- a/tools/py_stats/avm_stats/proto_helpers.py +++ b/tools/py_stats/avm_stats/proto_helpers.py
@@ -160,9 +160,13 @@ pixels_width = superblock_plane.width pixels_height = superblock_plane.height - superblock_samples = np.array(superblock_plane.pixels).reshape( + superblock_samples = np.array(superblock_plane.pixels, dtype=dtype).reshape( (pixels_height, pixels_width) )[:sb_height_clipped, :sb_width_clipped] + if superblock_plane.bit_depth > frame.bit_depth: + superblock_samples //= int(2 ** (superblock_plane.bit_depth - frame.bit_depth)) + elif superblock_plane.bit_depth < frame.bit_depth: + superblock_samples *= int(2 ** (frame.bit_depth - superblock_plane.bit_depth)) samples[sb_y: sb_y + sb_height_clipped, sb_x: sb_x + sb_width_clipped] = ( superblock_samples @@ -194,11 +198,11 @@ assert reconstruction is not None # Even for 8-bit frames, 16 bits are needed for the deltas, since the range is # [-255, 255]. - residual = pre_filtered.astype(np.int16) - prediction - filter_delta = reconstruction.astype(np.int16) - pre_filtered + residual = pre_filtered.astype(np.int16) - prediction.astype(np.int16) + filter_delta = reconstruction.astype(np.int16) - pre_filtered.astype(np.int16) # Since distortion is computed from the original, this might also be None. distortion = ( - original.astype(np.int16) - reconstruction + original.astype(np.int16) - reconstruction.astype(np.int16) if original is not None else None ) @@ -563,13 +567,22 @@ return self.proto.frame_params.bit_depth @property + def pixel_scale(self) -> int: + if self.bit_depth == 8: + return 1 + elif self.bit_depth == 10: + return 4 + else: + raise RuntimeError(f"Unsupported bit depth: {self.bit_depth}") + + @property def is_intra_frame(self) -> bool: return self.proto.frame_params.frame_type == 0 @property def pixels(self) -> list[PlaneBuffer]: if self._pixels is None: - self.pixels = [ + self._pixels = [ _create_plane_buffer(self, p) for p in (Plane.Y, Plane.U, Plane.V) ] return self._pixels @@ -578,9 +591,9 @@ def original_rgb(self) -> np.ndarray: if self._original_rgb is None and self.pixels[0].original is not None: self._original_rgb = yuv_tools.yuv_to_rgb( - self.pixels[0].original, - yuv_tools.upscale(self.pixels[1].original, 2), - yuv_tools.upscale(self.pixels[2].original, 2), + self.pixels[0].original // self.pixel_scale, + yuv_tools.upscale(self.pixels[1].original, 2) // self.pixel_scale, + yuv_tools.upscale(self.pixels[2].original, 2) // self.pixel_scale, ) return self._original_rgb @@ -588,9 +601,11 @@ def reconstruction_rgb(self) -> np.ndarray: if self._reconstruction_rgb is None: self._reconstruction_rgb = yuv_tools.yuv_to_rgb( - self.pixels[0].reconstruction, - yuv_tools.upscale(self.pixels[1].reconstruction, 2), - yuv_tools.upscale(self.pixels[2].reconstruction, 2), + self.pixels[0].reconstruction // self.pixel_scale, + yuv_tools.upscale( + self.pixels[1].reconstruction, 2) // self.pixel_scale, + yuv_tools.upscale( + self.pixels[2].reconstruction, 2) // self.pixel_scale, ) return self._reconstruction_rgb
diff --git a/tools/py_stats/examples/aggregate_from_extractor.py b/tools/py_stats/examples/aggregate_from_extractor.py index 2ae1747..f99e531 100755 --- a/tools/py_stats/examples/aggregate_from_extractor.py +++ b/tools/py_stats/examples/aggregate_from_extractor.py
@@ -29,14 +29,12 @@ import multiprocessing import pandas as pd -from extractors.partition_similarity_extractor import PartitionSimilarityExtractor from extractors.partition_type_extractor import PartitionTypeExtractor from extractors.prediction_mode_extractor import PredictionModeExtractor from extractors.symbol_bits_extractor import SymbolBitsExtractor from extractors.tx_type_extractor import TxTypeExtractor _EXTRACTORS = { - "partition_similarity": PartitionSimilarityExtractor, "partition_type": PartitionTypeExtractor, "prediction_mode": PredictionModeExtractor, "symbol_bits": SymbolBitsExtractor, @@ -163,7 +161,9 @@ extract_proto_path=extract_proto_path, stream_path=stream_path, output_path=output_path, - frame_limit=frame_limit + frame_limit=frame_limit, + # Preserve the entire stream path so that the CTC config and class can be retrieved if available from the path. + extra_args = ["--preserve_stream_path_depth", "-1"] )
diff --git a/tools/py_stats/examples/extractors/partition_type_extractor.py b/tools/py_stats/examples/extractors/partition_type_extractor.py index c972d5f..bd7f8c76 100644 --- a/tools/py_stats/examples/extractors/partition_type_extractor.py +++ b/tools/py_stats/examples/extractors/partition_type_extractor.py
@@ -24,11 +24,11 @@ class PartitionTypeExtractor(SuperblockExtractor): PartitionType = collections.namedtuple( "PartitionType", - ["width", "height", "block_size", "partition_type", "is_intra_frame", "stream_name"], + ["width", "height", "block_size", "partition_type", "is_intra_frame", "stream_path"], ) def sample(self, superblock: Superblock): - stream_name = superblock.frame.proto.stream_params.stream_name.removesuffix(".bin") + stream_path = superblock.frame.proto.stream_params.stream_path is_intra_frame = superblock.frame.is_intra_frame for partition_block in iter_partitions(superblock.proto.luma_partition_tree): width = partition_block.size.width @@ -36,4 +36,4 @@ block_size = f"{width}x{height}" partition_type = partition_block.partition_type partition_type_name = superblock.frame.proto.enum_mappings.partition_type_mapping[partition_type] - yield self.PartitionType(width, height, block_size, partition_type_name, is_intra_frame, stream_name) + yield self.PartitionType(width, height, block_size, partition_type_name, is_intra_frame, stream_path)
diff --git a/tools/py_stats/examples/extractors/prediction_mode_extractor.py b/tools/py_stats/examples/extractors/prediction_mode_extractor.py index 6333989..43eec36 100644 --- a/tools/py_stats/examples/extractors/prediction_mode_extractor.py +++ b/tools/py_stats/examples/extractors/prediction_mode_extractor.py
@@ -26,11 +26,11 @@ class PredictionModeExtractor(CodingUnitExtractor): PredictionMode = collections.namedtuple( "PredictionMode", - ["width", "height", "block_size", "mode", "mode_bits", "is_intra_frame", "stream_name"], + ["width", "height", "block_size", "mode", "mode_bits", "is_intra_frame", "stream_path"], ) def sample(self, coding_unit: CodingUnit): - stream_name = coding_unit.frame.proto.stream_params.stream_name.removesuffix(".bin") + stream_path = coding_unit.frame.proto.stream_params.stream_path width = coding_unit.rect.width height = coding_unit.rect.height block_size = f"{width}x{height}" @@ -40,4 +40,4 @@ for sym in coding_unit.get_symbols(prediction_mode_symbol_filter) ) is_intra_frame = coding_unit.frame.is_intra_frame - yield self.PredictionMode(width, height, block_size, mode, mode_bits, is_intra_frame, stream_name) + yield self.PredictionMode(width, height, block_size, mode, mode_bits, is_intra_frame, stream_path)
diff --git a/tools/py_stats/examples/extractors/symbol_bits_extractor.py b/tools/py_stats/examples/extractors/symbol_bits_extractor.py index 2a70bd5..b889f6f 100644 --- a/tools/py_stats/examples/extractors/symbol_bits_extractor.py +++ b/tools/py_stats/examples/extractors/symbol_bits_extractor.py
@@ -15,32 +15,16 @@ from avm_stats.proto_helpers import * from avm_stats.stats_aggregation import * -CTC_CLASSES = [ - 2160, - 1080, - 720, - 360, - 270, -] - class SymbolBitsExtractor(SuperblockExtractor): SymbolBits = collections.namedtuple( "SymbolBits", - ["symbol_name", "symbol_tags", "bits", "is_intra_frame", "stream_name", "qp", "ctc_class", "ctc_config"], + ["symbol_name", "symbol_tags", "bits", "is_intra_frame", "stream_path"], ) def sample(self, superblock: Superblock): - stream_name = superblock.frame.proto.stream_params.stream_name.removesuffix(".bin") - qp = stream_name.split("_")[-1] - ctc_config = stream_name.split("_")[-3] - ctc_class = None - for i, c in enumerate(CTC_CLASSES): - if str(c) in stream_name: - assert ctc_class is None - ctc_class = f"A{i+1}" - assert ctc_class is not None + stream_path = superblock.frame.proto.stream_params.stream_path is_intra_frame = superblock.frame.is_intra_frame for symbol in superblock.proto.symbols: symbol_info = superblock.frame.proto.symbol_info[symbol.info_id] symbol_tags = "/".join(symbol_info.tags) - yield self.SymbolBits(symbol_info.source_function, symbol_tags, symbol.bits, is_intra_frame, stream_name, qp, ctc_class, ctc_config) + yield self.SymbolBits(symbol_info.source_function, symbol_tags, symbol.bits, is_intra_frame, stream_path)
diff --git a/tools/py_stats/examples/extractors/tx_type_extractor.py b/tools/py_stats/examples/extractors/tx_type_extractor.py index bf9ae93..a07697d 100644 --- a/tools/py_stats/examples/extractors/tx_type_extractor.py +++ b/tools/py_stats/examples/extractors/tx_type_extractor.py
@@ -19,11 +19,11 @@ class TxTypeExtractor(CodingUnitExtractor): TxSize = collections.namedtuple( "TxSize", - ["width", "height", "tx_size", "tx_type", "is_intra_frame", "stream_name"], + ["width", "height", "tx_size", "tx_type", "is_intra_frame", "stream_path"], ) def sample(self, coding_unit: CodingUnit): - stream_name = coding_unit.frame.proto.stream_params.stream_name.removesuffix(".bin") + stream_path = coding_unit.frame.proto.stream_params.stream_path # Luma only is_chroma = len(coding_unit.proto.transform_planes) == 2 if is_chroma: @@ -36,4 +36,4 @@ # TODO(comc): Add method on transform_unit for this. tx_type = transform_unit.tx_type & 0xF tx_type_name = coding_unit.frame.proto.enum_mappings.transform_type_mapping[tx_type] - yield self.TxSize(width, height, tx_size, tx_type_name, is_intra_frame, stream_name) + yield self.TxSize(width, height, tx_size, tx_type_name, is_intra_frame, stream_path)
diff --git a/tools/py_stats/examples/pixel_pipeline.py b/tools/py_stats/examples/pixel_pipeline.py index 0f06e45..fe26b21 100755 --- a/tools/py_stats/examples/pixel_pipeline.py +++ b/tools/py_stats/examples/pixel_pipeline.py
@@ -57,6 +57,7 @@ output_path=tmp_path, skip_if_output_already_exists=False, yuv_path=yuv_path, + frame_limit=_FRAME.value, ) seq = list(frames) num_frames = len(seq)