130 lines
4.0 KiB
Python
130 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
import os
|
|
import re
|
|
import struct
|
|
import sys
|
|
import zlib
|
|
|
|
|
|
def write_png_rgb(path, width, height, rgb):
|
|
def chunk(tag, data):
|
|
body = tag + data
|
|
return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body) & 0xFFFFFFFF)
|
|
|
|
rows = bytearray()
|
|
stride = width * 3
|
|
for y in range(height):
|
|
rows.append(0)
|
|
start = y * stride
|
|
rows.extend(rgb[start:start + stride])
|
|
|
|
ihdr = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0)
|
|
with open(path, "wb") as f:
|
|
f.write(b"\x89PNG\r\n\x1a\n")
|
|
f.write(chunk(b"IHDR", ihdr))
|
|
f.write(chunk(b"IDAT", zlib.compress(bytes(rows), 6)))
|
|
f.write(chunk(b"IEND", b""))
|
|
|
|
|
|
def write_png_gray(path, width, height, gray):
|
|
def chunk(tag, data):
|
|
body = tag + data
|
|
return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body) & 0xFFFFFFFF)
|
|
|
|
rows = bytearray()
|
|
for y in range(height):
|
|
rows.append(0)
|
|
start = y * width
|
|
rows.extend(gray[start:start + width])
|
|
|
|
ihdr = struct.pack(">IIBBBBB", width, height, 8, 0, 0, 0, 0)
|
|
with open(path, "wb") as f:
|
|
f.write(b"\x89PNG\r\n\x1a\n")
|
|
f.write(chunk(b"IHDR", ihdr))
|
|
f.write(chunk(b"IDAT", zlib.compress(bytes(rows), 6)))
|
|
f.write(chunk(b"IEND", b""))
|
|
|
|
|
|
def clamp(v):
|
|
if v < 0:
|
|
return 0
|
|
if v > 255:
|
|
return 255
|
|
return v
|
|
|
|
|
|
def yuv420sp_to_rgb(y_plane, uv_plane, width, height, nv21=False):
|
|
rgb = bytearray(width * height * 3)
|
|
for y in range(height):
|
|
y_row = y * width
|
|
uv_row = (y // 2) * width
|
|
out_row = y_row * 3
|
|
for x in range(width):
|
|
yy = y_plane[y_row + x]
|
|
uv_index = uv_row + (x & ~1)
|
|
if nv21:
|
|
v = uv_plane[uv_index] - 128
|
|
u = uv_plane[uv_index + 1] - 128
|
|
else:
|
|
u = uv_plane[uv_index] - 128
|
|
v = uv_plane[uv_index + 1] - 128
|
|
|
|
c = yy - 16
|
|
if c < 0:
|
|
c = 0
|
|
r = (298 * c + 409 * v + 128) >> 8
|
|
g = (298 * c - 100 * u - 208 * v + 128) >> 8
|
|
b = (298 * c + 516 * u + 128) >> 8
|
|
|
|
out = out_row + x * 3
|
|
rgb[out] = clamp(r)
|
|
rgb[out + 1] = clamp(g)
|
|
rgb[out + 2] = clamp(b)
|
|
return rgb
|
|
|
|
|
|
def main():
|
|
dump_dir = sys.argv[1] if len(sys.argv) > 1 else "/home/user/dump"
|
|
out_dir = os.path.join(dump_dir, "png")
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
pattern = re.compile(r"^(\d{8})_(\d+)X(\d+)_y\.raw$")
|
|
converted = 0
|
|
for name in sorted(os.listdir(dump_dir)):
|
|
m = pattern.match(name)
|
|
if not m:
|
|
continue
|
|
frame, width_s, height_s = m.groups()
|
|
width = int(width_s)
|
|
height = int(height_s)
|
|
y_path = os.path.join(dump_dir, name)
|
|
uv_path = os.path.join(dump_dir, f"{frame}_{width}X{height}_uv.raw")
|
|
if not os.path.exists(uv_path):
|
|
print(f"skip {frame}: missing {uv_path}")
|
|
continue
|
|
|
|
with open(y_path, "rb") as f:
|
|
y_plane = f.read()
|
|
with open(uv_path, "rb") as f:
|
|
uv_plane = f.read()
|
|
|
|
expected_y = width * height
|
|
expected_uv = width * height // 2
|
|
if len(y_plane) != expected_y or len(uv_plane) != expected_uv:
|
|
print(f"skip {frame}: bad size y={len(y_plane)} uv={len(uv_plane)} expected {expected_y}/{expected_uv}")
|
|
continue
|
|
|
|
base = os.path.join(out_dir, f"{frame}_{width}x{height}")
|
|
write_png_gray(base + "_gray.png", width, height, y_plane)
|
|
write_png_rgb(base + "_nv12.png", width, height, yuv420sp_to_rgb(y_plane, uv_plane, width, height, nv21=False))
|
|
write_png_rgb(base + "_nv21.png", width, height, yuv420sp_to_rgb(y_plane, uv_plane, width, height, nv21=True))
|
|
converted += 1
|
|
print(f"converted frame {frame}: gray/nv12/nv21")
|
|
|
|
print(f"output: {out_dir}")
|
|
print(f"frames: {converted}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|