ref: 7e4b8a86b20a967ea0f2eb875137d69e0c5ce2b9
parent: 729da3df8dc5cff807cb3278f419dd1601a673db
author: Dan Zhu <[email protected]>
date: Thu Jun 20 12:05:30 EDT 2019
script to compact frames to y4m video Change-Id: I2d8c3ccf49c172a54181aeb2e2b8169bf5402456
--- /dev/null
+++ b/tools/3D-Reconstruction/genY4M/genY4M.py
@@ -1,0 +1,76 @@
+import argparse
+from os import listdir, path
+from PIL import Image
+import sys
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--frame_path", default="../data/frame/", type=str)+parser.add_argument("--frame_rate", default="25:1", type=str)+parser.add_argument("--interlacing", default="Ip", type=str)+parser.add_argument("--pix_ratio", default="0:0", type=str)+parser.add_argument("--color_space", default="4:2:0", type=str)+parser.add_argument("--output", default="output.y4m", type=str)+
+
+def generate(args, frames):
+ if len(frames) == 0:
+ return
+ #sort the frames based on the frame index
+ frames = sorted(frames, key=lambda x: x[0])
+ #convert the frames to YUV form
+ frames = [f.convert("YCbCr") for _, f in frames]+ #write the header
+ header = "YUV4MPEG2 W%d H%d F%s %s A%s" % (frames[0].width, frames[0].height,
+ args.frame_rate, args.interlacing,
+ args.pix_ratio)
+ cs = args.color_space.split(":")+ header += " C%s%s%s\n" % (cs[0], cs[1], cs[2])
+ #estimate the sample step based on subsample value
+ subsamples = [int(c) for c in cs]
+ r_step = [1, int(subsamples[2] == 0) + 1, int(subsamples[2] == 0) + 1]
+ c_step = [1, 4 // subsamples[1], 4 // subsamples[1]]
+ #write in frames
+ with open(args.output, "wb") as y4m:
+ y4m.write(header)
+ for f in frames:
+ y4m.write("FRAME\n")+ px = f.load()
+ for k in xrange(3):
+ for i in xrange(0, f.height, r_step[k]):
+ for j in xrange(0, f.width, c_step[k]):
+ yuv = px[j, i]
+ y4m.write(chr(yuv[k]))
+
+
+if __name__ == "__main__":
+ args = parser.parse_args()
+ frames = []
+ frames_mv = []
+ for filename in listdir(args.frame_path):
+ name, ext = filename.split(".")+ if ext == "png":
+ name_parse = name.split("_")+ idx = int(name_parse[-1])
+ img = Image.open(path.join(args.frame_path, filename))
+ if name_parse[-2] == "mv":
+ frames_mv.append((idx, img))
+ else:
+ frames.append((idx, img))
+ if len(frames) == 0:
+ print "No frames in directory: " + args.frame_path
+ sys.exit()
+ print("----------------------Y4M Info----------------------")+ print("width: %d" % frames[0][1].width)+ print("height: %d" % frames[0][1].height)+ print("#frame: %d" % len(frames))+ print("frame rate: %s" % args.frame_rate)+ print("interlacing: %s" % args.interlacing)+ print("pixel ratio: %s" % args.pix_ratio)+ print("color space: %s" % args.color_space)+ print("----------------------------------------------------")+
+ print("Generating ...")+ generate(args, frames)
+ if len(frames_mv) != 0:
+ args.output = args.output.replace(".y4m", "_mv.y4m")+ generate(args, frames_mv)