neural_codec/DCVC/src/utils/video_writer.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import os

from PIL import Image


class PNGWriter():
    def __init__(self, dst_path, width, height):
        self.dst_path = dst_path
        self.width = width
        self.height = height
        self.padding = 5
        self.current_frame_index = 1
        os.makedirs(dst_path, exist_ok=True)

    def write_one_frame(self, rgb):
        # rgb: 3xhxw uint8 numpy array
        rgb = rgb.transpose(1, 2, 0)

        png_path = os.path.join(self.dst_path,
                                f"im{str(self.current_frame_index).zfill(self.padding)}.png"
                                )
        Image.fromarray(rgb).save(png_path)

        self.current_frame_index += 1

    def close(self):
        self.current_frame_index = 1


class YUV420Writer():
    def __init__(self, dst_path, width, height):
        if not dst_path.endswith('.yuv'):
            dst_path = dst_path + '/out.yuv'
        self.dst_path = dst_path
        self.width = width
        self.height = height

        # pylint: disable=R1732
        self.file = open(dst_path, "wb")
        # pylint: enable=R1732

    def write_one_frame(self, y, uv):
        # y: 1xhxw uint8 numpy array
        # uv: 2x(h/2)x(w/2) uint8 numpy array
        self.file.write(y.tobytes())
        self.file.write(uv.tobytes())

    def close(self):
        self.file.close()