File size: 1,402 Bytes
12acbba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | # 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()
|