1## Copyright (c) 2020 The WebM project authors. All Rights Reserved. 2## 3## Use of this source code is governed by a BSD-style license 4## that can be found in the LICENSE file in the root of the source 5## tree. An additional intellectual property rights grant can be found 6## in the file PATENTS. All contributing project authors may 7## be found in the AUTHORS file in the root of the source tree. 8## 9 10import argparse 11from os import listdir, path 12from PIL import Image 13import sys 14 15parser = argparse.ArgumentParser() 16parser.add_argument("--frame_path", default="../data/frame/", type=str) 17parser.add_argument("--frame_rate", default="25:1", type=str) 18parser.add_argument("--interlacing", default="Ip", type=str) 19parser.add_argument("--pix_ratio", default="0:0", type=str) 20parser.add_argument("--color_space", default="4:2:0", type=str) 21parser.add_argument("--output", default="output.y4m", type=str) 22 23 24def generate(args, frames): 25 if len(frames) == 0: 26 return 27 #sort the frames based on the frame index 28 frames = sorted(frames, key=lambda x: x[0]) 29 #convert the frames to YUV form 30 frames = [f.convert("YCbCr") for _, f in frames] 31 #write the header 32 header = "YUV4MPEG2 W%d H%d F%s %s A%s" % (frames[0].width, frames[0].height, 33 args.frame_rate, args.interlacing, 34 args.pix_ratio) 35 cs = args.color_space.split(":") 36 header += " C%s%s%s\n" % (cs[0], cs[1], cs[2]) 37 #estimate the sample step based on subsample value 38 subsamples = [int(c) for c in cs] 39 r_step = [1, int(subsamples[2] == 0) + 1, int(subsamples[2] == 0) + 1] 40 c_step = [1, 4 // subsamples[1], 4 // subsamples[1]] 41 #write in frames 42 with open(args.output, "wb") as y4m: 43 y4m.write(header) 44 for f in frames: 45 y4m.write("FRAME\n") 46 px = f.load() 47 for k in xrange(3): 48 for i in xrange(0, f.height, r_step[k]): 49 for j in xrange(0, f.width, c_step[k]): 50 yuv = px[j, i] 51 y4m.write(chr(yuv[k])) 52 53 54if __name__ == "__main__": 55 args = parser.parse_args() 56 frames = [] 57 frames_mv = [] 58 for filename in listdir(args.frame_path): 59 name, ext = filename.split(".") 60 if ext == "png": 61 name_parse = name.split("_") 62 idx = int(name_parse[-1]) 63 img = Image.open(path.join(args.frame_path, filename)) 64 if name_parse[-2] == "mv": 65 frames_mv.append((idx, img)) 66 else: 67 frames.append((idx, img)) 68 if len(frames) == 0: 69 print("No frames in directory: " + args.frame_path) 70 sys.exit() 71 print("----------------------Y4M Info----------------------") 72 print("width: %d" % frames[0][1].width) 73 print("height: %d" % frames[0][1].height) 74 print("#frame: %d" % len(frames)) 75 print("frame rate: %s" % args.frame_rate) 76 print("interlacing: %s" % args.interlacing) 77 print("pixel ratio: %s" % args.pix_ratio) 78 print("color space: %s" % args.color_space) 79 print("----------------------------------------------------") 80 81 print("Generating ...") 82 generate(args, frames) 83 if len(frames_mv) != 0: 84 args.output = args.output.replace(".y4m", "_mv.y4m") 85 generate(args, frames_mv) 86