Spaces:
Runtime error
Runtime error
File size: 2,191 Bytes
d380b77 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
#!/usr/bin/env python3
import os
import cv2
import numpy as np
import tqdm
from saicinpainting.evaluation.data import PrecomputedInpaintingResultsDataset
from saicinpainting.evaluation.utils import load_yaml
def main(args):
config = load_yaml(args.config)
if not args.predictdir.endswith('/'):
args.predictdir += '/'
dataset = PrecomputedInpaintingResultsDataset(args.datadir, args.predictdir, **config.dataset_kwargs)
os.makedirs(os.path.dirname(args.outpath), exist_ok=True)
for img_i in tqdm.trange(len(dataset)):
pred_fname = dataset.pred_filenames[img_i]
cur_out_fname = os.path.join(args.outpath, pred_fname[len(args.predictdir):])
os.makedirs(os.path.dirname(cur_out_fname), exist_ok=True)
sample = dataset[img_i]
img = sample['image']
mask = sample['mask']
inpainted = sample['inpainted']
inpainted_blurred = cv2.GaussianBlur(np.transpose(inpainted, (1, 2, 0)),
ksize=(args.k, args.k),
sigmaX=args.s, sigmaY=args.s,
borderType=cv2.BORDER_REFLECT)
cur_res = (1 - mask) * np.transpose(img, (1, 2, 0)) + mask * inpainted_blurred
cur_res = np.clip(cur_res * 255, 0, 255).astype('uint8')
cur_res = cv2.cvtColor(cur_res, cv2.COLOR_RGB2BGR)
cv2.imwrite(cur_out_fname, cur_res)
if __name__ == '__main__':
import argparse
aparser = argparse.ArgumentParser()
aparser.add_argument('config', type=str, help='Path to evaluation config')
aparser.add_argument('datadir', type=str,
help='Path to folder with images and masks (output of gen_mask_dataset.py)')
aparser.add_argument('predictdir', type=str,
help='Path to folder with predicts (e.g. predict_hifill_baseline.py)')
aparser.add_argument('outpath', type=str, help='Where to put results')
aparser.add_argument('-s', type=float, default=0.1, help='Gaussian blur sigma')
aparser.add_argument('-k', type=int, default=5, help='Kernel size in gaussian blur')
main(aparser.parse_args())
|