123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146 |
- """
- python generate_imagematte_with_background_image.py \
- --imagematte-dir ../matting-data/Distinctions/test \
- --background-dir ../matting-data/Backgrounds/valid \
- --resolution 512 \
- --out-dir ../matting-data/evaluation/distinction_static_sd/ \
- --random-seed 10
-
- Seed:
- 10 - distinction-static
- 11 - distinction-motion
- 12 - adobe-static
- 13 - adobe-motion
-
- """
- import argparse
- import os
- import pims
- import numpy as np
- import random
- from PIL import Image
- from tqdm import tqdm
- from tqdm.contrib.concurrent import process_map
- from torchvision import transforms
- from torchvision.transforms import functional as F
- parser = argparse.ArgumentParser()
- parser.add_argument('--imagematte-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('--resolution', type=int, required=True)
- parser.add_argument('--out-dir', type=str, required=True)
- parser.add_argument('--random-seed', type=int)
- parser.add_argument('--extension', type=str, default='.png')
- args = parser.parse_args()
-
- random.seed(args.random_seed)
- imagematte_filenames = os.listdir(os.path.join(args.imagematte_dir, 'fgr'))
- background_filenames = os.listdir(args.background_dir)
- random.shuffle(imagematte_filenames)
- random.shuffle(background_filenames)
- def lerp(a, b, percentage):
- return a * (1 - percentage) + b * percentage
- def motion_affine(*imgs):
- config = dict(degrees=(-10, 10), translate=(0.1, 0.1),
- scale_ranges=(0.9, 1.1), shears=(-5, 5), img_size=imgs[0][0].size)
- angleA, (transXA, transYA), scaleA, (shearXA, shearYA) = transforms.RandomAffine.get_params(**config)
- angleB, (transXB, transYB), scaleB, (shearXB, shearYB) = transforms.RandomAffine.get_params(**config)
- T = len(imgs[0])
- variation_over_time = random.random()
- for t in range(T):
- percentage = (t / (T - 1)) * variation_over_time
- angle = lerp(angleA, angleB, percentage)
- transX = lerp(transXA, transXB, percentage)
- transY = lerp(transYA, transYB, percentage)
- scale = lerp(scaleA, scaleB, percentage)
- shearX = lerp(shearXA, shearXB, percentage)
- shearY = lerp(shearYA, shearYB, percentage)
- for img in imgs:
- img[t] = F.affine(img[t], angle, (transX, transY), scale, (shearX, shearY), F.InterpolationMode.BILINEAR)
- return imgs
-
- def process(i):
- imagematte_filename = imagematte_filenames[i % len(imagematte_filenames)]
- background_filename = background_filenames[i % len(background_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)
-
- with Image.open(os.path.join(args.background_dir, background_filename)) as bgr:
- bgr = bgr.convert('RGB')
-
- w, h = bgr.size
- scale = args.resolution / min(h, w)
- w, h = int(w * scale), int(h * scale)
- bgr = bgr.resize((w, h))
- bgr = F.center_crop(bgr, (args.resolution, args.resolution))
- with Image.open(os.path.join(args.imagematte_dir, 'fgr', imagematte_filename)) as fgr, \
- Image.open(os.path.join(args.imagematte_dir, 'pha', imagematte_filename)) as pha:
- fgr = fgr.convert('RGB')
- pha = pha.convert('L')
-
- fgrs = [fgr] * args.num_frames
- phas = [pha] * args.num_frames
- fgrs, phas = motion_affine(fgrs, phas)
-
- for t in tqdm(range(args.num_frames), desc=str(i).zfill(4)):
- fgr = fgrs[t]
- pha = phas[t]
-
- w, h = fgr.size
- scale = args.resolution / max(h, w)
- w, h = int(w * scale), int(h * scale)
-
- fgr = fgr.resize((w, h))
- pha = pha.resize((w, h))
-
- if h < args.resolution:
- pt = (args.resolution - h) // 2
- pb = args.resolution - h - pt
- else:
- pt = 0
- pb = 0
-
- if w < args.resolution:
- pl = (args.resolution - w) // 2
- pr = args.resolution - w - pl
- else:
- pl = 0
- pr = 0
-
- fgr = F.pad(fgr, [pl, pt, pr, pb])
- pha = F.pad(pha, [pl, pt, pr, pb])
-
- if i // len(imagematte_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) + args.extension))
- pha.save(os.path.join(out_path, 'pha', str(t).zfill(4) + args.extension))
-
- if t == 0:
- bgr.save(os.path.join(out_path, 'bgr', str(t).zfill(4) + args.extension))
- else:
- os.symlink(str(0).zfill(4) + args.extension, os.path.join(out_path, 'bgr', str(t).zfill(4) + args.extension))
-
- 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) + args.extension))
- if __name__ == '__main__':
- r = process_map(process, range(args.num_samples), max_workers=32)
|