lama

1
#!/usr/bin/env python3
2
import glob
3
import logging
4
import os
5
import shutil
6
import sys
7
import traceback
8

9
from saicinpainting.evaluation.data import load_image
10
from saicinpainting.evaluation.utils import move_to_device
11

12
os.environ['OMP_NUM_THREADS'] = '1'
13
os.environ['OPENBLAS_NUM_THREADS'] = '1'
14
os.environ['MKL_NUM_THREADS'] = '1'
15
os.environ['VECLIB_MAXIMUM_THREADS'] = '1'
16
os.environ['NUMEXPR_NUM_THREADS'] = '1'
17

18
import cv2
19
import hydra
20
import numpy as np
21
import torch
22
import tqdm
23
import yaml
24
from omegaconf import OmegaConf
25
from torch.utils.data._utils.collate import default_collate
26

27
from saicinpainting.training.data.datasets import make_default_val_dataset
28
from saicinpainting.training.trainers import load_checkpoint
29
from saicinpainting.utils import register_debug_signal_handlers
30

31
LOGGER = logging.getLogger(__name__)
32

33

34
def main(args):
35
    try:
36
        if not args.indir.endswith('/'):
37
            args.indir += '/'
38

39
        for in_img in glob.glob(os.path.join(args.indir, '**', '*' + args.img_suffix), recursive=True):
40
            if 'mask' in os.path.basename(in_img):
41
                continue
42

43
            out_img_path = os.path.join(args.outdir, os.path.splitext(in_img[len(args.indir):])[0] + '.png')
44
            out_mask_path = f'{os.path.splitext(out_img_path)[0]}_mask.png'
45

46
            os.makedirs(os.path.dirname(out_img_path), exist_ok=True)
47

48
            img = load_image(in_img)
49
            height, width = img.shape[1:]
50
            pad_h, pad_w = int(height * args.coef / 2), int(width * args.coef / 2)
51

52
            mask = np.zeros((height, width), dtype='uint8')
53

54
            if args.expand:
55
                img = np.pad(img, ((0, 0), (pad_h, pad_h), (pad_w, pad_w)))
56
                mask = np.pad(mask, ((pad_h, pad_h), (pad_w, pad_w)), mode='constant', constant_values=255)
57
            else:
58
                mask[:pad_h] = 255
59
                mask[-pad_h:] = 255
60
                mask[:, :pad_w] = 255
61
                mask[:, -pad_w:] = 255
62

63
            # img = np.pad(img, ((0, 0), (pad_h * 2, pad_h * 2), (pad_w * 2, pad_w * 2)), mode='symmetric')
64
            # mask = np.pad(mask, ((pad_h * 2, pad_h * 2), (pad_w * 2, pad_w * 2)), mode = 'symmetric')
65

66
            img = np.clip(np.transpose(img, (1, 2, 0)) * 255, 0, 255).astype('uint8')
67
            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
68
            cv2.imwrite(out_img_path, img)
69

70
            cv2.imwrite(out_mask_path, mask)
71
    except KeyboardInterrupt:
72
        LOGGER.warning('Interrupted by user')
73
    except Exception as ex:
74
        LOGGER.critical(f'Prediction failed due to {ex}:\n{traceback.format_exc()}')
75
        sys.exit(1)
76

77

78
if __name__ == '__main__':
79
    import argparse
80

81
    aparser = argparse.ArgumentParser()
82
    aparser.add_argument('indir', type=str, help='Root directory with images')
83
    aparser.add_argument('outdir', type=str, help='Where to store results')
84
    aparser.add_argument('--img-suffix', type=str, default='.png', help='Input image extension')
85
    aparser.add_argument('--expand', action='store_true', help='Generate mask by padding (true) or by cropping (false)')
86
    aparser.add_argument('--coef', type=float, default=0.2, help='How much to crop/expand in order to get masks')
87

88
    main(aparser.parse_args())
89