lama

1
#!/usr/bin/env python3
2

3
import glob
4
import os
5

6
import PIL.Image as Image
7
import cv2
8
import numpy as np
9
import tqdm
10
import shutil
11

12

13
from saicinpainting.evaluation.utils import load_yaml
14

15

16
def generate_masks_for_img(infile, outmask_pattern, mask_size=200, step=0.5):
17
    inimg = Image.open(infile)
18
    width, height = inimg.size
19
    step_abs = int(mask_size * step)
20

21
    mask = np.zeros((height, width), dtype='uint8')
22
    mask_i = 0
23

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] = 255
27

28
            cv2.imwrite(outmask_pattern.format(mask_i), mask)
29

30
            mask[start_vertical:start_vertical + mask_size, start_horizontal:start_horizontal + mask_size] = 0
31
            mask_i += 1
32

33

34
def 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

53
if __name__ == '__main__':
54
    import argparse
55

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