lama

1
#!/usr/bin/env python3
2

3
# Example command:
4
# ./bin/predict.py \
5
#       model.path=<path to checkpoint, prepared by make_checkpoint.py> \
6
#       indir=<path to input data> \
7
#       outdir=<where to store predicts>
8

9
import logging
10
import os
11
import sys
12
import traceback
13

14
from saicinpainting.evaluation.utils import move_to_device
15
from saicinpainting.evaluation.refinement import refine_predict
16
os.environ['OMP_NUM_THREADS'] = '1'
17
os.environ['OPENBLAS_NUM_THREADS'] = '1'
18
os.environ['MKL_NUM_THREADS'] = '1'
19
os.environ['VECLIB_MAXIMUM_THREADS'] = '1'
20
os.environ['NUMEXPR_NUM_THREADS'] = '1'
21

22
import cv2
23
import hydra
24
import numpy as np
25
import torch
26
import tqdm
27
import yaml
28
from omegaconf import OmegaConf
29
from torch.utils.data._utils.collate import default_collate
30

31
from saicinpainting.training.data.datasets import make_default_val_dataset
32
from saicinpainting.training.trainers import load_checkpoint
33
from saicinpainting.utils import register_debug_signal_handlers
34

35
LOGGER = logging.getLogger(__name__)
36

37

38
@hydra.main(config_path='../configs/prediction', config_name='default.yaml')
39
def main(predict_config: OmegaConf):
40
    try:
41
        if sys.platform != 'win32':
42
            register_debug_signal_handlers()  # kill -10 <pid> will result in traceback dumped into log
43

44
        device = torch.device("cpu")
45

46
        train_config_path = os.path.join(predict_config.model.path, 'config.yaml')
47
        with open(train_config_path, 'r') as f:
48
            train_config = OmegaConf.create(yaml.safe_load(f))
49
        
50
        train_config.training_model.predict_only = True
51
        train_config.visualizer.kind = 'noop'
52

53
        out_ext = predict_config.get('out_ext', '.png')
54

55
        checkpoint_path = os.path.join(predict_config.model.path, 
56
                                       'models', 
57
                                       predict_config.model.checkpoint)
58
        model = load_checkpoint(train_config, checkpoint_path, strict=False, map_location='cpu')
59
        model.freeze()
60
        if not predict_config.get('refine', False):
61
            model.to(device)
62

63
        if not predict_config.indir.endswith('/'):
64
            predict_config.indir += '/'
65

66
        dataset = make_default_val_dataset(predict_config.indir, **predict_config.dataset)
67
        for img_i in tqdm.trange(len(dataset)):
68
            mask_fname = dataset.mask_filenames[img_i]
69
            cur_out_fname = os.path.join(
70
                predict_config.outdir, 
71
                os.path.splitext(mask_fname[len(predict_config.indir):])[0] + out_ext
72
            )
73
            os.makedirs(os.path.dirname(cur_out_fname), exist_ok=True)
74
            batch = default_collate([dataset[img_i]])
75
            if predict_config.get('refine', False):
76
                assert 'unpad_to_size' in batch, "Unpadded size is required for the refinement"
77
                # image unpadding is taken care of in the refiner, so that output image
78
                # is same size as the input image
79
                cur_res = refine_predict(batch, model, **predict_config.refiner)
80
                cur_res = cur_res[0].permute(1,2,0).detach().cpu().numpy()
81
            else:
82
                with torch.no_grad():
83
                    batch = move_to_device(batch, device)
84
                    batch['mask'] = (batch['mask'] > 0) * 1
85
                    batch = model(batch)                    
86
                    cur_res = batch[predict_config.out_key][0].permute(1, 2, 0).detach().cpu().numpy()
87
                    unpad_to_size = batch.get('unpad_to_size', None)
88
                    if unpad_to_size is not None:
89
                        orig_height, orig_width = unpad_to_size
90
                        cur_res = cur_res[:orig_height, :orig_width]
91

92
            cur_res = np.clip(cur_res * 255, 0, 255).astype('uint8')
93
            cur_res = cv2.cvtColor(cur_res, cv2.COLOR_RGB2BGR)
94
            cv2.imwrite(cur_out_fname, cur_res)
95

96
    except KeyboardInterrupt:
97
        LOGGER.warning('Interrupted by user')
98
    except Exception as ex:
99
        LOGGER.critical(f'Prediction failed due to {ex}:\n{traceback.format_exc()}')
100
        sys.exit(1)
101

102

103
if __name__ == '__main__':
104
    main()
105