CVPR/lama-example

1#!/usr/bin/env python32 3import glob4import os5 6import PIL.Image as Image7import cv28import numpy as np9import tqdm10import shutil11 12 13from saicinpainting.evaluation.utils import load_yaml14 15 16def generate_masks_for_img(infile, outmask_pattern, mask_size=200, step=0.5):17    inimg = Image.open(infile)18    width, height = inimg.size19    step_abs = int(mask_size * step)20 21    mask = np.zeros((height, width), dtype='uint8')22    mask_i = 023 24    for start_vertical in range(0, height - step_abs, step_abs):25        for start_horizontal in range(0, width - step_abs, step_abs):26            mask[start_vertical:start_vertical + mask_size, start_horizontal:start_horizontal + mask_size] = 25527 28            cv2.imwrite(outmask_pattern.format(mask_i), mask)29 30            mask[start_vertical:start_vertical + mask_size, start_horizontal:start_horizontal + mask_size] = 031            mask_i += 132 33 34def main(args):35    if not args.indir.endswith('/'):36        args.indir += '/'37    if not args.outdir.endswith('/'):38        args.outdir += '/'39 40    config = load_yaml(args.config)41 42    in_files = list(glob.glob(os.path.join(args.indir, '**', f'*{config.img_ext}'), recursive=True))43    for infile in tqdm.tqdm(in_files):44        outimg = args.outdir + infile[len(args.indir):]45        outmask_pattern = outimg[:-len(config.img_ext)] + '_mask{:04d}.png'46 47        os.makedirs(os.path.dirname(outimg), exist_ok=True)48        shutil.copy2(infile, outimg)49 50        generate_masks_for_img(infile, outmask_pattern, **config.gen_kwargs)51 52 53if __name__ == '__main__':54    import argparse55 56    aparser = argparse.ArgumentParser()57    aparser.add_argument('config', type=str, help='Path to config for dataset generation')58    aparser.add_argument('indir', type=str, help='Path to folder with images')59    aparser.add_argument('outdir', type=str, help='Path to folder to store aligned images and masks to')60 61    main(aparser.parse_args())62