| #!/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/. |
| ## |
| """Helper tool to dump pixel data from a frame proto as YUVs.""" |
| |
| from collections.abc import Sequence |
| import pathlib |
| |
| from absl import app |
| from absl import flags |
| from absl import logging |
| from avm_stats import proto_helpers |
| |
| _PROTO = flags.DEFINE_string("proto", None, "Path to AVM frame protobuf.") |
| flags.mark_flag_as_required("proto") |
| |
| _OUTPUT_FOLDER = flags.DEFINE_string( |
| "output_folder", None, "Path to YUV output folder." |
| ) |
| flags.mark_flag_as_required("output_folder") |
| |
| _PIXEL_FIELDS = [ |
| "original", |
| "prediction", |
| "pre_filtered", |
| "reconstruction", |
| ] |
| |
| |
| def dump_yuv( |
| frame: proto_helpers.Frame, field_name: str, output_path: pathlib.Path |
| ): |
| """Dumps pixels from a frame as raw YUV. |
| |
| Args: |
| frame: Frame to extract pixels from. |
| field_name: Name of the pixel field to extract, e.g. "original" or |
| "prediction". |
| output_path: Path to output YUV. |
| """ |
| with output_path.open("wb") as f: |
| for plane_id in range(3): |
| if not hasattr(frame.pixels[plane_id], field_name): |
| logging.warning( |
| "Skipping %s (pixels do not exist on proto).", output_path |
| ) |
| return |
| plane = getattr(frame.pixels[plane_id], field_name) |
| plane.tofile(f) |
| logging.info("Generated %s.", output_path) |
| |
| |
| def _create_output_path( |
| proto_path: pathlib.Path, output_folder: pathlib.Path, pixel_field: str |
| ) -> pathlib.Path: |
| proto_name = proto_path.stem |
| return output_folder / f"{proto_name}_{pixel_field}.yuv" |
| |
| |
| def dump_all_yuvs(proto_path: pathlib.Path, output_folder: pathlib.Path): |
| """Loads a proto frame and dumps each class of pixels as raw YUVs. |
| |
| Args: |
| proto_path: Path to proto frame to load and extract pixels from. |
| output_folder: Path to output folder. Output filename will be the same as |
| the input proto, but with a .yuv extension instead. |
| """ |
| frame = proto_helpers.load_frame_from_path(proto_path) |
| for pixel_field in _PIXEL_FIELDS: |
| output_path = _create_output_path(proto_path, output_folder, pixel_field) |
| dump_yuv(frame, pixel_field, output_path) |
| |
| |
| def main(argv: Sequence[str]) -> None: |
| if len(argv) > 1: |
| raise app.UsageError("Too many command-line arguments.") |
| |
| proto_path = pathlib.Path(_PROTO.value) |
| output_folder = pathlib.Path(_OUTPUT_FOLDER.value) |
| dump_all_yuvs(proto_path, output_folder) |
| |
| |
| if __name__ == "__main__": |
| app.run(main) |