CVPR/lama-example

1#!/usr/bin/env python32import os3import random4 5import cv26import numpy as np7 8from saicinpainting.evaluation.data import PrecomputedInpaintingResultsDataset9from saicinpainting.evaluation.utils import load_yaml10from saicinpainting.training.visualizers.base import visualize_mask_and_images11 12 13def 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}) == 119    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] += 135                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] = 140 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 61if __name__ == '__main__':62    import argparse63 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