Python stats aggregation improvements. Added better examples and updated README.
diff --git a/tools/py_stats/README.md b/tools/py_stats/README.md index aa6ad85..0bccb86 100644 --- a/tools/py_stats/README.md +++ b/tools/py_stats/README.md
@@ -197,30 +197,105 @@ Bits for read_intra_luma_mode: 760.7987823486328 (9.55%) ``` -### Prediction mode / block size aggregation -All previous examples focused on visualizing a single frame from a single stream. This example shows how to compute aggregated stats across an arbitrary number of AVM streams. -A typical use case might be to compare how bits are spent in CTC frames under different test conditions. +### Statistics aggregation +All previous examples focused on visualizing a single frame from a single stream. The following examples show how to compute aggregated stats across an arbitrary number of AVM streams. A typical use case might be to compare how bits are spent in CTC frames under different test conditions. -This example script takes one or more glob paths and will extract all frames from all streams found: +Each of the following examples uses the `aggregate_from_extractor.py` script to plot different types of data from the streams, in addition to dumping that data as a .csv file. + +### Block size distribution (A5, QP=235, 10 frames); Plot Intra vs Inter ```bash -python3 ${AOM_ROOT}/tools/py_stats/examples/aggregate_prediction_modes.py \ - --stream_glob "/path/to/some/interesting/streams/*.ivf" \ - --stream_glob "/path/to/some/other/streams/*.ivf" \ - --extract_proto_bin ${AOM_BUILD_DIR}/extract_proto \ - --output_csv /tmp/aggregate_dump.csv +python3 ${AOM_ROOT}/tools/py_stats/examples/aggregate_from_extractor.py \ + --extract_proto_bin ${AOM_BUILD_DIR}/extract_proto \ + --stream_glob "${CTC_STREAMS}/A5_Natural_270p/LowDelay/*_QP_235.bin" \ + --threads 32 \ + --frame_limit 10 \ + --extractor "partition_type" \ + --group_by block_size,partition_type,is_intra_frame \ + --output_csv /tmp/aggregate_block_sizes.csv \ + --plot 'title:"Intra Block Sizes", field:block_size, filter:"is_intra_frame"' \ + --plot 'title:"Inter Block Sizes", field:block_size, filter:"not is_intra_frame"' ``` + -The script will produce three different plots: -1. The overall distribution of block sizes, comparing intra and inter frames: - -2. The overall distribution of prediction modes, comparing intra and inter frames: - +### Prediction mode distribution (A5, QP=235, 10 frames); Plot Intra vs Inter +```bash +python3 ${AOM_ROOT}/tools/py_stats/examples/aggregate_from_extractor.py \ + --extract_proto_bin ${AOM_BUILD_DIR}/extract_proto \ + --stream_glob "${CTC_STREAMS}/A5_Natural_270p/LowDelay/*_QP_235.bin" \ + --threads 32 \ + --frame_limit 10 \ + --extractor "prediction_mode" \ + --group_by mode,is_intra_frame \ + --output_csv /tmp/prediction_modes.csv \ + --plot 'title:"Intra Prediction Modes", field:mode, filter:"is_intra_frame"' \ + --plot 'title:"Inter Prediction Modes", field:mode, filter:"not is_intra_frame"' +``` + -3. Prediction mode distribution, weighted by the number of bits used to code each mode, comparing intra and inter frames: - -The script also takes an optional `--output_csv` argument that will dump the final Pandas dataframe to a CSV file. +### Prediction modes weighted by bits used (A5, QP=235, 10 frames); Plot Intra vs Inter +```bash +python3 ${AOM_ROOT}/tools/py_stats/examples/aggregate_from_extractor.py \ + --extract_proto_bin ${AOM_BUILD_DIR}/extract_proto \ + --stream_glob "${CTC_STREAMS}/A5_Natural_270p/LowDelay/*_QP_235.bin" \ + --threads 32 \ + --frame_limit 10 \ + --extractor "prediction_mode" \ + --aggregated_field mode_bits \ + --group_by mode,is_intra_frame \ + --output_csv /tmp/prediction_mode_bits.csv \ + --plot 'title:"Intra Prediction Modes by bits", field:mode, filter:"is_intra_frame"' \ + --plot 'title:"Inter Prediction Modes by bits", field:mode, filter:"not is_intra_frame"' +``` + + +### TX types (A5, QP=235, 10 frames); Plot Intra vs Inter +```bash +python3 ${AOM_ROOT}/tools/py_stats/examples/aggregate_from_extractor.py \ + --extract_proto_bin ${AOM_BUILD_DIR}/extract_proto \ + --stream_glob "${CTC_STREAMS}/A5_Natural_270p/LowDelay/*_QP_235.bin" \ + --threads 32 \ + --frame_limit 10 \ + --extractor "tx_type" \ + --group_by tx_type,is_intra_frame \ + --output_csv /tmp/aggregate_tx_types.csv \ + --plot 'title:"Intra TX Types", field:tx_type, filter:"is_intra_frame"' \ + --plot 'title:"Inter TX Types", field:tx_type, filter:"not is_intra_frame"' +``` + + +### Intra partition types (A5, QP=210, 10 frames); Block sizes 64x64 vs 32x32 +```bash +python3 ${AOM_ROOT}/tools/py_stats/examples/aggregate_from_extractor.py \ + --extract_proto_bin ${AOM_BUILD_DIR}/extract_proto \ + --stream_glob "${CTC_STREAMS}/A5_Natural_270p/AllIntra/*_QP_210.bin" \ + --threads 32 \ + --frame_limit 10 \ + --extractor "partition_type" \ + --group_by partition_type,block_size \ + --output_csv /tmp/aggregate_partition_types.csv \ + --plot 'title:"Partition types (64x64)", field:partition_type, filter:block_size=="64x64"' \ + --plot 'title:"Partition types (32x32)", field:partition_type, filter:block_size=="32x32"' +``` + + +### Top 10 symbols types weighted by bits used (A5, QP=235, 10 frames); Intra vs Inter frames +```bash +python3 ${AOM_ROOT}/tools/py_stats/examples/aggregate_from_extractor.py \ + --extract_proto_bin ${AOM_BUILD_DIR}/extract_proto \ + --stream_glob "${CTC_STREAMS}/A5_Natural_270p/LowDelay/*_QP_210.bin" \ + --threads 32 \ + --frame_limit 10 \ + --extractor "symbol_bits" \ + --aggregated_field bits \ + --group_by symbol_name,is_intra_frame \ + --output_csv /tmp/aggregate_symbols.csv \ + --plot 'title:"Symbols (Intra frames)", field:symbol_name, filter:"is_intra_frame", limit:10' \ + --plot 'title:"Symbols (Inter frames)", field:symbol_name, filter:"not is_intra_frame", limit:10' +``` + + ### Launch an example Jupyter notebook This demonstrates many of the same visualizations listed above, but in a single Jupyter notebook.
diff --git a/tools/py_stats/avm_stats/proto_helpers.py b/tools/py_stats/avm_stats/proto_helpers.py index b012ec7..b6c861a 100644 --- a/tools/py_stats/avm_stats/proto_helpers.py +++ b/tools/py_stats/avm_stats/proto_helpers.py
@@ -164,7 +164,7 @@ (pixels_height, pixels_width) )[:sb_height_clipped, :sb_width_clipped] - samples[sb_y : sb_y + sb_height_clipped, sb_x : sb_x + sb_width_clipped] = ( + samples[sb_y: sb_y + sb_height_clipped, sb_x: sb_x + sb_width_clipped] = ( superblock_samples ) @@ -541,23 +541,10 @@ self.superblocks = [ Superblock(self, sb_proto) for sb_proto in proto.superblocks ] - self.pixels = [ - _create_plane_buffer(self, p) for p in (Plane.Y, Plane.U, Plane.V) - ] - - if self.pixels[0].original is None: - self.original_rgb = None - else: - 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.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), - ) + # Pixel data is created lazily. + self._pixels = None + self._original_rgb = None + self._reconstruction_rgb = None @property def frame_id(self) -> int: @@ -575,6 +562,38 @@ def bit_depth(self) -> int: return self.proto.frame_params.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 = [ + _create_plane_buffer(self, p) for p in (Plane.Y, Plane.U, Plane.V) + ] + return self._pixels + + @property + 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), + ) + return self._original_rgb + + @property + 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), + ) + return self._reconstruction_rgb + def clip_rect(self, rect: Rectangle) -> Rectangle: """Clips a rectangle to be contained with the frame boundaries.
diff --git a/tools/py_stats/examples/aggregate_from_extractor.py b/tools/py_stats/examples/aggregate_from_extractor.py new file mode 100755 index 0000000..2ae1747 --- /dev/null +++ b/tools/py_stats/examples/aggregate_from_extractor.py
@@ -0,0 +1,232 @@ +#!/usr/bin/env python3 +## Copyright (c) 2023, Alliance for Open Media. All rights reserved +## +## This source code is subject to the terms of the BSD 3-Clause Clear License and the +## Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was +## not distributed with this source code in the LICENSE file, you can obtain it +## at aomedia.org/license/software-license/bsd-3-c-c/. If the Alliance for Open Media Patent +## License 1.0 was not distributed with this source code in the PATENTS file, you +## can obtain it at aomedia.org/license/patent-license/. +## +from __future__ import annotations + +from functools import partial +import gc +import glob +import pathlib +import sys +import tempfile +from typing import Sequence, Type + +from absl import app +from absl import flags +from absl import logging +from avm_stats.extract_proto import * +from avm_stats.frame_visualizations import * +from avm_stats.proto_helpers import * +from avm_stats.stats_aggregation import * +import matplotlib.pyplot as plt +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, + "tx_type": TxTypeExtractor, +} + +_STREAM_GLOB = flags.DEFINE_multi_string( + "stream_glob", None, "Path to AVM streams." +) +flags.mark_flag_as_required("stream_glob") + +_EXTRACT_PROTO_BIN = flags.DEFINE_string( + "extract_proto_bin", None, "Path to extract_proto binary." +) +flags.mark_flag_as_required("extract_proto_bin") + +_OUTPUT_CSV = flags.DEFINE_string( + "output_csv", None, "Path to output CSV file (optional)." +) + +_AGGREGATED_FIELD = flags.DEFINE_string( + "aggregated_field", "count", "Field to aggregate on. By default will use the count." +) + +_GROUP_BY = flags.DEFINE_string( + "group_by", None, "Group by these fields (comma separated)." +) + +_EXTRACTOR = flags.DEFINE_enum( + "extractor", None, _EXTRACTORS.keys(), "Data extractor to use." +) +flags.mark_flag_as_required("extractor") + +_THREADS = flags.DEFINE_integer( + "threads", 1, "Number of parallel workers to spawn." +) + +_FRAME_LIMIT = flags.DEFINE_integer( + "frame_limit", None, "Use at most this many frames from each stream." +) + +_PLOT = flags.DEFINE_multi_string( + "plot", None, """Plot command args, e.g.: --plot 'title:"Block Sizes", field:"block_size", filter:"not is_intra_frame", limit:10'""" +) + +_TMP_DIR = flags.DEFINE_string( + "tmp_dir", None, "Temp working dir." +) + +@dataclasses.dataclass +class PlotArgs: + title: str = "" + field: str = "" + filter: str = "" + limit: str = "" + + def __init__(self, args_str: str): + def kv_pair(kv: str) -> tuple[str, str]: + k, v = [i.strip() for i in kv.split(":")] + if v.startswith('"') and v.endswith('"'): + v = v[1:-1] + return (k, v) + + for k, v in [kv_pair(kv) for kv in args_str.split(",")]: + if not hasattr(self, k): + raise ValueError(f"Unknown plot arg: {k}") + setattr(self, k, v) + + +def filter_dataframe(df: pd.DataFrame, *, group_by: list[str], aggregated_field: str = "count", filt: str = "", limit: int | None = None): + if filt: + filtered_df = df.query(filt)[group_by + [aggregated_field]].groupby(group_by, as_index=False) + else: + filtered_df = df[group_by + [aggregated_field]].groupby(group_by, as_index=False) + filtered_df = filtered_df.sum() + filtered_df["percent"] = filtered_df[aggregated_field].transform(lambda x: x / x.sum() * 100) + filtered_df = filtered_df.sort_values(by=[aggregated_field], ascending=False) + if limit is not None: + filtered_df_top_n = filtered_df.iloc[:limit, :] + other_total = filtered_df.iloc[limit:, :][aggregated_field].sum() + other_percent = filtered_df.iloc[limit:, :]["percent"].sum() + others = pd.DataFrame({group_by[0]: ["Other"], aggregated_field: other_total, "percent": other_percent}) + filtered_df = pd.concat([filtered_df_top_n, others]) + return filtered_df + + +def sum_dataframe(df: pd.DataFrame, *, group_by: list[str], aggregated_field: str = "count"): + aggregated_df = df[group_by + [aggregated_field]].groupby(group_by, as_index=False) + aggregated_df = aggregated_df.sum() + return aggregated_df + + +def create_plot( + df: pd.DataFrame, + plot_args: PlotArgs, + ax: plt.Axes, + legend_ax: plt.Axes, + aggregated_field: str, +): + limit = int(plot_args.limit) if plot_args.limit else None + aggregated_df = filter_dataframe( + df, group_by=[plot_args.field], aggregated_field=aggregated_field, filt=plot_args.filter, limit=limit) + plot_title = plot_args.title or plot_args.field + ax.set_title(f"{plot_title}") + patches, _ = ax.pie( + x=aggregated_df[aggregated_field], labels=aggregated_df[plot_args.field], autopct=None) + labels = [f"{n} - {p:.1f}% ({v:.1f})" for n, p, v in zip(aggregated_df[plot_args.field], + aggregated_df["percent"], aggregated_df[aggregated_field])] + legend_ax.legend(patches, labels, loc="best", ncol=2) + legend_ax.axis("off") + + +def extract_to_temp_dir(stream_path: pathlib.Path, frame_limit: int | None = None) -> Iterator[proto_helpers.Frame]: + with tempfile.TemporaryDirectory(dir=_TMP_DIR.value) as tmp_dir: + tmp_path = pathlib.Path(tmp_dir) + extract_proto_path = pathlib.Path(_EXTRACT_PROTO_BIN.value) + stream_name = stream_path.stem + output_path = tmp_path / stream_name + try: + output_path.mkdir() + except FileExistsError: + logging.fatal(f"Duplicate stream name: {stream_name}") + yield from extract_and_load_protos( + extract_proto_path=extract_proto_path, + stream_path=stream_path, + output_path=output_path, + frame_limit=frame_limit + ) + + +def process_stream(stream_path: pathlib.Path, *, extractor_class: Type[Extractor]) -> pd.DataFrame: + group_by = _GROUP_BY.value.split(",") + frames = extract_to_temp_dir(stream_path, _FRAME_LIMIT.value) + df = None + for frame in frames: + frame_df = aggregate_to_dataframe([frame], extractor_class()) + frame_df["count"] = 1 + if df is None: + df = frame_df + else: + df = pd.concat([df, frame_df]) + if len(df): + df = sum_dataframe(df, group_by=group_by, aggregated_field=_AGGREGATED_FIELD.value) + gc.collect() + if len(df): + df = sum_dataframe(df, group_by=group_by, aggregated_field=_AGGREGATED_FIELD.value) + gc.collect() + with open(f"/opt/tmp/progress/{stream_path.stem}.txt", "w") as f: + f.write("Done") + return df + + +def main(argv: Sequence[str]) -> None: + if len(argv) > 1: + raise app.UsageError("Too many command-line arguments.") + + extractor_name = _EXTRACTOR.value + try: + extractor_class = _EXTRACTORS[extractor_name] + except KeyError: + print(f"Unknown extractor: {extractor_name}", file=sys.stderr) + sys.exit(1) + + stream_paths = [] + for stream_glob in _STREAM_GLOB.value: + for stream in glob.glob(stream_glob, recursive=True): + path = pathlib.Path(stream) + stream_paths.append(path) + + stream_paths = sorted(stream_paths) + with multiprocessing.Pool(_THREADS.value) as pool: + dfs = pool.map( + partial(process_stream, extractor_class=extractor_class), stream_paths) + + df = pd.concat(dfs) + group_by = _GROUP_BY.value.split(",") + aggregated_df = sum_dataframe(df, group_by=group_by, aggregated_field=_AGGREGATED_FIELD.value) + + if _OUTPUT_CSV.value: + aggregated_df.to_csv(_OUTPUT_CSV.value) + + plots = _PLOT.value + if plots: + _, axes = plt.subplots(2, len(plots), squeeze=False, figsize=(10, 5)) + for i, plot in enumerate(plots): + args = PlotArgs(plot) + create_plot(df, args, axes[0, i], axes[1, i], aggregated_field=_AGGREGATED_FIELD.value) + plt.tight_layout() + plt.show() + + +if __name__ == "__main__": + app.run(main)
diff --git a/tools/py_stats/examples/aggregate_prediction_modes.py b/tools/py_stats/examples/aggregate_prediction_modes.py deleted file mode 100644 index 9d73fd3..0000000 --- a/tools/py_stats/examples/aggregate_prediction_modes.py +++ /dev/null
@@ -1,154 +0,0 @@ -#!/usr/bin/env python3 -## Copyright (c) 2023, Alliance for Open Media. All rights reserved -## -## This source code is subject to the terms of the BSD 3-Clause Clear License and the -## Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was -## not distributed with this source code in the LICENSE file, you can obtain it -## at aomedia.org/license/software-license/bsd-3-c-c/. If the Alliance for Open Media Patent -## License 1.0 was not distributed with this source code in the PATENTS file, you -## can obtain it at aomedia.org/license/patent-license/. -## -from __future__ import annotations - -import collections -import glob -import itertools -import pathlib -import tempfile -from typing import Sequence - -from absl import app -from absl import flags -from absl import logging -from avm_stats.extract_proto import * -from avm_stats.frame_visualizations import * -from avm_stats.proto_helpers import * -from avm_stats.stats_aggregation import * -import matplotlib.pyplot as plt -import pandas as pd - -_STREAM_GLOB = flags.DEFINE_multi_string( - "stream_glob", None, "Path to AVM stream." -) -flags.mark_flag_as_required("stream_glob") - -_EXTRACT_PROTO_BIN = flags.DEFINE_string( - "extract_proto_bin", None, "Path to extract_proto binary." -) -flags.mark_flag_as_required("extract_proto_bin") - -_OUTPUT_CSV = flags.DEFINE_string( - "output_csv", None, "Path to output CSV file (optional)." -) - - -def prediction_mode_symbol_filter(symbol: Symbol): - return symbol.source_function in ( - "read_inter_mode", - "read_intra_luma_mode", - "read_drl_index", - "read_inter_compound_mode", - ) - - -class PredictionModeExtractor(CodingUnitExtractor): - PredictionMode = collections.namedtuple( - "PredictionMode", - ["width", "height", "mode", "mode_bits", "is_intra_frame"], - ) - - def sample(self, coding_unit: CodingUnit): - width = coding_unit.rect.width - height = coding_unit.rect.height - mode = coding_unit.get_prediction_mode() - mode_bits = sum( - sym.bits - for sym in coding_unit.get_symbols(prediction_mode_symbol_filter) - ) - is_intra_frame = coding_unit.frame.proto.frame_params.frame_type == 0 - yield self.PredictionMode(width, height, mode, mode_bits, is_intra_frame) - - -def compare_intra_inter( - df: pd.DataFrame, - column: str, - *, - aggregated_field: str = "count", - aggregator: str = "count", - plot_title: str | None = None, -): - df_intra = df.query("is_intra_frame")[[column, aggregated_field]].groupby( - column, as_index=False - ) - df_intra = getattr(df_intra, aggregator)() - df_inter = df.query("not is_intra_frame")[[column, aggregated_field]].groupby( - column, as_index=False - ) - df_inter = getattr(df_inter, aggregator)() - _, axes = plt.subplots(1, 2) - plot_title = plot_title or column - axes[0].pie( - x=df_intra[aggregated_field], labels=df_intra[column], autopct="%.1f%%" - ) - axes[0].set_title(f"{plot_title} (Intra frames)") - axes[1].pie( - x=df_inter[aggregated_field], labels=df_inter[column], autopct="%.1f%%" - ) - axes[1].set_title(f"{plot_title} (Inter frames)") - plt.show() - - -def main(argv: Sequence[str]) -> None: - if len(argv) > 1: - raise app.UsageError("Too many command-line arguments.") - - stream_paths = [] - for stream_glob in _STREAM_GLOB.value: - for stream in glob.glob(stream_glob, recursive=True): - path = pathlib.Path(stream) - stream_paths.append(path) - - with tempfile.TemporaryDirectory() as tmp_dir: - tmp_path = pathlib.Path(tmp_dir) - extract_proto_path = pathlib.Path(_EXTRACT_PROTO_BIN.value) - - def extract_to_temp_dir(stream_path: pathlib.Path) -> Frame: - stream_name = stream_path.stem - output_path = tmp_path / stream_name - try: - output_path.mkdir() - except FileExistsError: - logging.fatal(f"Duplicate stream name: {stream_name}") - yield from extract_and_load_protos( - extract_proto_path=extract_proto_path, - stream_path=stream_path, - output_path=output_path, - ) - - all_frames = itertools.chain.from_iterable( - map(extract_to_temp_dir, stream_paths) - ) - df = aggregate_to_dataframe(all_frames, PredictionModeExtractor()) - df["count"] = 1 - df["width_height"] = df.apply( - lambda row: f"{row.width}x{row.height}", axis=1 - ) - - compare_intra_inter( - df, "width_height", plot_title="Block size distribution" - ) - compare_intra_inter(df, "mode", plot_title="Prediction mode distribution") - compare_intra_inter( - df, - "mode", - aggregated_field="mode_bits", - aggregator="sum", - plot_title="Prediction modes by bits spent", - ) - - if _OUTPUT_CSV.value: - df.to_csv(_OUTPUT_CSV.value) - - -if __name__ == "__main__": - app.run(main)
diff --git a/tools/py_stats/examples/extractors/partition_type_extractor.py b/tools/py_stats/examples/extractors/partition_type_extractor.py new file mode 100644 index 0000000..c972d5f --- /dev/null +++ b/tools/py_stats/examples/extractors/partition_type_extractor.py
@@ -0,0 +1,39 @@ +#!/usr/bin/env python3 +## Copyright (c) 2023, Alliance for Open Media. All rights reserved +## +## This source code is subject to the terms of the BSD 3-Clause Clear License and the +## Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was +## not distributed with this source code in the LICENSE file, you can obtain it +## at aomedia.org/license/software-license/bsd-3-c-c/. If the Alliance for Open Media Patent +## License 1.0 was not distributed with this source code in the PATENTS file, you +## can obtain it at aomedia.org/license/patent-license/. +## +import collections + +from avm_stats.extract_proto import * +from avm_stats.frame_visualizations import * +from avm_stats.proto_helpers import * +from avm_stats.stats_aggregation import * + +def iter_partitions(partition_block): + yield partition_block + if not partition_block.is_leaf_node: + for child in partition_block.children: + yield from iter_partitions(child) + +class PartitionTypeExtractor(SuperblockExtractor): + PartitionType = collections.namedtuple( + "PartitionType", + ["width", "height", "block_size", "partition_type", "is_intra_frame", "stream_name"], + ) + + def sample(self, superblock: Superblock): + stream_name = superblock.frame.proto.stream_params.stream_name.removesuffix(".bin") + is_intra_frame = superblock.frame.is_intra_frame + for partition_block in iter_partitions(superblock.proto.luma_partition_tree): + width = partition_block.size.width + height = partition_block.size.height + 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)
diff --git a/tools/py_stats/examples/extractors/prediction_mode_extractor.py b/tools/py_stats/examples/extractors/prediction_mode_extractor.py new file mode 100644 index 0000000..6333989 --- /dev/null +++ b/tools/py_stats/examples/extractors/prediction_mode_extractor.py
@@ -0,0 +1,43 @@ +#!/usr/bin/env python3 +## Copyright (c) 2023, Alliance for Open Media. All rights reserved +## +## This source code is subject to the terms of the BSD 3-Clause Clear License and the +## Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was +## not distributed with this source code in the LICENSE file, you can obtain it +## at aomedia.org/license/software-license/bsd-3-c-c/. If the Alliance for Open Media Patent +## License 1.0 was not distributed with this source code in the PATENTS file, you +## can obtain it at aomedia.org/license/patent-license/. +## +import collections + +from avm_stats.extract_proto import * +from avm_stats.frame_visualizations import * +from avm_stats.proto_helpers import * +from avm_stats.stats_aggregation import * + +def prediction_mode_symbol_filter(symbol: Symbol): + return symbol.source_function in ( + "read_inter_mode", + "read_intra_luma_mode", + "read_drl_index", + "read_inter_compound_mode", + ) + +class PredictionModeExtractor(CodingUnitExtractor): + PredictionMode = collections.namedtuple( + "PredictionMode", + ["width", "height", "block_size", "mode", "mode_bits", "is_intra_frame", "stream_name"], + ) + + def sample(self, coding_unit: CodingUnit): + stream_name = coding_unit.frame.proto.stream_params.stream_name.removesuffix(".bin") + width = coding_unit.rect.width + height = coding_unit.rect.height + block_size = f"{width}x{height}" + mode = coding_unit.get_prediction_mode() + mode_bits = sum( + sym.bits + 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)
diff --git a/tools/py_stats/examples/extractors/symbol_bits_extractor.py b/tools/py_stats/examples/extractors/symbol_bits_extractor.py new file mode 100644 index 0000000..2a70bd5 --- /dev/null +++ b/tools/py_stats/examples/extractors/symbol_bits_extractor.py
@@ -0,0 +1,46 @@ +#!/usr/bin/env python3 +## Copyright (c) 2023, Alliance for Open Media. All rights reserved +## +## This source code is subject to the terms of the BSD 3-Clause Clear License and the +## Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was +## not distributed with this source code in the LICENSE file, you can obtain it +## at aomedia.org/license/software-license/bsd-3-c-c/. If the Alliance for Open Media Patent +## License 1.0 was not distributed with this source code in the PATENTS file, you +## can obtain it at aomedia.org/license/patent-license/. +## +import collections + +from avm_stats.extract_proto import * +from avm_stats.frame_visualizations import * +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"], + ) + + 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 + 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)
diff --git a/tools/py_stats/examples/extractors/tx_type_extractor.py b/tools/py_stats/examples/extractors/tx_type_extractor.py new file mode 100644 index 0000000..bf9ae93 --- /dev/null +++ b/tools/py_stats/examples/extractors/tx_type_extractor.py
@@ -0,0 +1,39 @@ +#!/usr/bin/env python3 +## Copyright (c) 2023, Alliance for Open Media. All rights reserved +## +## This source code is subject to the terms of the BSD 3-Clause Clear License and the +## Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was +## not distributed with this source code in the LICENSE file, you can obtain it +## at aomedia.org/license/software-license/bsd-3-c-c/. If the Alliance for Open Media Patent +## License 1.0 was not distributed with this source code in the PATENTS file, you +## can obtain it at aomedia.org/license/patent-license/. +## +import collections + +from avm_stats.extract_proto import * +from avm_stats.frame_visualizations import * +from avm_stats.proto_helpers import * +from avm_stats.stats_aggregation import * + + +class TxTypeExtractor(CodingUnitExtractor): + TxSize = collections.namedtuple( + "TxSize", + ["width", "height", "tx_size", "tx_type", "is_intra_frame", "stream_name"], + ) + + def sample(self, coding_unit: CodingUnit): + stream_name = coding_unit.frame.proto.stream_params.stream_name.removesuffix(".bin") + # Luma only + is_chroma = len(coding_unit.proto.transform_planes) == 2 + if is_chroma: + return + is_intra_frame = coding_unit.frame.is_intra_frame + for transform_unit in coding_unit.proto.transform_planes[0].transform_units: + width = transform_unit.size.width + height = transform_unit.size.height + tx_size = f"{width}x{height}" + # 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)
diff --git a/tools/py_stats/examples/partition_tree.py b/tools/py_stats/examples/partition_tree.py index 3a26160..d2c32ad 100755 --- a/tools/py_stats/examples/partition_tree.py +++ b/tools/py_stats/examples/partition_tree.py
@@ -1,12 +1,12 @@ #!/usr/bin/env python3 -# Copyright (c) 2023, Alliance for Open Media. All rights reserved +## Copyright (c) 2023, Alliance for Open Media. All rights reserved ## -# This source code is subject to the terms of the BSD 3-Clause Clear License and the -# Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was -# not distributed with this source code in the LICENSE file, you can obtain it -# at aomedia.org/license/software-license/bsd-3-c-c/. If the Alliance for Open Media Patent -# License 1.0 was not distributed with this source code in the PATENTS file, you -# can obtain it at aomedia.org/license/patent-license/. +## This source code is subject to the terms of the BSD 3-Clause Clear License and the +## Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was +## not distributed with this source code in the LICENSE file, you can obtain it +## at aomedia.org/license/software-license/bsd-3-c-c/. If the Alliance for Open Media Patent +## License 1.0 was not distributed with this source code in the PATENTS file, you +## can obtain it at aomedia.org/license/patent-license/. ## from __future__ import annotations
diff --git a/tools/py_stats/img/aggregation_block_sizes.png b/tools/py_stats/img/aggregation_block_sizes.png new file mode 100644 index 0000000..8064457 --- /dev/null +++ b/tools/py_stats/img/aggregation_block_sizes.png Binary files differ
diff --git a/tools/py_stats/img/aggregation_partition_types.png b/tools/py_stats/img/aggregation_partition_types.png new file mode 100644 index 0000000..f632dea --- /dev/null +++ b/tools/py_stats/img/aggregation_partition_types.png Binary files differ
diff --git a/tools/py_stats/img/aggregation_prediction_modes.png b/tools/py_stats/img/aggregation_prediction_modes.png new file mode 100644 index 0000000..81dafce --- /dev/null +++ b/tools/py_stats/img/aggregation_prediction_modes.png Binary files differ
diff --git a/tools/py_stats/img/aggregation_prediction_modes_by_bits.png b/tools/py_stats/img/aggregation_prediction_modes_by_bits.png new file mode 100644 index 0000000..e87f9d9 --- /dev/null +++ b/tools/py_stats/img/aggregation_prediction_modes_by_bits.png Binary files differ
diff --git a/tools/py_stats/img/aggregation_symbols.png b/tools/py_stats/img/aggregation_symbols.png new file mode 100644 index 0000000..cc6e226 --- /dev/null +++ b/tools/py_stats/img/aggregation_symbols.png Binary files differ
diff --git a/tools/py_stats/img/aggregation_tx_types.png b/tools/py_stats/img/aggregation_tx_types.png new file mode 100644 index 0000000..6c63310 --- /dev/null +++ b/tools/py_stats/img/aggregation_tx_types.png Binary files differ
diff --git a/tools/py_stats/img/bits_heatmap.png b/tools/py_stats/img/bits_heatmap.png new file mode 100644 index 0000000..bf7e6e5 --- /dev/null +++ b/tools/py_stats/img/bits_heatmap.png Binary files differ
diff --git a/tools/py_stats/img/bits_heatmap_read_intra_luma_mode.png b/tools/py_stats/img/bits_heatmap_read_intra_luma_mode.png new file mode 100644 index 0000000..4db7259 --- /dev/null +++ b/tools/py_stats/img/bits_heatmap_read_intra_luma_mode.png Binary files differ
diff --git a/tools/py_stats/img/partition_tree_luma_inter.png b/tools/py_stats/img/partition_tree_luma_inter.png new file mode 100644 index 0000000..e4ee744 --- /dev/null +++ b/tools/py_stats/img/partition_tree_luma_inter.png Binary files differ
diff --git a/tools/py_stats/img/partition_tree_luma_intra.png b/tools/py_stats/img/partition_tree_luma_intra.png new file mode 100644 index 0000000..169dffc --- /dev/null +++ b/tools/py_stats/img/partition_tree_luma_intra.png Binary files differ
diff --git a/tools/py_stats/img/pixel_pipeline.png b/tools/py_stats/img/pixel_pipeline.png new file mode 100644 index 0000000..7fcac2d --- /dev/null +++ b/tools/py_stats/img/pixel_pipeline.png Binary files differ
diff --git a/tools/py_stats/img/prediction_modes_luma_intra.png b/tools/py_stats/img/prediction_modes_luma_intra.png new file mode 100644 index 0000000..d27c7b8 --- /dev/null +++ b/tools/py_stats/img/prediction_modes_luma_intra.png Binary files differ
diff --git a/tools/py_stats/pyproject.toml b/tools/py_stats/pyproject.toml index 469ba6a..6d94bab 100644 --- a/tools/py_stats/pyproject.toml +++ b/tools/py_stats/pyproject.toml
@@ -7,7 +7,7 @@ license = "BSD-3-Clause" readme = "README.md" # By default Poetry will use .gitignore, which excludes the generated protobuf bindings. -include = ["*.py"] +include = ["**/*.py"] [tool.poetry.dependencies] absl-py = "^2.0.0"
