| """ |
| python generate_videomatte_with_background_video.py \ |
| --videomatte-dir ../matting-data/VideoMatte240K_JPEG_HD/test \ |
| --background-dir ../matting-data/BackgroundVideos_mp4/test \ |
| --resize 512 288 \ |
| --out-dir ../matting-data/evaluation/vidematte_motion_sd/ |
| """ |
|
|
| import argparse |
| import os |
| import pims |
| import numpy as np |
| import random |
| from PIL import Image |
| from tqdm import tqdm |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument('--videomatte-dir', type=str, required=True) |
| parser.add_argument('--background-dir', type=str, required=True) |
| parser.add_argument('--num-samples', type=int, default=20) |
| parser.add_argument('--num-frames', type=int, default=100) |
| parser.add_argument('--resize', type=int, default=None, nargs=2) |
| parser.add_argument('--out-dir', type=str, required=True) |
| args = parser.parse_args() |
|
|
| |
| background_filenames = [ |
| "0000.mp4", |
| "0007.mp4", |
| "0008.mp4", |
| "0010.mp4", |
| "0013.mp4", |
| "0015.mp4", |
| "0016.mp4", |
| "0018.mp4", |
| "0021.mp4", |
| "0029.mp4", |
| "0033.mp4", |
| "0035.mp4", |
| "0039.mp4", |
| "0050.mp4", |
| "0052.mp4", |
| "0055.mp4", |
| "0060.mp4", |
| "0063.mp4", |
| "0087.mp4", |
| "0086.mp4", |
| "0090.mp4", |
| "0101.mp4", |
| "0110.mp4", |
| "0117.mp4", |
| "0120.mp4", |
| "0122.mp4", |
| "0123.mp4", |
| "0125.mp4", |
| "0128.mp4", |
| "0131.mp4", |
| "0172.mp4", |
| "0176.mp4", |
| "0181.mp4", |
| "0187.mp4", |
| "0193.mp4", |
| "0198.mp4", |
| "0220.mp4", |
| "0221.mp4", |
| "0224.mp4", |
| "0229.mp4", |
| "0233.mp4", |
| "0238.mp4", |
| "0241.mp4", |
| "0245.mp4", |
| "0246.mp4" |
| ] |
|
|
| random.seed(10) |
| |
| videomatte_filenames = [(clipname, sorted(os.listdir(os.path.join(args.videomatte_dir, 'fgr', clipname)))) |
| for clipname in sorted(os.listdir(os.path.join(args.videomatte_dir, 'fgr')))] |
|
|
| random.shuffle(background_filenames) |
|
|
| for i in range(args.num_samples): |
| bgrs = pims.PyAVVideoReader(os.path.join(args.background_dir, background_filenames[i % len(background_filenames)])) |
| clipname, framenames = videomatte_filenames[i % len(videomatte_filenames)] |
| |
| out_path = os.path.join(args.out_dir, str(i).zfill(4)) |
| os.makedirs(os.path.join(out_path, 'fgr'), exist_ok=True) |
| os.makedirs(os.path.join(out_path, 'pha'), exist_ok=True) |
| os.makedirs(os.path.join(out_path, 'com'), exist_ok=True) |
| os.makedirs(os.path.join(out_path, 'bgr'), exist_ok=True) |
| |
| base_t = random.choice(range(len(framenames) - args.num_frames)) |
| |
| for t in tqdm(range(args.num_frames), desc=str(i).zfill(4)): |
| with Image.open(os.path.join(args.videomatte_dir, 'fgr', clipname, framenames[base_t + t])) as fgr, \ |
| Image.open(os.path.join(args.videomatte_dir, 'pha', clipname, framenames[base_t + t])) as pha: |
| fgr = fgr.convert('RGB') |
| pha = pha.convert('L') |
| |
| if args.resize is not None: |
| fgr = fgr.resize(args.resize, Image.BILINEAR) |
| pha = pha.resize(args.resize, Image.BILINEAR) |
| |
| |
| if i // len(videomatte_filenames) % 2 == 1: |
| fgr = fgr.transpose(Image.FLIP_LEFT_RIGHT) |
| pha = pha.transpose(Image.FLIP_LEFT_RIGHT) |
| |
| fgr.save(os.path.join(out_path, 'fgr', str(t).zfill(4) + '.png')) |
| pha.save(os.path.join(out_path, 'pha', str(t).zfill(4) + '.png')) |
| |
| bgr = Image.fromarray(bgrs[t]) |
| bgr = bgr.resize(fgr.size, Image.BILINEAR) |
| bgr.save(os.path.join(out_path, 'bgr', str(t).zfill(4) + '.png')) |
| |
| pha = np.asarray(pha).astype(float)[:, :, None] / 255 |
| com = Image.fromarray(np.uint8(np.asarray(fgr) * pha + np.asarray(bgr) * (1 - pha))) |
| com.save(os.path.join(out_path, 'com', str(t).zfill(4) + '.png')) |
|
|
