|
|
|
|
| import os
|
|
|
| import numpy as np
|
| import tqdm
|
| from skimage import io
|
| from skimage.segmentation import mark_boundaries
|
|
|
| from saicinpainting.evaluation.data import InpaintingDataset
|
| from saicinpainting.evaluation.vis import save_item_for_vis
|
|
|
| def save_mask_for_sidebyside(item, out_file):
|
| mask = item['mask']
|
| if mask.ndim == 3:
|
| mask = mask[0]
|
| mask = np.clip(mask * 255, 0, 255).astype('uint8')
|
| io.imsave(out_file, mask)
|
|
|
| def save_img_for_sidebyside(item, out_file):
|
| img = np.transpose(item['image'], (1, 2, 0))
|
| img = np.clip(img * 255, 0, 255).astype('uint8')
|
| io.imsave(out_file, img)
|
|
|
| def save_masked_img_for_sidebyside(item, out_file):
|
| mask = item['mask']
|
| img = item['image']
|
|
|
| img = (1-mask) * img + mask
|
| img = np.transpose(img, (1, 2, 0))
|
|
|
| img = np.clip(img * 255, 0, 255).astype('uint8')
|
| io.imsave(out_file, img)
|
|
|
| def main(args):
|
| dataset = InpaintingDataset(args.datadir, img_suffix='.png')
|
|
|
| area_bins = np.linspace(0, 1, args.area_bins + 1)
|
|
|
| heights = []
|
| widths = []
|
| image_areas = []
|
| hole_areas = []
|
| hole_area_percents = []
|
| area_bins_count = np.zeros(args.area_bins)
|
| area_bin_titles = [f'{area_bins[i] * 100:.0f}-{area_bins[i + 1] * 100:.0f}' for i in range(args.area_bins)]
|
|
|
| bin2i = [[] for _ in range(args.area_bins)]
|
|
|
| for i, item in enumerate(tqdm.tqdm(dataset)):
|
| h, w = item['image'].shape[1:]
|
| heights.append(h)
|
| widths.append(w)
|
| full_area = h * w
|
| image_areas.append(full_area)
|
| hole_area = (item['mask'] == 1).sum()
|
| hole_areas.append(hole_area)
|
| hole_percent = hole_area / full_area
|
| hole_area_percents.append(hole_percent)
|
| bin_i = np.clip(np.searchsorted(area_bins, hole_percent) - 1, 0, len(area_bins_count) - 1)
|
| area_bins_count[bin_i] += 1
|
| bin2i[bin_i].append(i)
|
|
|
| os.makedirs(args.outdir, exist_ok=True)
|
|
|
| for bin_i in range(args.area_bins):
|
| bindir = os.path.join(args.outdir, area_bin_titles[bin_i])
|
| os.makedirs(bindir, exist_ok=True)
|
| bin_idx = bin2i[bin_i]
|
| for sample_i in np.random.choice(bin_idx, size=min(len(bin_idx), args.samples_n), replace=False):
|
| item = dataset[sample_i]
|
| path = os.path.join(bindir, dataset.img_filenames[sample_i].split('/')[-1])
|
| save_masked_img_for_sidebyside(item, path)
|
|
|
|
|
| if __name__ == '__main__':
|
| import argparse
|
|
|
| aparser = argparse.ArgumentParser()
|
| aparser.add_argument('--datadir', type=str,
|
| help='Path to folder with images and masks (output of gen_mask_dataset.py)')
|
| aparser.add_argument('--outdir', type=str, help='Where to put results')
|
| aparser.add_argument('--samples-n', type=int, default=10,
|
| help='Number of sample images with masks to copy for visualization for each area bin')
|
| aparser.add_argument('--area-bins', type=int, default=10, help='How many area bins to have')
|
|
|
| main(aparser.parse_args())
|
|
|
