CVPR/lama-example

1#!/usr/bin/env python32 3import os4 5import numpy as np6import tqdm7from skimage import io8from skimage.segmentation import mark_boundaries9 10from saicinpainting.evaluation.data import InpaintingDataset11from saicinpainting.evaluation.vis import save_item_for_vis12 13def save_mask_for_sidebyside(item, out_file):14    mask = item['mask']# > 0.515    if mask.ndim == 3:16        mask = mask[0]17    mask = np.clip(mask * 255, 0, 255).astype('uint8')18    io.imsave(out_file, mask)19 20def save_img_for_sidebyside(item, out_file):21    img = np.transpose(item['image'], (1, 2, 0))22    img = np.clip(img * 255, 0, 255).astype('uint8')23    io.imsave(out_file, img)24 25def save_masked_img_for_sidebyside(item, out_file):26    mask = item['mask']27    img  = item['image']28 29    img = (1-mask) * img + mask30    img = np.transpose(img, (1, 2, 0))31 32    img = np.clip(img * 255, 0, 255).astype('uint8')33    io.imsave(out_file, img)34 35def main(args):36    dataset = InpaintingDataset(args.datadir, img_suffix='.png')37 38    area_bins = np.linspace(0, 1, args.area_bins + 1)39 40    heights = []41    widths = []42    image_areas = []43    hole_areas = []44    hole_area_percents = []45    area_bins_count = np.zeros(args.area_bins)46    area_bin_titles = [f'{area_bins[i] * 100:.0f}-{area_bins[i + 1] * 100:.0f}' for i in range(args.area_bins)]47 48    bin2i = [[] for _ in range(args.area_bins)]49 50    for i, item in enumerate(tqdm.tqdm(dataset)):51        h, w = item['image'].shape[1:]52        heights.append(h)53        widths.append(w)54        full_area = h * w55        image_areas.append(full_area)56        hole_area = (item['mask'] == 1).sum()57        hole_areas.append(hole_area)58        hole_percent = hole_area / full_area59        hole_area_percents.append(hole_percent)60        bin_i = np.clip(np.searchsorted(area_bins, hole_percent) - 1, 0, len(area_bins_count) - 1)61        area_bins_count[bin_i] += 162        bin2i[bin_i].append(i)63 64    os.makedirs(args.outdir, exist_ok=True)65   66    for bin_i in range(args.area_bins):67        bindir = os.path.join(args.outdir, area_bin_titles[bin_i])68        os.makedirs(bindir, exist_ok=True)69        bin_idx = bin2i[bin_i]70        for sample_i in np.random.choice(bin_idx, size=min(len(bin_idx), args.samples_n), replace=False):71            item = dataset[sample_i]72            path = os.path.join(bindir, dataset.img_filenames[sample_i].split('/')[-1])73            save_masked_img_for_sidebyside(item, path)74           75 76if __name__ == '__main__':77    import argparse78 79    aparser = argparse.ArgumentParser()80    aparser.add_argument('--datadir', type=str,81                         help='Path to folder with images and masks (output of gen_mask_dataset.py)')82    aparser.add_argument('--outdir', type=str, help='Where to put results')83    aparser.add_argument('--samples-n', type=int, default=10,84                         help='Number of sample images with masks to copy for visualization for each area bin')85    aparser.add_argument('--area-bins', type=int, default=10, help='How many area bins to have')86 87    main(aparser.parse_args())88