|
|
|
import glob |
|
import logging |
|
import os |
|
import shutil |
|
import sys |
|
import traceback |
|
|
|
from saicinpainting.evaluation.data import load_image |
|
from saicinpainting.evaluation.utils import move_to_device |
|
|
|
os.environ['OMP_NUM_THREADS'] = '1' |
|
os.environ['OPENBLAS_NUM_THREADS'] = '1' |
|
os.environ['MKL_NUM_THREADS'] = '1' |
|
os.environ['VECLIB_MAXIMUM_THREADS'] = '1' |
|
os.environ['NUMEXPR_NUM_THREADS'] = '1' |
|
|
|
import cv2 |
|
import hydra |
|
import numpy as np |
|
import torch |
|
import tqdm |
|
import yaml |
|
from omegaconf import OmegaConf |
|
from torch.utils.data._utils.collate import default_collate |
|
|
|
from saicinpainting.training.data.datasets import make_default_val_dataset |
|
from saicinpainting.training.trainers import load_checkpoint |
|
from saicinpainting.utils import register_debug_signal_handlers |
|
|
|
LOGGER = logging.getLogger(__name__) |
|
|
|
|
|
def main(args): |
|
try: |
|
if not args.indir.endswith('/'): |
|
args.indir += '/' |
|
|
|
for in_img in glob.glob(os.path.join(args.indir, '**', '*' + args.img_suffix), recursive=True): |
|
if 'mask' in os.path.basename(in_img): |
|
continue |
|
|
|
out_img_path = os.path.join(args.outdir, os.path.splitext(in_img[len(args.indir):])[0] + '.png') |
|
out_mask_path = f'{os.path.splitext(out_img_path)[0]}_mask.png' |
|
|
|
os.makedirs(os.path.dirname(out_img_path), exist_ok=True) |
|
|
|
img = load_image(in_img) |
|
height, width = img.shape[1:] |
|
pad_h, pad_w = int(height * args.coef / 2), int(width * args.coef / 2) |
|
|
|
mask = np.zeros((height, width), dtype='uint8') |
|
|
|
if args.expand: |
|
img = np.pad(img, ((0, 0), (pad_h, pad_h), (pad_w, pad_w))) |
|
mask = np.pad(mask, ((pad_h, pad_h), (pad_w, pad_w)), mode='constant', constant_values=255) |
|
else: |
|
mask[:pad_h] = 255 |
|
mask[-pad_h:] = 255 |
|
mask[:, :pad_w] = 255 |
|
mask[:, -pad_w:] = 255 |
|
|
|
|
|
|
|
|
|
img = np.clip(np.transpose(img, (1, 2, 0)) * 255, 0, 255).astype('uint8') |
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) |
|
cv2.imwrite(out_img_path, img) |
|
|
|
cv2.imwrite(out_mask_path, mask) |
|
except KeyboardInterrupt: |
|
LOGGER.warning('Interrupted by user') |
|
except Exception as ex: |
|
LOGGER.critical(f'Prediction failed due to {ex}:\n{traceback.format_exc()}') |
|
sys.exit(1) |
|
|
|
|
|
if __name__ == '__main__': |
|
import argparse |
|
|
|
aparser = argparse.ArgumentParser() |
|
aparser.add_argument('indir', type=str, help='Root directory with images') |
|
aparser.add_argument('outdir', type=str, help='Where to store results') |
|
aparser.add_argument('--img-suffix', type=str, default='.png', help='Input image extension') |
|
aparser.add_argument('--expand', action='store_true', help='Generate mask by padding (true) or by cropping (false)') |
|
aparser.add_argument('--coef', type=float, default=0.2, help='How much to crop/expand in order to get masks') |
|
|
|
main(aparser.parse_args()) |
|
|
