lama

1
#!/usr/bin/env python3
2
import os
3
import random
4

5
import cv2
6
import numpy as np
7

8
from saicinpainting.evaluation.data import PrecomputedInpaintingResultsDataset
9
from saicinpainting.evaluation.utils import load_yaml
10
from saicinpainting.training.visualizers.base import visualize_mask_and_images
11

12

13
def main(args):
14
    config = load_yaml(args.config)
15

16
    datasets = [PrecomputedInpaintingResultsDataset(args.datadir, cur_predictdir, **config.dataset_kwargs)
17
                for cur_predictdir in args.predictdirs]
18
    assert len({len(ds) for ds in datasets}) == 1
19
    len_first = len(datasets[0])
20

21
    indices = list(range(len_first))
22
    if len_first > args.max_n:
23
        indices = sorted(random.sample(indices, args.max_n))
24

25
    os.makedirs(args.outpath, exist_ok=True)
26

27
    filename2i = {}
28

29
    keys = ['image'] + [i for i in range(len(datasets))]
30
    for img_i in indices:
31
        try:
32
            mask_fname = os.path.basename(datasets[0].mask_filenames[img_i])
33
            if mask_fname in filename2i:
34
                filename2i[mask_fname] += 1
35
                idx = filename2i[mask_fname]
36
                mask_fname_only, ext = os.path.split(mask_fname)
37
                mask_fname = f'{mask_fname_only}_{idx}{ext}'
38
            else:
39
                filename2i[mask_fname] = 1
40

41
            cur_vis_dict = datasets[0][img_i]
42
            for ds_i, ds in enumerate(datasets):
43
                cur_vis_dict[ds_i] = ds[img_i]['inpainted']
44

45
            vis_img = visualize_mask_and_images(cur_vis_dict, keys,
46
                                                last_without_mask=False,
47
                                                mask_only_first=True,
48
                                                black_mask=args.black)
49
            vis_img = np.clip(vis_img * 255, 0, 255).astype('uint8')
50

51
            out_fname = os.path.join(args.outpath, mask_fname)
52

53

54

55
            vis_img = cv2.cvtColor(vis_img, cv2.COLOR_RGB2BGR)
56
            cv2.imwrite(out_fname, vis_img)
57
        except Exception as ex:
58
            print(f'Could not process {img_i} due to {ex}')
59

60

61
if __name__ == '__main__':
62
    import argparse
63

64
    aparser = argparse.ArgumentParser()
65
    aparser.add_argument('--max-n', type=int, default=100, help='Maximum number of images to print')
66
    aparser.add_argument('--black', action='store_true', help='Whether to fill mask on GT with black')
67
    aparser.add_argument('config', type=str, help='Path to evaluation config (e.g. configs/eval1.yaml)')
68
    aparser.add_argument('outpath', type=str, help='Where to put results')
69
    aparser.add_argument('datadir', type=str,
70
                         help='Path to folder with images and masks')
71
    aparser.add_argument('predictdirs', type=str,
72
                         nargs='+',
73
                         help='Path to folders with predicts')
74

75

76
    main(aparser.parse_args())
77