files

bin/analyze_errors.py

@@ -0,0 +1,316 @@
+#!/usr/bin/env python3
+import cv2
+import numpy as np
+import sklearn
+import torch
+import os
+import pickle
+import pandas as pd
+import matplotlib.pyplot as plt
+from joblib import Parallel, delayed
+
+from saicinpainting.evaluation.data import PrecomputedInpaintingResultsDataset, load_image
+from saicinpainting.evaluation.losses.fid.inception import InceptionV3
+from saicinpainting.evaluation.utils import load_yaml
+from saicinpainting.training.visualizers.base import visualize_mask_and_images
+
+
+def draw_score(img, score):
+    img = np.transpose(img, (1, 2, 0))
+    cv2.putText(img, f'{score:.2f}',
+                (40, 40),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                1,
+                (0, 1, 0),
+                thickness=3)
+    img = np.transpose(img, (2, 0, 1))
+    return img
+
+
+def save_global_samples(global_mask_fnames, mask2real_fname, mask2fake_fname, out_dir, real_scores_by_fname, fake_scores_by_fname):
+    for cur_mask_fname in global_mask_fnames:
+        cur_real_fname = mask2real_fname[cur_mask_fname]
+        orig_img = load_image(cur_real_fname, mode='RGB')
+        fake_img = load_image(mask2fake_fname[cur_mask_fname], mode='RGB')[:, :orig_img.shape[1], :orig_img.shape[2]]
+        mask = load_image(cur_mask_fname, mode='L')[None, ...]
+
+        draw_score(orig_img, real_scores_by_fname.loc[cur_real_fname, 'real_score'])
+        draw_score(fake_img, fake_scores_by_fname.loc[cur_mask_fname, 'fake_score'])
+
+        cur_grid = visualize_mask_and_images(dict(image=orig_img, mask=mask, fake=fake_img),
+                                             keys=['image', 'fake'],
+                                             last_without_mask=True)
+        cur_grid = np.clip(cur_grid * 255, 0, 255).astype('uint8')
+        cur_grid = cv2.cvtColor(cur_grid, cv2.COLOR_RGB2BGR)
+        cv2.imwrite(os.path.join(out_dir, os.path.splitext(os.path.basename(cur_mask_fname))[0] + '.jpg'),
+                    cur_grid)
+
+
+def save_samples_by_real(worst_best_by_real, mask2fake_fname, fake_info, out_dir):
+    for real_fname in worst_best_by_real.index:
+        worst_mask_path = worst_best_by_real.loc[real_fname, 'worst']
+        best_mask_path = worst_best_by_real.loc[real_fname, 'best']
+        orig_img = load_image(real_fname, mode='RGB')
+        worst_mask_img = load_image(worst_mask_path, mode='L')[None, ...]
+        worst_fake_img = load_image(mask2fake_fname[worst_mask_path], mode='RGB')[:, :orig_img.shape[1], :orig_img.shape[2]]
+        best_mask_img = load_image(best_mask_path, mode='L')[None, ...]
+        best_fake_img = load_image(mask2fake_fname[best_mask_path], mode='RGB')[:, :orig_img.shape[1], :orig_img.shape[2]]
+
+        draw_score(orig_img, worst_best_by_real.loc[real_fname, 'real_score'])
+        draw_score(worst_fake_img, worst_best_by_real.loc[real_fname, 'worst_score'])
+        draw_score(best_fake_img, worst_best_by_real.loc[real_fname, 'best_score'])
+
+        cur_grid = visualize_mask_and_images(dict(image=orig_img, mask=np.zeros_like(worst_mask_img),
+                                                  worst_mask=worst_mask_img, worst_img=worst_fake_img,
+                                                  best_mask=best_mask_img, best_img=best_fake_img),
+                                             keys=['image', 'worst_mask', 'worst_img', 'best_mask', 'best_img'],
+                                             rescale_keys=['worst_mask', 'best_mask'],
+                                             last_without_mask=True)
+        cur_grid = np.clip(cur_grid * 255, 0, 255).astype('uint8')
+        cur_grid = cv2.cvtColor(cur_grid, cv2.COLOR_RGB2BGR)
+        cv2.imwrite(os.path.join(out_dir,
+                                 os.path.splitext(os.path.basename(real_fname))[0] + '.jpg'),
+                    cur_grid)
+
+        fig, (ax1, ax2) = plt.subplots(1, 2)
+        cur_stat = fake_info[fake_info['real_fname'] == real_fname]
+        cur_stat['fake_score'].hist(ax=ax1)
+        cur_stat['real_score'].hist(ax=ax2)
+        fig.tight_layout()
+        fig.savefig(os.path.join(out_dir,
+                                 os.path.splitext(os.path.basename(real_fname))[0] + '_scores.png'))
+        plt.close(fig)
+
+
+def extract_overlapping_masks(mask_fnames, cur_i, fake_scores_table, max_overlaps_n=2):
+    result_pairs = []
+    result_scores = []
+    mask_fname_a = mask_fnames[cur_i]
+    mask_a = load_image(mask_fname_a, mode='L')[None, ...] > 0.5
+    cur_score_a = fake_scores_table.loc[mask_fname_a, 'fake_score']
+    for mask_fname_b in mask_fnames[cur_i + 1:]:
+        mask_b = load_image(mask_fname_b, mode='L')[None, ...] > 0.5
+        if not np.any(mask_a & mask_b):
+            continue
+        cur_score_b = fake_scores_table.loc[mask_fname_b, 'fake_score']
+        result_pairs.append((mask_fname_a, mask_fname_b))
+        result_scores.append(cur_score_b - cur_score_a)
+        if len(result_pairs) >= max_overlaps_n:
+            break
+    return result_pairs, result_scores
+
+
+def main(args):
+    config = load_yaml(args.config)
+
+    latents_dir = os.path.join(args.outpath, 'latents')
+    os.makedirs(latents_dir, exist_ok=True)
+    global_worst_dir = os.path.join(args.outpath, 'global_worst')
+    os.makedirs(global_worst_dir, exist_ok=True)
+    global_best_dir = os.path.join(args.outpath, 'global_best')
+    os.makedirs(global_best_dir, exist_ok=True)
+    worst_best_by_best_worst_score_diff_max_dir = os.path.join(args.outpath, 'worst_best_by_real', 'best_worst_score_diff_max')
+    os.makedirs(worst_best_by_best_worst_score_diff_max_dir, exist_ok=True)
+    worst_best_by_best_worst_score_diff_min_dir = os.path.join(args.outpath, 'worst_best_by_real', 'best_worst_score_diff_min')
+    os.makedirs(worst_best_by_best_worst_score_diff_min_dir, exist_ok=True)
+    worst_best_by_real_best_score_diff_max_dir = os.path.join(args.outpath, 'worst_best_by_real', 'real_best_score_diff_max')
+    os.makedirs(worst_best_by_real_best_score_diff_max_dir, exist_ok=True)
+    worst_best_by_real_best_score_diff_min_dir = os.path.join(args.outpath, 'worst_best_by_real', 'real_best_score_diff_min')
+    os.makedirs(worst_best_by_real_best_score_diff_min_dir, exist_ok=True)
+    worst_best_by_real_worst_score_diff_max_dir = os.path.join(args.outpath, 'worst_best_by_real', 'real_worst_score_diff_max')
+    os.makedirs(worst_best_by_real_worst_score_diff_max_dir, exist_ok=True)
+    worst_best_by_real_worst_score_diff_min_dir = os.path.join(args.outpath, 'worst_best_by_real', 'real_worst_score_diff_min')
+    os.makedirs(worst_best_by_real_worst_score_diff_min_dir, exist_ok=True)
+
+    if not args.only_report:
+        block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048]
+        inception_model = InceptionV3([block_idx]).eval().cuda()
+
+        dataset = PrecomputedInpaintingResultsDataset(args.datadir, args.predictdir, **config.dataset_kwargs)
+
+        real2vector_cache = {}
+
+        real_features = []
+        fake_features = []
+
+        orig_fnames = []
+        mask_fnames = []
+        mask2real_fname = {}
+        mask2fake_fname = {}
+
+        for batch_i, batch in enumerate(dataset):
+            orig_img_fname = dataset.img_filenames[batch_i]
+            mask_fname = dataset.mask_filenames[batch_i]
+            fake_fname = dataset.pred_filenames[batch_i]
+            mask2real_fname[mask_fname] = orig_img_fname
+            mask2fake_fname[mask_fname] = fake_fname
+
+            cur_real_vector = real2vector_cache.get(orig_img_fname, None)
+            if cur_real_vector is None:
+                with torch.no_grad():
+                    in_img = torch.from_numpy(batch['image'][None, ...]).cuda()
+                    cur_real_vector = inception_model(in_img)[0].squeeze(-1).squeeze(-1).cpu().numpy()
+                real2vector_cache[orig_img_fname] = cur_real_vector
+
+            pred_img = torch.from_numpy(batch['inpainted'][None, ...]).cuda()
+            cur_fake_vector = inception_model(pred_img)[0].squeeze(-1).squeeze(-1).cpu().numpy()
+
+            real_features.append(cur_real_vector)
+            fake_features.append(cur_fake_vector)
+
+            orig_fnames.append(orig_img_fname)
+            mask_fnames.append(mask_fname)
+
+        ids_features = np.concatenate(real_features + fake_features, axis=0)
+        ids_labels = np.array(([1] * len(real_features)) + ([0] * len(fake_features)))
+
+        with open(os.path.join(latents_dir, 'featues.pkl'), 'wb') as f:
+            pickle.dump(ids_features, f, protocol=3)
+        with open(os.path.join(latents_dir, 'labels.pkl'), 'wb') as f:
+            pickle.dump(ids_labels, f, protocol=3)
+        with open(os.path.join(latents_dir, 'orig_fnames.pkl'), 'wb') as f:
+            pickle.dump(orig_fnames, f, protocol=3)
+        with open(os.path.join(latents_dir, 'mask_fnames.pkl'), 'wb') as f:
+            pickle.dump(mask_fnames, f, protocol=3)
+        with open(os.path.join(latents_dir, 'mask2real_fname.pkl'), 'wb') as f:
+            pickle.dump(mask2real_fname, f, protocol=3)
+        with open(os.path.join(latents_dir, 'mask2fake_fname.pkl'), 'wb') as f:
+            pickle.dump(mask2fake_fname, f, protocol=3)
+
+        svm = sklearn.svm.LinearSVC(dual=False)
+        svm.fit(ids_features, ids_labels)
+
+        pred_scores = svm.decision_function(ids_features)
+        real_scores = pred_scores[:len(real_features)]
+        fake_scores = pred_scores[len(real_features):]
+
+        with open(os.path.join(latents_dir, 'pred_scores.pkl'), 'wb') as f:
+            pickle.dump(pred_scores, f, protocol=3)
+        with open(os.path.join(latents_dir, 'real_scores.pkl'), 'wb') as f:
+            pickle.dump(real_scores, f, protocol=3)
+        with open(os.path.join(latents_dir, 'fake_scores.pkl'), 'wb') as f:
+            pickle.dump(fake_scores, f, protocol=3)
+    else:
+        with open(os.path.join(latents_dir, 'orig_fnames.pkl'), 'rb') as f:
+            orig_fnames = pickle.load(f)
+        with open(os.path.join(latents_dir, 'mask_fnames.pkl'), 'rb') as f:
+            mask_fnames = pickle.load(f)
+        with open(os.path.join(latents_dir, 'mask2real_fname.pkl'), 'rb') as f:
+            mask2real_fname = pickle.load(f)
+        with open(os.path.join(latents_dir, 'mask2fake_fname.pkl'), 'rb') as f:
+            mask2fake_fname = pickle.load(f)
+        with open(os.path.join(latents_dir, 'real_scores.pkl'), 'rb') as f:
+            real_scores = pickle.load(f)
+        with open(os.path.join(latents_dir, 'fake_scores.pkl'), 'rb') as f:
+            fake_scores = pickle.load(f)
+
+    real_info = pd.DataFrame(data=[dict(real_fname=fname,
+                                        real_score=score)
+                                   for fname, score
+                                   in zip(orig_fnames, real_scores)])
+    real_info.set_index('real_fname', drop=True, inplace=True)
+
+    fake_info = pd.DataFrame(data=[dict(mask_fname=fname,
+                                        fake_fname=mask2fake_fname[fname],
+                                        real_fname=mask2real_fname[fname],
+                                        fake_score=score)
+                                   for fname, score
+                                   in zip(mask_fnames, fake_scores)])
+    fake_info = fake_info.join(real_info, on='real_fname', how='left')
+    fake_info.drop_duplicates(['fake_fname', 'real_fname'], inplace=True)
+
+    fake_stats_by_real = fake_info.groupby('real_fname')['fake_score'].describe()[['mean', 'std']].rename(
+        {'mean': 'mean_fake_by_real', 'std': 'std_fake_by_real'}, axis=1)
+    fake_info = fake_info.join(fake_stats_by_real, on='real_fname', rsuffix='stat_by_real')
+    fake_info.drop_duplicates(['fake_fname', 'real_fname'], inplace=True)
+    fake_info.to_csv(os.path.join(latents_dir, 'join_scores_table.csv'), sep='\t', index=False)
+
+    fake_scores_table = fake_info.set_index('mask_fname')['fake_score'].to_frame()
+    real_scores_table = fake_info.set_index('real_fname')['real_score'].drop_duplicates().to_frame()
+
+    fig, (ax1, ax2) = plt.subplots(1, 2)
+    ax1.hist(fake_scores)
+    ax2.hist(real_scores)
+    fig.tight_layout()
+    fig.savefig(os.path.join(args.outpath, 'global_scores_hist.png'))
+    plt.close(fig)
+
+    global_worst_masks = fake_info.sort_values('fake_score', ascending=True)['mask_fname'].iloc[:config.take_global_top].to_list()
+    global_best_masks = fake_info.sort_values('fake_score', ascending=False)['mask_fname'].iloc[:config.take_global_top].to_list()
+    save_global_samples(global_worst_masks, mask2real_fname, mask2fake_fname, global_worst_dir, real_scores_table, fake_scores_table)
+    save_global_samples(global_best_masks, mask2real_fname, mask2fake_fname, global_best_dir, real_scores_table, fake_scores_table)
+
+    # grouped by real
+    worst_samples_by_real = fake_info.groupby('real_fname').apply(
+        lambda d: d.set_index('mask_fname')['fake_score'].idxmin()).to_frame().rename({0: 'worst'}, axis=1)
+    best_samples_by_real = fake_info.groupby('real_fname').apply(
+        lambda d: d.set_index('mask_fname')['fake_score'].idxmax()).to_frame().rename({0: 'best'}, axis=1)
+    worst_best_by_real = pd.concat([worst_samples_by_real, best_samples_by_real], axis=1)
+
+    worst_best_by_real = worst_best_by_real.join(fake_scores_table.rename({'fake_score': 'worst_score'}, axis=1),
+                                                 on='worst')
+    worst_best_by_real = worst_best_by_real.join(fake_scores_table.rename({'fake_score': 'best_score'}, axis=1),
+                                                 on='best')
+    worst_best_by_real = worst_best_by_real.join(real_scores_table)
+
+    worst_best_by_real['best_worst_score_diff'] = worst_best_by_real['best_score'] - worst_best_by_real['worst_score']
+    worst_best_by_real['real_best_score_diff'] = worst_best_by_real['real_score'] - worst_best_by_real['best_score']
+    worst_best_by_real['real_worst_score_diff'] = worst_best_by_real['real_score'] - worst_best_by_real['worst_score']
+
+    worst_best_by_best_worst_score_diff_min = worst_best_by_real.sort_values('best_worst_score_diff', ascending=True).iloc[:config.take_worst_best_top]
+    worst_best_by_best_worst_score_diff_max = worst_best_by_real.sort_values('best_worst_score_diff', ascending=False).iloc[:config.take_worst_best_top]
+    save_samples_by_real(worst_best_by_best_worst_score_diff_min, mask2fake_fname, fake_info, worst_best_by_best_worst_score_diff_min_dir)
+    save_samples_by_real(worst_best_by_best_worst_score_diff_max, mask2fake_fname, fake_info, worst_best_by_best_worst_score_diff_max_dir)
+
+    worst_best_by_real_best_score_diff_min = worst_best_by_real.sort_values('real_best_score_diff', ascending=True).iloc[:config.take_worst_best_top]
+    worst_best_by_real_best_score_diff_max = worst_best_by_real.sort_values('real_best_score_diff', ascending=False).iloc[:config.take_worst_best_top]
+    save_samples_by_real(worst_best_by_real_best_score_diff_min, mask2fake_fname, fake_info, worst_best_by_real_best_score_diff_min_dir)
+    save_samples_by_real(worst_best_by_real_best_score_diff_max, mask2fake_fname, fake_info, worst_best_by_real_best_score_diff_max_dir)
+
+    worst_best_by_real_worst_score_diff_min = worst_best_by_real.sort_values('real_worst_score_diff', ascending=True).iloc[:config.take_worst_best_top]
+    worst_best_by_real_worst_score_diff_max = worst_best_by_real.sort_values('real_worst_score_diff', ascending=False).iloc[:config.take_worst_best_top]
+    save_samples_by_real(worst_best_by_real_worst_score_diff_min, mask2fake_fname, fake_info, worst_best_by_real_worst_score_diff_min_dir)
+    save_samples_by_real(worst_best_by_real_worst_score_diff_max, mask2fake_fname, fake_info, worst_best_by_real_worst_score_diff_max_dir)
+
+    # analyze what change of mask causes bigger change of score
+    overlapping_mask_fname_pairs = []
+    overlapping_mask_fname_score_diffs = []
+    for cur_real_fname in orig_fnames:
+        cur_fakes_info = fake_info[fake_info['real_fname'] == cur_real_fname]
+        cur_mask_fnames = sorted(cur_fakes_info['mask_fname'].unique())
+
+        cur_mask_pairs_and_scores = Parallel(args.n_jobs)(
+            delayed(extract_overlapping_masks)(cur_mask_fnames, i, fake_scores_table)
+            for i in range(len(cur_mask_fnames) - 1)
+        )
+        for cur_pairs, cur_scores in cur_mask_pairs_and_scores:
+            overlapping_mask_fname_pairs.extend(cur_pairs)
+            overlapping_mask_fname_score_diffs.extend(cur_scores)
+
+    overlapping_mask_fname_pairs = np.asarray(overlapping_mask_fname_pairs)
+    overlapping_mask_fname_score_diffs = np.asarray(overlapping_mask_fname_score_diffs)
+    overlapping_sort_idx = np.argsort(overlapping_mask_fname_score_diffs)
+    overlapping_mask_fname_pairs = overlapping_mask_fname_pairs[overlapping_sort_idx]
+    overlapping_mask_fname_score_diffs = overlapping_mask_fname_score_diffs[overlapping_sort_idx]
+
+
+
+
+
+
+if __name__ == '__main__':
+    import argparse
+
+    aparser = argparse.ArgumentParser()
+    aparser.add_argument('config', type=str, help='Path to config for dataset generation')
+    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('--only-report', action='store_true',
+                         help='Whether to skip prediction and feature extraction, '
+                              'load all the possible latents and proceed with report only')
+    aparser.add_argument('--n-jobs', type=int, default=8, help='how many processes to use for pair mask mining')
+
+    main(aparser.parse_args())

bin/blur_predicts.py

@@ -0,0 +1,57 @@
+#!/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())

bin/calc_dataset_stats.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+
+import os
+
+import numpy as np
+import tqdm
+from scipy.ndimage.morphology import distance_transform_edt
+
+from saicinpainting.evaluation.data import InpaintingDataset
+from saicinpainting.evaluation.vis import save_item_for_vis
+
+
+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 = []
+    known_pixel_distances = []
+
+    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)
+        bin_mask = item['mask'] > 0.5
+        hole_area = bin_mask.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)
+
+        cur_dist = distance_transform_edt(bin_mask)
+        cur_dist_inside_mask = cur_dist[bin_mask]
+        known_pixel_distances.append(cur_dist_inside_mask.mean())
+
+    os.makedirs(args.outdir, exist_ok=True)
+    with open(os.path.join(args.outdir, 'summary.txt'), 'w') as f:
+        f.write(f'''Location:          {args.datadir}
+
+Number of samples: {len(dataset)}
+
+Image height: min {min(heights):5d} max {max(heights):5d} mean {np.mean(heights):.2f}
+Image width:  min {min(widths):5d} max {max(widths):5d} mean {np.mean(widths):.2f}
+Image area:   min {min(image_areas):7d} max {max(image_areas):7d} mean {np.mean(image_areas):.2f}
+Hole area:    min {min(hole_areas):7d} max {max(hole_areas):7d} mean {np.mean(hole_areas):.2f}
+Hole area %:  min {min(hole_area_percents) * 100:2.2f} max {max(hole_area_percents) * 100:2.2f} mean {np.mean(hole_area_percents) * 100:2.2f}
+Dist 2known:  min {min(known_pixel_distances):2.2f} max {max(known_pixel_distances):2.2f} mean {np.mean(known_pixel_distances):2.2f} median {np.median(known_pixel_distances):2.2f}
+
+Stats by hole area %:
+''')
+        for bin_i in range(args.area_bins):
+            f.write(f'{area_bin_titles[bin_i]}%: '
+                    f'samples number {area_bins_count[bin_i]}, '
+                    f'{area_bins_count[bin_i] / len(dataset) * 100:.1f}%\n')
+
+    for bin_i in range(args.area_bins):
+        bindir = os.path.join(args.outdir, 'samples', 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):
+            save_item_for_vis(dataset[sample_i], os.path.join(bindir, f'{sample_i}.png'))
+
+
+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())

bin/debug/analyze_overlapping_masks.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+
+BASEDIR="$(dirname $0)"
+
+# paths are valid for mml7
+
+# select images
+#ls /data/inpainting/work/data/train | shuf | head -2000 | xargs -n1 -I{} cp {} /data/inpainting/mask_analysis/src
+
+# generate masks
+#"$BASEDIR/../gen_debug_mask_dataset.py" \
+#    "$BASEDIR/../../configs/debug_mask_gen.yaml" \
+#    "/data/inpainting/mask_analysis/src" \
+#    "/data/inpainting/mask_analysis/generated"
+
+# predict
+#"$BASEDIR/../predict.py" \
+#    model.path="simple_pix2pix2_gap_sdpl_novgg_large_b18_ffc075_batch8x15/saved_checkpoint/r.suvorov_2021-04-30_14-41-12_train_simple_pix2pix2_gap_sdpl_novgg_large_b18_ffc075_batch8x15_epoch22-step-574999" \
+#    indir="/data/inpainting/mask_analysis/generated" \
+#    outdir="/data/inpainting/mask_analysis/predicted" \
+#    dataset.img_suffix=.jpg \
+#    +out_ext=.jpg
+
+# analyze good and bad samples
+"$BASEDIR/../analyze_errors.py" \
+    --only-report \
+    --n-jobs 8 \
+    "$BASEDIR/../../configs/analyze_mask_errors.yaml" \
+    "/data/inpainting/mask_analysis/small/generated" \
+    "/data/inpainting/mask_analysis/small/predicted" \
+    "/data/inpainting/mask_analysis/small/report"

bin/evaluate_predicts.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+
+import os
+
+import pandas as pd
+
+from saicinpainting.evaluation.data import PrecomputedInpaintingResultsDataset
+from saicinpainting.evaluation.evaluator import InpaintingEvaluator, lpips_fid100_f1
+from saicinpainting.evaluation.losses.base_loss import SegmentationAwareSSIM, \
+    SegmentationClassStats, SSIMScore, LPIPSScore, FIDScore, SegmentationAwareLPIPS, SegmentationAwareFID
+from saicinpainting.evaluation.utils import load_yaml
+
+
+def main(args):
+    config = load_yaml(args.config)
+
+    dataset = PrecomputedInpaintingResultsDataset(args.datadir, args.predictdir, **config.dataset_kwargs)
+
+    metrics = {
+        'ssim': SSIMScore(),
+        'lpips': LPIPSScore(),
+        'fid': FIDScore()
+    }
+    enable_segm = config.get('segmentation', dict(enable=False)).get('enable', False)
+    if enable_segm:
+        weights_path = os.path.expandvars(config.segmentation.weights_path)
+        metrics.update(dict(
+            segm_stats=SegmentationClassStats(weights_path=weights_path),
+            segm_ssim=SegmentationAwareSSIM(weights_path=weights_path),
+            segm_lpips=SegmentationAwareLPIPS(weights_path=weights_path),
+            segm_fid=SegmentationAwareFID(weights_path=weights_path)
+        ))
+    evaluator = InpaintingEvaluator(dataset, scores=metrics,
+                                    integral_title='lpips_fid100_f1', integral_func=lpips_fid100_f1,
+                                    **config.evaluator_kwargs)
+
+    os.makedirs(os.path.dirname(args.outpath), exist_ok=True)
+
+    results = evaluator.evaluate()
+
+    results = pd.DataFrame(results).stack(1).unstack(0)
+    results.dropna(axis=1, how='all', inplace=True)
+    results.to_csv(args.outpath, sep='\t', float_format='%.4f')
+
+    if enable_segm:
+        only_short_results = results[[c for c in results.columns if not c[0].startswith('segm_')]].dropna(axis=1, how='all')
+        only_short_results.to_csv(args.outpath + '_short', sep='\t', float_format='%.4f')
+
+        print(only_short_results)
+
+        segm_metrics_results = results[['segm_ssim', 'segm_lpips', 'segm_fid']].dropna(axis=1, how='all').transpose().unstack(0).reorder_levels([1, 0], axis=1)
+        segm_metrics_results.drop(['mean', 'std'], axis=0, inplace=True)
+
+        segm_stats_results = results['segm_stats'].dropna(axis=1, how='all').transpose()
+        segm_stats_results.index = pd.MultiIndex.from_tuples(n.split('/') for n in segm_stats_results.index)
+        segm_stats_results = segm_stats_results.unstack(0).reorder_levels([1, 0], axis=1)
+        segm_stats_results.sort_index(axis=1, inplace=True)
+        segm_stats_results.dropna(axis=0, how='all', inplace=True)
+
+        segm_results = pd.concat([segm_metrics_results, segm_stats_results], axis=1, sort=True)
+        segm_results.sort_values(('mask_freq', 'total'), ascending=False, inplace=True)
+
+        segm_results.to_csv(args.outpath + '_segm', sep='\t', float_format='%.4f')
+    else:
+        print(results)
+
+
+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')
+
+    main(aparser.parse_args())

bin/evaluator_example.py

@@ -0,0 +1,76 @@
+import os
+
+import cv2
+import numpy as np
+import torch
+from skimage import io
+from skimage.transform import resize
+from torch.utils.data import Dataset
+
+from saicinpainting.evaluation.evaluator import InpaintingEvaluator
+from saicinpainting.evaluation.losses.base_loss import SSIMScore, LPIPSScore, FIDScore
+
+
+class SimpleImageDataset(Dataset):
+    def __init__(self, root_dir, image_size=(400, 600)):
+        self.root_dir = root_dir
+        self.files = sorted(os.listdir(root_dir))
+        self.image_size = image_size
+
+    def __getitem__(self, index):
+        img_name = os.path.join(self.root_dir, self.files[index])
+        image = io.imread(img_name)
+        image = resize(image, self.image_size, anti_aliasing=True)
+        image = torch.FloatTensor(image).permute(2, 0, 1)
+        return image
+
+    def __len__(self):
+        return len(self.files)
+
+
+def create_rectangle_mask(height, width):
+    mask = np.ones((height, width))
+    up_left_corner = width // 4, height // 4
+    down_right_corner = (width - up_left_corner[0] - 1, height - up_left_corner[1] - 1)
+    cv2.rectangle(mask, up_left_corner, down_right_corner, (0, 0, 0), thickness=cv2.FILLED)
+    return mask
+
+
+class Model():
+    def __call__(self, img_batch, mask_batch):
+        mean = (img_batch * mask_batch[:, None, :, :]).sum(dim=(2, 3)) / mask_batch.sum(dim=(1, 2))[:, None]
+        inpainted = mean[:, :, None, None] * (1 - mask_batch[:, None, :, :]) + img_batch * mask_batch[:, None, :, :]
+        return inpainted
+
+
+class SimpleImageSquareMaskDataset(Dataset):
+    def __init__(self, dataset):
+        self.dataset = dataset
+        self.mask = torch.FloatTensor(create_rectangle_mask(*self.dataset.image_size))
+        self.model = Model()
+
+    def __getitem__(self, index):
+        img = self.dataset[index]
+        mask = self.mask.clone()
+        inpainted = self.model(img[None, ...], mask[None, ...])
+        return dict(image=img, mask=mask, inpainted=inpainted)
+
+    def __len__(self):
+        return len(self.dataset)
+
+
+dataset = SimpleImageDataset('imgs')
+mask_dataset = SimpleImageSquareMaskDataset(dataset)
+model = Model()
+metrics = {
+    'ssim': SSIMScore(),
+    'lpips': LPIPSScore(),
+    'fid': FIDScore()
+}
+
+evaluator = InpaintingEvaluator(
+    mask_dataset, scores=metrics, batch_size=3, area_grouping=True
+)
+
+results = evaluator.evaluate(model)
+print(results)

bin/extract_masks.py

@@ -0,0 +1,63 @@
+import PIL.Image as Image
+import numpy as np
+import os
+
+
+def main(args):
+    if not args.indir.endswith('/'):
+        args.indir += '/'
+    os.makedirs(args.outdir, exist_ok=True)
+
+    src_images = [
+        args.indir+fname for fname in  os.listdir(args.indir)]
+
+    tgt_masks = [
+        args.outdir+fname[:-4] + f'_mask000.png' 
+            for fname in  os.listdir(args.indir)]
+
+    for img_name, msk_name in zip(src_images, tgt_masks):
+        #print(img)
+        #print(msk)
+
+        image = Image.open(img_name).convert('RGB')
+        image = np.transpose(np.array(image), (2, 0, 1))
+
+        mask = (image == 255).astype(int)
+
+        print(mask.dtype, mask.shape)
+
+
+        Image.fromarray(
+            np.clip(mask[0,:,:] * 255, 0, 255).astype('uint8'),mode='L'
+        ).save(msk_name)
+
+
+
+
+    '''
+    for infile in src_images:
+        try:
+            file_relpath = infile[len(indir):]
+            img_outpath = os.path.join(outdir, file_relpath)
+            os.makedirs(os.path.dirname(img_outpath), exist_ok=True)
+
+            image = Image.open(infile).convert('RGB')
+
+            mask = 
+
+            Image.fromarray(
+                np.clip(
+                    cur_mask * 255, 0, 255).astype('uint8'),
+                    mode='L'
+                ).save(cur_basename + f'_mask{i:03d}.png')
+    '''
+
+
+
+if __name__ == '__main__':
+    import argparse
+    aparser = argparse.ArgumentParser()
+    aparser.add_argument('--indir', type=str, help='Path to folder with images')
+    aparser.add_argument('--outdir', type=str, help='Path to folder to store aligned images and masks to')
+
+    main(aparser.parse_args())

bin/filter_sharded_dataset.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+
+
+import math
+import os
+import random
+
+import braceexpand
+import webdataset as wds
+
+DEFAULT_CATS_FILE = os.path.join(os.path.dirname(__file__), '..', 'configs', 'places2-categories_157.txt')
+
+def is_good_key(key, cats):
+    return any(c in key for c in cats)
+
+
+def main(args):
+    if args.categories == 'nofilter':
+        good_categories = None
+    else:
+        with open(args.categories, 'r') as f:
+            good_categories = set(line.strip().split(' ')[0] for line in f if line.strip())
+
+    all_input_files = list(braceexpand.braceexpand(args.infile))
+    chunk_size = int(math.ceil(len(all_input_files) / args.n_read_streams))
+
+    input_iterators = [iter(wds.Dataset(all_input_files[start : start + chunk_size]).shuffle(args.shuffle_buffer))
+                       for start in range(0, len(all_input_files), chunk_size)]
+    output_datasets = [wds.ShardWriter(args.outpattern.format(i)) for i in range(args.n_write_streams)]
+
+    good_readers = list(range(len(input_iterators)))
+    step_i = 0
+    good_samples = 0
+    bad_samples = 0
+    while len(good_readers) > 0:
+        if step_i % args.print_freq == 0:
+            print(f'Iterations done {step_i}; readers alive {good_readers}; good samples {good_samples}; bad samples {bad_samples}')
+
+        step_i += 1
+
+        ri = random.choice(good_readers)
+        try:
+            sample = next(input_iterators[ri])
+        except StopIteration:
+            good_readers = list(set(good_readers) - {ri})
+            continue
+
+        if good_categories is not None and not is_good_key(sample['__key__'], good_categories):
+            bad_samples += 1
+            continue
+
+        wi = random.randint(0, args.n_write_streams - 1)
+        output_datasets[wi].write(sample)
+        good_samples += 1
+
+
+if __name__ == '__main__':
+    import argparse
+
+    aparser = argparse.ArgumentParser()
+    aparser.add_argument('--categories', type=str, default=DEFAULT_CATS_FILE)
+    aparser.add_argument('--shuffle-buffer', type=int, default=10000)
+    aparser.add_argument('--n-read-streams', type=int, default=10)
+    aparser.add_argument('--n-write-streams', type=int, default=10)
+    aparser.add_argument('--print-freq', type=int, default=1000)
+    aparser.add_argument('infile', type=str)
+    aparser.add_argument('outpattern', type=str)
+
+    main(aparser.parse_args())

bin/gen_debug_mask_dataset.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+
+import glob
+import os
+
+import PIL.Image as Image
+import cv2
+import numpy as np
+import tqdm
+import shutil
+
+
+from saicinpainting.evaluation.utils import load_yaml
+
+
+def generate_masks_for_img(infile, outmask_pattern, mask_size=200, step=0.5):
+    inimg = Image.open(infile)
+    width, height = inimg.size
+    step_abs = int(mask_size * step)
+
+    mask = np.zeros((height, width), dtype='uint8')
+    mask_i = 0
+
+    for start_vertical in range(0, height - step_abs, step_abs):
+        for start_horizontal in range(0, width - step_abs, step_abs):
+            mask[start_vertical:start_vertical + mask_size, start_horizontal:start_horizontal + mask_size] = 255
+
+            cv2.imwrite(outmask_pattern.format(mask_i), mask)
+
+            mask[start_vertical:start_vertical + mask_size, start_horizontal:start_horizontal + mask_size] = 0
+            mask_i += 1
+
+
+def main(args):
+    if not args.indir.endswith('/'):
+        args.indir += '/'
+    if not args.outdir.endswith('/'):
+        args.outdir += '/'
+
+    config = load_yaml(args.config)
+
+    in_files = list(glob.glob(os.path.join(args.indir, '**', f'*{config.img_ext}'), recursive=True))
+    for infile in tqdm.tqdm(in_files):
+        outimg = args.outdir + infile[len(args.indir):]
+        outmask_pattern = outimg[:-len(config.img_ext)] + '_mask{:04d}.png'
+
+        os.makedirs(os.path.dirname(outimg), exist_ok=True)
+        shutil.copy2(infile, outimg)
+
+        generate_masks_for_img(infile, outmask_pattern, **config.gen_kwargs)
+
+
+if __name__ == '__main__':
+    import argparse
+
+    aparser = argparse.ArgumentParser()
+    aparser.add_argument('config', type=str, help='Path to config for dataset generation')
+    aparser.add_argument('indir', type=str, help='Path to folder with images')
+    aparser.add_argument('outdir', type=str, help='Path to folder to store aligned images and masks to')
+
+    main(aparser.parse_args())

bin/gen_mask_dataset.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+
+import glob
+import os
+import shutil
+import traceback
+
+import PIL.Image as Image
+import numpy as np
+from joblib import Parallel, delayed
+
+from saicinpainting.evaluation.masks.mask import SegmentationMask, propose_random_square_crop
+from saicinpainting.evaluation.utils import load_yaml, SmallMode
+from saicinpainting.training.data.masks import MixedMaskGenerator
+
+
+class MakeManyMasksWrapper:
+    def __init__(self, impl, variants_n=2):
+        self.impl = impl
+        self.variants_n = variants_n
+
+    def get_masks(self, img):
+        img = np.transpose(np.array(img), (2, 0, 1))
+        return [self.impl(img)[0] for _ in range(self.variants_n)]
+
+
+def process_images(src_images, indir, outdir, config):
+    if config.generator_kind == 'segmentation':
+        mask_generator = SegmentationMask(**config.mask_generator_kwargs)
+    elif config.generator_kind == 'random':
+        variants_n = config.mask_generator_kwargs.pop('variants_n', 2)
+        mask_generator = MakeManyMasksWrapper(MixedMaskGenerator(**config.mask_generator_kwargs),
+                                              variants_n=variants_n)
+    else:
+        raise ValueError(f'Unexpected generator kind: {config.generator_kind}')
+
+    max_tamper_area = config.get('max_tamper_area', 1)
+
+    for infile in src_images:
+        try:
+            file_relpath = infile[len(indir):]
+            img_outpath = os.path.join(outdir, file_relpath)
+            os.makedirs(os.path.dirname(img_outpath), exist_ok=True)
+
+            image = Image.open(infile).convert('RGB')
+
+            # scale input image to output resolution and filter smaller images
+            if min(image.size) < config.cropping.out_min_size:
+                handle_small_mode = SmallMode(config.cropping.handle_small_mode)
+                if handle_small_mode == SmallMode.DROP:
+                    continue
+                elif handle_small_mode == SmallMode.UPSCALE:
+                    factor = config.cropping.out_min_size / min(image.size)
+                    out_size = (np.array(image.size) * factor).round().astype('uint32')
+                    image = image.resize(out_size, resample=Image.BICUBIC)
+            else:
+                factor = config.cropping.out_min_size / min(image.size)
+                out_size = (np.array(image.size) * factor).round().astype('uint32')
+                image = image.resize(out_size, resample=Image.BICUBIC)
+
+            # generate and select masks
+            src_masks = mask_generator.get_masks(image)
+
+            filtered_image_mask_pairs = []
+            for cur_mask in src_masks:
+                if config.cropping.out_square_crop:
+                    (crop_left,
+                     crop_top,
+                     crop_right,
+                     crop_bottom) = propose_random_square_crop(cur_mask,
+                                                               min_overlap=config.cropping.crop_min_overlap)
+                    cur_mask = cur_mask[crop_top:crop_bottom, crop_left:crop_right]
+                    cur_image = image.copy().crop((crop_left, crop_top, crop_right, crop_bottom))
+                else:
+                    cur_image = image
+
+                if len(np.unique(cur_mask)) == 0 or cur_mask.mean() > max_tamper_area:
+                    continue
+
+                filtered_image_mask_pairs.append((cur_image, cur_mask))
+
+            mask_indices = np.random.choice(len(filtered_image_mask_pairs),
+                                            size=min(len(filtered_image_mask_pairs), config.max_masks_per_image),
+                                            replace=False)
+
+            # crop masks; save masks together with input image
+            mask_basename = os.path.join(outdir, os.path.splitext(file_relpath)[0])
+            for i, idx in enumerate(mask_indices):
+                cur_image, cur_mask = filtered_image_mask_pairs[idx]
+                cur_basename = mask_basename + f'_crop{i:03d}'
+                Image.fromarray(np.clip(cur_mask * 255, 0, 255).astype('uint8'),
+                                mode='L').save(cur_basename + f'_mask{i:03d}.png')
+                cur_image.save(cur_basename + '.png')
+        except KeyboardInterrupt:
+            return
+        except Exception as ex:
+            print(f'Could not make masks for {infile} due to {ex}:\n{traceback.format_exc()}')
+
+
+def main(args):
+    if not args.indir.endswith('/'):
+        args.indir += '/'
+
+    os.makedirs(args.outdir, exist_ok=True)
+
+    config = load_yaml(args.config)
+
+    in_files = list(glob.glob(os.path.join(args.indir, '**', f'*.{args.ext}'), recursive=True))
+    if args.n_jobs == 0:
+        process_images(in_files, args.indir, args.outdir, config)
+    else:
+        in_files_n = len(in_files)
+        chunk_size = in_files_n // args.n_jobs + (1 if in_files_n % args.n_jobs > 0 else 0)
+        Parallel(n_jobs=args.n_jobs)(
+            delayed(process_images)(in_files[start:start+chunk_size], args.indir, args.outdir, config)
+            for start in range(0, len(in_files), chunk_size)
+        )
+
+
+if __name__ == '__main__':
+    import argparse
+
+    aparser = argparse.ArgumentParser()
+    aparser.add_argument('config', type=str, help='Path to config for dataset generation')
+    aparser.add_argument('indir', type=str, help='Path to folder with images')
+    aparser.add_argument('outdir', type=str, help='Path to folder to store aligned images and masks to')
+    aparser.add_argument('--n-jobs', type=int, default=0, help='How many processes to use')
+    aparser.add_argument('--ext', type=str, default='jpg', help='Input image extension')
+
+    main(aparser.parse_args())

bin/gen_mask_dataset_hydra.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+
+import glob
+import os
+import shutil
+import traceback
+import hydra
+from omegaconf import OmegaConf
+
+import PIL.Image as Image
+import numpy as np
+from joblib import Parallel, delayed
+
+from saicinpainting.evaluation.masks.mask import SegmentationMask, propose_random_square_crop
+from saicinpainting.evaluation.utils import load_yaml, SmallMode
+from saicinpainting.training.data.masks import MixedMaskGenerator
+
+
+class MakeManyMasksWrapper:
+    def __init__(self, impl, variants_n=2):
+        self.impl = impl
+        self.variants_n = variants_n
+
+    def get_masks(self, img):
+        img = np.transpose(np.array(img), (2, 0, 1))
+        return [self.impl(img)[0] for _ in range(self.variants_n)]
+
+
+def process_images(src_images, indir, outdir, config):
+    if config.generator_kind == 'segmentation':
+        mask_generator = SegmentationMask(**config.mask_generator_kwargs)
+    elif config.generator_kind == 'random':
+        mask_generator_kwargs = OmegaConf.to_container(config.mask_generator_kwargs, resolve=True)
+        variants_n = mask_generator_kwargs.pop('variants_n', 2)
+        mask_generator = MakeManyMasksWrapper(MixedMaskGenerator(**mask_generator_kwargs),
+                                              variants_n=variants_n)
+    else:
+        raise ValueError(f'Unexpected generator kind: {config.generator_kind}')
+
+    max_tamper_area = config.get('max_tamper_area', 1)
+
+    for infile in src_images:
+        try:
+            file_relpath = infile[len(indir):]
+            img_outpath = os.path.join(outdir, file_relpath)
+            os.makedirs(os.path.dirname(img_outpath), exist_ok=True)
+
+            image = Image.open(infile).convert('RGB')
+
+            # scale input image to output resolution and filter smaller images
+            if min(image.size) < config.cropping.out_min_size:
+                handle_small_mode = SmallMode(config.cropping.handle_small_mode)
+                if handle_small_mode == SmallMode.DROP:
+                    continue
+                elif handle_small_mode == SmallMode.UPSCALE:
+                    factor = config.cropping.out_min_size / min(image.size)
+                    out_size = (np.array(image.size) * factor).round().astype('uint32')
+                    image = image.resize(out_size, resample=Image.BICUBIC)
+            else:
+                factor = config.cropping.out_min_size / min(image.size)
+                out_size = (np.array(image.size) * factor).round().astype('uint32')
+                image = image.resize(out_size, resample=Image.BICUBIC)
+
+            # generate and select masks
+            src_masks = mask_generator.get_masks(image)
+
+            filtered_image_mask_pairs = []
+            for cur_mask in src_masks:
+                if config.cropping.out_square_crop:
+                    (crop_left,
+                     crop_top,
+                     crop_right,
+                     crop_bottom) = propose_random_square_crop(cur_mask,
+                                                               min_overlap=config.cropping.crop_min_overlap)
+                    cur_mask = cur_mask[crop_top:crop_bottom, crop_left:crop_right]
+                    cur_image = image.copy().crop((crop_left, crop_top, crop_right, crop_bottom))
+                else:
+                    cur_image = image
+
+                if len(np.unique(cur_mask)) == 0 or cur_mask.mean() > max_tamper_area:
+                    continue
+
+                filtered_image_mask_pairs.append((cur_image, cur_mask))
+
+            mask_indices = np.random.choice(len(filtered_image_mask_pairs),
+                                            size=min(len(filtered_image_mask_pairs), config.max_masks_per_image),
+                                            replace=False)
+
+            # crop masks; save masks together with input image
+            mask_basename = os.path.join(outdir, os.path.splitext(file_relpath)[0])
+            for i, idx in enumerate(mask_indices):
+                cur_image, cur_mask = filtered_image_mask_pairs[idx]
+                cur_basename = mask_basename + f'_crop{i:03d}'
+                Image.fromarray(np.clip(cur_mask * 255, 0, 255).astype('uint8'),
+                                mode='L').save(cur_basename + f'_mask{i:03d}.png')
+                cur_image.save(cur_basename + '.png')
+        except KeyboardInterrupt:
+            return
+        except Exception as ex:
+            print(f'Could not make masks for {infile} due to {ex}:\n{traceback.format_exc()}')
+
+
+@hydra.main(config_path='../configs/data_gen/whydra', config_name='random_medium_256.yaml')
+def main(config: OmegaConf):
+    if not config.indir.endswith('/'):
+        config.indir += '/'
+
+    os.makedirs(config.outdir, exist_ok=True)
+
+    in_files = list(glob.glob(os.path.join(config.indir, '**', f'*.{config.location.extension}'),
+                              recursive=True))
+    if config.n_jobs == 0:
+        process_images(in_files, config.indir, config.outdir, config)
+    else:
+        in_files_n = len(in_files)
+        chunk_size = in_files_n // config.n_jobs + (1 if in_files_n % config.n_jobs > 0 else 0)
+        Parallel(n_jobs=config.n_jobs)(
+            delayed(process_images)(in_files[start:start+chunk_size], config.indir, config.outdir, config)
+            for start in range(0, len(in_files), chunk_size)
+        )
+
+
+if __name__ == '__main__':
+    main()

bin/gen_outpainting_dataset.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+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.pad(img, ((0, 0), (pad_h * 2, pad_h * 2), (pad_w * 2, pad_w * 2)), mode='symmetric')
+            # mask = np.pad(mask, ((pad_h * 2, pad_h * 2), (pad_w * 2, pad_w * 2)), mode = 'symmetric')
+
+            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())

bin/make_checkpoint.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+
+import os
+import shutil
+
+import torch
+
+
+def get_checkpoint_files(s):
+    s = s.strip()
+    if ',' in s:
+        return [get_checkpoint_files(chunk) for chunk in s.split(',')]
+    return 'last.ckpt' if s == 'last' else f'{s}.ckpt'
+
+
+def main(args):
+    checkpoint_fnames = get_checkpoint_files(args.epochs)
+    if isinstance(checkpoint_fnames, str):
+        checkpoint_fnames = [checkpoint_fnames]
+    assert len(checkpoint_fnames) >= 1
+
+    checkpoint_path = os.path.join(args.indir, 'models', checkpoint_fnames[0])
+    checkpoint = torch.load(checkpoint_path, map_location='cpu')
+    del checkpoint['optimizer_states']
+
+    if len(checkpoint_fnames) > 1:
+        for fname in checkpoint_fnames[1:]:
+            print('sum', fname)
+            sum_tensors_cnt = 0
+            other_cp = torch.load(os.path.join(args.indir, 'models', fname), map_location='cpu')
+            for k in checkpoint['state_dict'].keys():
+                if checkpoint['state_dict'][k].dtype is torch.float:
+                    checkpoint['state_dict'][k].data.add_(other_cp['state_dict'][k].data)
+                    sum_tensors_cnt += 1
+            print('summed', sum_tensors_cnt, 'tensors')
+
+        for k in checkpoint['state_dict'].keys():
+            if checkpoint['state_dict'][k].dtype is torch.float:
+                checkpoint['state_dict'][k].data.mul_(1 / float(len(checkpoint_fnames)))
+
+    state_dict = checkpoint['state_dict']
+
+    if not args.leave_discriminators:
+        for k in list(state_dict.keys()):
+            if k.startswith('discriminator.'):
+                del state_dict[k]
+
+    if not args.leave_losses:
+        for k in list(state_dict.keys()):
+            if k.startswith('loss_'):
+                del state_dict[k]
+
+    out_checkpoint_path = os.path.join(args.outdir, 'models', 'best.ckpt')
+    os.makedirs(os.path.dirname(out_checkpoint_path), exist_ok=True)
+
+    torch.save(checkpoint, out_checkpoint_path)
+
+    shutil.copy2(os.path.join(args.indir, 'config.yaml'),
+                 os.path.join(args.outdir, 'config.yaml'))
+
+
+if __name__ == '__main__':
+    import argparse
+
+    aparser = argparse.ArgumentParser()
+    aparser.add_argument('indir',
+                         help='Path to directory with output of training '
+                              '(i.e. directory, which has samples, modules, config.yaml and train.log')
+    aparser.add_argument('outdir',
+                         help='Where to put minimal checkpoint, which can be consumed by "bin/predict.py"')
+    aparser.add_argument('--epochs', type=str, default='last',
+                         help='Which checkpoint to take. '
+                              'Can be "last" or integer - number of epoch')
+    aparser.add_argument('--leave-discriminators', action='store_true',
+                         help='If enabled, the state of discriminators will not be removed from the checkpoint')
+    aparser.add_argument('--leave-losses', action='store_true',
+                         help='If enabled, weights of nn-based losses (e.g. perceptual) will not be removed')
+
+    main(aparser.parse_args())

bin/mask_example.py

@@ -0,0 +1,14 @@
+import matplotlib.pyplot as plt
+from skimage import io
+from skimage.transform import resize
+
+from saicinpainting.evaluation.masks.mask import SegmentationMask
+
+im = io.imread('imgs/ex4.jpg')
+im = resize(im, (512, 1024), anti_aliasing=True)
+mask_seg = SegmentationMask(num_variants_per_mask=10)
+mask_examples = mask_seg.get_masks(im)
+for i, example in enumerate(mask_examples):
+    plt.imshow(example)
+    plt.show()
+    plt.imsave(f'tmp/img_masks/{i}.png', example)

bin/paper_runfiles/blur_tests.sh

@@ -0,0 +1,37 @@
+##!/usr/bin/env bash
+#
+## !!! file set to make test_large_30k from the vanilla test_large: configs/test_large_30k.lst
+#
+## paths to data are valid for mml7
+#PLACES_ROOT="/data/inpainting/Places365"
+#OUT_DIR="/data/inpainting/paper_data/Places365_val_test"
+#
+#source "$(dirname $0)/env.sh"
+#
+#for datadir in test_large_30k  # val_large
+#do
+#    for conf in random_thin_256 random_medium_256 random_thick_256 random_thin_512 random_medium_512 random_thick_512
+#    do
+#        "$BINDIR/gen_mask_dataset.py" "$CONFIGDIR/data_gen/${conf}.yaml" \
+#            "$PLACES_ROOT/$datadir" "$OUT_DIR/$datadir/$conf" --n-jobs 8
+#
+#        "$BINDIR/calc_dataset_stats.py" --samples-n 20 "$OUT_DIR/$datadir/$conf" "$OUT_DIR/$datadir/${conf}_stats"
+#    done
+#
+#    for conf in segm_256 segm_512
+#    do
+#        "$BINDIR/gen_mask_dataset.py" "$CONFIGDIR/data_gen/${conf}.yaml" \
+#            "$PLACES_ROOT/$datadir" "$OUT_DIR/$datadir/$conf" --n-jobs 2
+#
+#        "$BINDIR/calc_dataset_stats.py" --samples-n 20 "$OUT_DIR/$datadir/$conf" "$OUT_DIR/$datadir/${conf}_stats"
+#    done
+#done
+#
+#IN_DIR="/data/inpainting/paper_data/Places365_val_test/test_large_30k/random_medium_512"
+#PRED_DIR="/data/inpainting/predictions/final/images/r.suvorov_2021-03-05_17-08-35_train_ablv2_work_resume_epoch37/random_medium_512"
+#BLUR_OUT_DIR="/data/inpainting/predictions/final/blur/images"
+#
+#for b in 0.1
+#
+#"$BINDIR/blur_predicts.py" "$BASEDIR/../../configs/eval2.yaml" "$CUR_IN_DIR" "$CUR_OUT_DIR" "$CUR_EVAL_DIR"
+#

bin/paper_runfiles/env.sh

@@ -0,0 +1,8 @@
+DIRNAME="$(dirname $0)"
+DIRNAME="$(realpath ""$DIRNAME"")"
+
+BINDIR="$DIRNAME/.."
+SRCDIR="$BINDIR/.."
+CONFIGDIR="$SRCDIR/configs"
+
+export PYTHONPATH="$SRCDIR:$PYTHONPATH"

bin/paper_runfiles/find_best_checkpoint.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+
+
+import os
+from argparse import ArgumentParser
+
+
+def ssim_fid100_f1(metrics, fid_scale=100):
+    ssim = metrics.loc['total', 'ssim']['mean']
+    fid = metrics.loc['total', 'fid']['mean']
+    fid_rel = max(0, fid_scale - fid) / fid_scale
+    f1 = 2 * ssim * fid_rel / (ssim + fid_rel + 1e-3)
+    return f1
+
+
+def find_best_checkpoint(model_list, models_dir):
+    with open(model_list) as f:
+        models = [m.strip() for m in f.readlines()]
+    with open(f'{model_list}_best', 'w') as f:
+        for model in models:
+            print(model)
+            best_f1 = 0
+            best_epoch = 0
+            best_step = 0
+            with open(os.path.join(models_dir, model, 'train.log')) as fm:
+                lines = fm.readlines()
+                for line_index in range(len(lines)):
+                    line = lines[line_index]
+                    if 'Validation metrics after epoch' in line:
+                        sharp_index = line.index('#')
+                        cur_ep = line[sharp_index + 1:]
+                        comma_index = cur_ep.index(',')
+                        cur_ep = int(cur_ep[:comma_index])
+                        total_index = line.index('total ')
+                        step = int(line[total_index:].split()[1].strip())
+                        total_line = lines[line_index + 5]
+                        if not total_line.startswith('total'):
+                            continue
+                        words = total_line.strip().split()
+                        f1 = float(words[-1])
+                        print(f'\tEpoch: {cur_ep}, f1={f1}')
+                        if f1 > best_f1:
+                            best_f1 = f1
+                            best_epoch = cur_ep
+                            best_step = step
+            f.write(f'{model}\t{best_epoch}\t{best_step}\t{best_f1}\n')
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument('model_list')
+    parser.add_argument('models_dir')
+    args = parser.parse_args()
+    find_best_checkpoint(args.model_list, args.models_dir)

bin/paper_runfiles/generate_test_celeba-hq.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+
+# paths to data are valid for mml-ws01
+OUT_DIR="/media/inpainting/paper_data/CelebA-HQ_val_test"
+
+source "$(dirname $0)/env.sh"
+
+for datadir in "val" "test"
+do
+    for conf in random_thin_256 random_medium_256 random_thick_256 random_thin_512 random_medium_512 random_thick_512
+    do
+        "$BINDIR/gen_mask_dataset_hydra.py" -cn $conf datadir=$datadir location=mml-ws01-celeba-hq \
+         location.out_dir=$OUT_DIR cropping.out_square_crop=False
+
+        "$BINDIR/calc_dataset_stats.py" --samples-n 20 "$OUT_DIR/$datadir/$conf" "$OUT_DIR/$datadir/${conf}_stats"
+    done
+done

bin/paper_runfiles/generate_test_ffhq.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+
+# paths to data are valid for mml-ws01
+OUT_DIR="/media/inpainting/paper_data/FFHQ_val"
+
+source "$(dirname $0)/env.sh"
+
+for datadir in test
+do
+    for conf in random_thin_256 random_medium_256 random_thick_256 random_thin_512 random_medium_512 random_thick_512
+    do
+        "$BINDIR/gen_mask_dataset_hydra.py" -cn $conf datadir=$datadir location=mml-ws01-ffhq \
+         location.out_dir=$OUT_DIR cropping.out_square_crop=False
+
+        "$BINDIR/calc_dataset_stats.py" --samples-n 20 "$OUT_DIR/$datadir/$conf" "$OUT_DIR/$datadir/${conf}_stats"
+    done
+done

bin/paper_runfiles/generate_test_paris.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+
+# paths to data are valid for mml-ws01
+OUT_DIR="/media/inpainting/paper_data/Paris_StreetView_Dataset_val"
+
+source "$(dirname $0)/env.sh"
+
+for datadir in paris_eval_gt
+do
+    for conf in random_thin_256 random_medium_256 random_thick_256 segm_256
+    do
+        "$BINDIR/gen_mask_dataset_hydra.py" -cn $conf datadir=$datadir location=mml-ws01-paris \
+         location.out_dir=OUT_DIR cropping.out_square_crop=False cropping.out_min_size=227
+
+        "$BINDIR/calc_dataset_stats.py" --samples-n 20 "$OUT_DIR/$datadir/$conf" "$OUT_DIR/$datadir/${conf}_stats"
+    done
+done

bin/paper_runfiles/generate_test_paris_256.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+
+# paths to data are valid for mml-ws01
+OUT_DIR="/media/inpainting/paper_data/Paris_StreetView_Dataset_val_256"
+
+source "$(dirname $0)/env.sh"
+
+for datadir in paris_eval_gt
+do
+    for conf in random_thin_256 random_medium_256 random_thick_256 segm_256
+    do
+        "$BINDIR/gen_mask_dataset_hydra.py" -cn $conf datadir=$datadir location=mml-ws01-paris \
+         location.out_dir=$OUT_DIR cropping.out_square_crop=False cropping.out_min_size=256
+
+        "$BINDIR/calc_dataset_stats.py" --samples-n 20 "$OUT_DIR/$datadir/$conf" "$OUT_DIR/$datadir/${conf}_stats"
+    done
+done

bin/paper_runfiles/generate_val_test.sh

@@ -0,0 +1,28 @@
+#!/usr/bin/env bash
+
+# !!! file set to make test_large_30k from the vanilla test_large: configs/test_large_30k.lst
+
+# paths to data are valid for mml7
+PLACES_ROOT="/data/inpainting/Places365"
+OUT_DIR="/data/inpainting/paper_data/Places365_val_test"
+
+source "$(dirname $0)/env.sh"
+
+for datadir in test_large_30k  # val_large
+do
+    for conf in random_thin_256 random_medium_256 random_thick_256 random_thin_512 random_medium_512 random_thick_512
+    do
+        "$BINDIR/gen_mask_dataset.py" "$CONFIGDIR/data_gen/${conf}.yaml" \
+            "$PLACES_ROOT/$datadir" "$OUT_DIR/$datadir/$conf" --n-jobs 8
+
+        "$BINDIR/calc_dataset_stats.py" --samples-n 20 "$OUT_DIR/$datadir/$conf" "$OUT_DIR/$datadir/${conf}_stats"
+    done
+
+    for conf in segm_256 segm_512
+    do
+        "$BINDIR/gen_mask_dataset.py" "$CONFIGDIR/data_gen/${conf}.yaml" \
+            "$PLACES_ROOT/$datadir" "$OUT_DIR/$datadir/$conf" --n-jobs 2
+
+        "$BINDIR/calc_dataset_stats.py" --samples-n 20 "$OUT_DIR/$datadir/$conf" "$OUT_DIR/$datadir/${conf}_stats"
+    done
+done

bin/paper_runfiles/predict_inner_features.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+
+# paths to data are valid for mml7
+
+source "$(dirname $0)/env.sh"
+
+"$BINDIR/predict_inner_features.py" \
+    -cn default_inner_features_ffc \
+    model.path="/data/inpainting/paper_data/final_models/ours/r.suvorov_2021-03-05_17-34-05_train_ablv2_work_ffc075_resume_epoch39" \
+    indir="/data/inpainting/paper_data/inner_features_vis/input/" \
+    outdir="/data/inpainting/paper_data/inner_features_vis/output/ffc" \
+    dataset.img_suffix=.png
+
+
+"$BINDIR/predict_inner_features.py" \
+    -cn default_inner_features_work \
+    model.path="/data/inpainting/paper_data/final_models/ours/r.suvorov_2021-03-05_17-08-35_train_ablv2_work_resume_epoch37" \
+    indir="/data/inpainting/paper_data/inner_features_vis/input/" \
+    outdir="/data/inpainting/paper_data/inner_features_vis/output/work" \
+    dataset.img_suffix=.png

bin/paper_runfiles/update_test_data_stats.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+
+# paths to data are valid for mml7
+
+source "$(dirname $0)/env.sh"
+
+#INDIR="/data/inpainting/paper_data/Places365_val_test/test_large_30k"
+#
+#for dataset in random_medium_256 random_medium_512 random_thick_256 random_thick_512 random_thin_256 random_thin_512
+#do
+#    "$BINDIR/calc_dataset_stats.py" "$INDIR/$dataset" "$INDIR/${dataset}_stats2"
+#done
+#
+#"$BINDIR/calc_dataset_stats.py" "/data/inpainting/evalset2" "/data/inpainting/evalset2_stats2"
+
+
+INDIR="/data/inpainting/paper_data/CelebA-HQ_val_test/test"
+
+for dataset in random_medium_256 random_thick_256 random_thin_256
+do
+    "$BINDIR/calc_dataset_stats.py" "$INDIR/$dataset" "$INDIR/${dataset}_stats2"
+done
+
+
+INDIR="/data/inpainting/paper_data/Paris_StreetView_Dataset_val_256/paris_eval_gt"
+
+for dataset in random_medium_256 random_thick_256 random_thin_256
+do
+    "$BINDIR/calc_dataset_stats.py" "$INDIR/$dataset" "$INDIR/${dataset}_stats2"
+done

bin/predict.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+
+# Example command:
+# ./bin/predict.py \
+#       model.path=<path to checkpoint, prepared by make_checkpoint.py> \
+#       indir=<path to input data> \
+#       outdir=<where to store predicts>
+
+import logging
+import os
+import sys
+import traceback
+
+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__)
+
+
+@hydra.main(config_path='../configs/prediction', config_name='default.yaml')
+def main(predict_config: OmegaConf):
+    try:
+        register_debug_signal_handlers()  # kill -10 <pid> will result in traceback dumped into log
+
+        device = torch.device(predict_config.device)
+
+        train_config_path = os.path.join(predict_config.model.path, 'config.yaml')
+        with open(train_config_path, 'r') as f:
+            train_config = OmegaConf.create(yaml.safe_load(f))
+        
+        train_config.training_model.predict_only = True
+
+        out_ext = predict_config.get('out_ext', '.png')
+
+        checkpoint_path = os.path.join(predict_config.model.path, 
+                                       'models', 
+                                       predict_config.model.checkpoint)
+        model = load_checkpoint(train_config, checkpoint_path, strict=False, map_location='cpu')
+        model.freeze()
+        model.to(device)
+
+        if not predict_config.indir.endswith('/'):
+            predict_config.indir += '/'
+
+        dataset = make_default_val_dataset(predict_config.indir, **predict_config.dataset)
+        with torch.no_grad():
+            for img_i in tqdm.trange(len(dataset)):
+                mask_fname = dataset.mask_filenames[img_i]
+                cur_out_fname = os.path.join(
+                    predict_config.outdir, 
+                    os.path.splitext(mask_fname[len(predict_config.indir):])[0] + out_ext
+                )
+                os.makedirs(os.path.dirname(cur_out_fname), exist_ok=True)
+
+                batch = move_to_device(default_collate([dataset[img_i]]), device)
+                batch['mask'] = (batch['mask'] > 0) * 1
+                batch = model(batch)
+                cur_res = batch[predict_config.out_key][0].permute(1, 2, 0).detach().cpu().numpy()
+
+                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)
+    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__':
+    main()

bin/predict_inner_features.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+
+# Example command:
+# ./bin/predict.py \
+#       model.path=<path to checkpoint, prepared by make_checkpoint.py> \
+#       indir=<path to input data> \
+#       outdir=<where to store predicts>
+
+import logging
+import os
+import sys
+import traceback
+
+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, DefaultInpaintingTrainingModule
+from saicinpainting.utils import register_debug_signal_handlers, get_shape
+
+LOGGER = logging.getLogger(__name__)
+
+
+@hydra.main(config_path='../configs/prediction', config_name='default_inner_features.yaml')
+def main(predict_config: OmegaConf):
+    try:
+        register_debug_signal_handlers()  # kill -10 <pid> will result in traceback dumped into log
+
+        device = torch.device(predict_config.device)
+
+        train_config_path = os.path.join(predict_config.model.path, 'config.yaml')
+        with open(train_config_path, 'r') as f:
+            train_config = OmegaConf.create(yaml.safe_load(f))
+
+        checkpoint_path = os.path.join(predict_config.model.path, 'models', predict_config.model.checkpoint)
+        model = load_checkpoint(train_config, checkpoint_path, strict=False)
+        model.freeze()
+        model.to(device)
+
+        assert isinstance(model, DefaultInpaintingTrainingModule), 'Only DefaultInpaintingTrainingModule is supported'
+        assert isinstance(getattr(model.generator, 'model', None), torch.nn.Sequential)
+
+        if not predict_config.indir.endswith('/'):
+            predict_config.indir += '/'
+
+        dataset = make_default_val_dataset(predict_config.indir, **predict_config.dataset)
+
+        max_level = max(predict_config.levels)
+
+        with torch.no_grad():
+            for img_i in tqdm.trange(len(dataset)):
+                mask_fname = dataset.mask_filenames[img_i]
+                cur_out_fname = os.path.join(predict_config.outdir, os.path.splitext(mask_fname[len(predict_config.indir):])[0])
+                os.makedirs(os.path.dirname(cur_out_fname), exist_ok=True)
+
+                batch = move_to_device(default_collate([dataset[img_i]]), device)
+
+                img = batch['image']
+                mask = batch['mask']
+                mask[:] = 0
+                mask_h, mask_w = mask.shape[-2:]
+                mask[:, :,
+                    mask_h // 2 - predict_config.hole_radius : mask_h // 2 + predict_config.hole_radius,
+                    mask_w // 2 - predict_config.hole_radius : mask_w // 2 + predict_config.hole_radius] = 1
+
+                masked_img = torch.cat([img * (1 - mask), mask], dim=1)
+
+                feats = masked_img
+                for level_i, level in enumerate(model.generator.model):
+                    feats = level(feats)
+                    if level_i in predict_config.levels:
+                        cur_feats = torch.cat([f for f in feats if torch.is_tensor(f)], dim=1) \
+                            if isinstance(feats, tuple) else feats
+
+                        if predict_config.slice_channels:
+                            cur_feats = cur_feats[:, slice(*predict_config.slice_channels)]
+
+                        cur_feat = cur_feats.pow(2).mean(1).pow(0.5).clone()
+                        cur_feat -= cur_feat.min()
+                        cur_feat /= cur_feat.std()
+                        cur_feat = cur_feat.clamp(0, 1) / 1
+                        cur_feat = cur_feat.cpu().numpy()[0]
+                        cur_feat *= 255
+                        cur_feat = np.clip(cur_feat, 0, 255).astype('uint8')
+                        cv2.imwrite(cur_out_fname + f'_lev{level_i:02d}_norm.png', cur_feat)
+
+                        # for channel_i in predict_config.channels:
+                        #
+                        #     cur_feat = cur_feats[0, channel_i].clone().detach().cpu().numpy()
+                        #     cur_feat -= cur_feat.min()
+                        #     cur_feat /= cur_feat.max()
+                        #     cur_feat *= 255
+                        #     cur_feat = np.clip(cur_feat, 0, 255).astype('uint8')
+                        #     cv2.imwrite(cur_out_fname + f'_lev{level_i}_ch{channel_i}.png', cur_feat)
+                    elif level_i >= max_level:
+                        break
+    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__':
+    main()

bin/report_from_tb.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+
+import glob
+import os
+import re
+
+import tensorflow as tf
+from torch.utils.tensorboard import SummaryWriter
+
+
+GROUPING_RULES = [
+    re.compile(r'^(?P<group>train|test|val|extra_val_.*?(256|512))_(?P<title>.*)', re.I)
+]
+
+
+DROP_RULES = [
+    re.compile(r'_std$', re.I)
+]
+
+
+def need_drop(tag):
+    for rule in DROP_RULES:
+        if rule.search(tag):
+            return True
+    return False
+
+
+def get_group_and_title(tag):
+    for rule in GROUPING_RULES:
+        match = rule.search(tag)
+        if match is None:
+            continue
+        return match.group('group'), match.group('title')
+    return None, None
+
+
+def main(args):
+    os.makedirs(args.outdir, exist_ok=True)
+
+    ignored_events = set()
+
+    for orig_fname in glob.glob(args.inglob):
+        cur_dirpath = os.path.dirname(orig_fname)  # remove filename, this should point to "version_0" directory
+        subdirname = os.path.basename(cur_dirpath)  # == "version_0" most of time
+        exp_root_path = os.path.dirname(cur_dirpath)  # remove "version_0"
+        exp_name = os.path.basename(exp_root_path)
+
+        writers_by_group = {}
+
+        for e in tf.compat.v1.train.summary_iterator(orig_fname):
+            for v in e.summary.value:
+                if need_drop(v.tag):
+                    continue
+
+                cur_group, cur_title = get_group_and_title(v.tag)
+                if cur_group is None:
+                    if v.tag not in ignored_events:
+                        print(f'WARNING: Could not detect group for {v.tag}, ignoring it')
+                        ignored_events.add(v.tag)
+                    continue
+
+                cur_writer = writers_by_group.get(cur_group, None)
+                if cur_writer is None:
+                    if args.include_version:
+                        cur_outdir = os.path.join(args.outdir, exp_name, f'{subdirname}_{cur_group}')
+                    else:
+                        cur_outdir = os.path.join(args.outdir, exp_name, cur_group)
+                    cur_writer = SummaryWriter(cur_outdir)
+                    writers_by_group[cur_group] = cur_writer
+
+                cur_writer.add_scalar(cur_title, v.simple_value, global_step=e.step, walltime=e.wall_time)
+
+
+if __name__ == '__main__':
+    import argparse
+
+    aparser = argparse.ArgumentParser()
+    aparser.add_argument('inglob', type=str)
+    aparser.add_argument('outdir', type=str)
+    aparser.add_argument('--include-version', action='store_true',
+                         help='Include subdirectory name e.g. "version_0" into output path')
+
+    main(aparser.parse_args())

bin/sample_from_dataset.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+
+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']# > 0.5
+    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())

bin/side_by_side.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+import os
+import random
+
+import cv2
+import numpy as np
+
+from saicinpainting.evaluation.data import PrecomputedInpaintingResultsDataset
+from saicinpainting.evaluation.utils import load_yaml
+from saicinpainting.training.visualizers.base import visualize_mask_and_images
+
+
+def main(args):
+    config = load_yaml(args.config)
+
+    datasets = [PrecomputedInpaintingResultsDataset(args.datadir, cur_predictdir, **config.dataset_kwargs)
+                for cur_predictdir in args.predictdirs]
+    assert len({len(ds) for ds in datasets}) == 1
+    len_first = len(datasets[0])
+
+    indices = list(range(len_first))
+    if len_first > args.max_n:
+        indices = sorted(random.sample(indices, args.max_n))
+
+    os.makedirs(args.outpath, exist_ok=True)
+
+    filename2i = {}
+
+    keys = ['image'] + [i for i in range(len(datasets))]
+    for img_i in indices:
+        try:
+            mask_fname = os.path.basename(datasets[0].mask_filenames[img_i])
+            if mask_fname in filename2i:
+                filename2i[mask_fname] += 1
+                idx = filename2i[mask_fname]
+                mask_fname_only, ext = os.path.split(mask_fname)
+                mask_fname = f'{mask_fname_only}_{idx}{ext}'
+            else:
+                filename2i[mask_fname] = 1
+
+            cur_vis_dict = datasets[0][img_i]
+            for ds_i, ds in enumerate(datasets):
+                cur_vis_dict[ds_i] = ds[img_i]['inpainted']
+
+            vis_img = visualize_mask_and_images(cur_vis_dict, keys,
+                                                last_without_mask=False,
+                                                mask_only_first=True,
+                                                black_mask=args.black)
+            vis_img = np.clip(vis_img * 255, 0, 255).astype('uint8')
+
+            out_fname = os.path.join(args.outpath, mask_fname)
+
+
+
+            vis_img = cv2.cvtColor(vis_img, cv2.COLOR_RGB2BGR)
+            cv2.imwrite(out_fname, vis_img)
+        except Exception as ex:
+            print(f'Could not process {img_i} due to {ex}')
+
+
+if __name__ == '__main__':
+    import argparse
+
+    aparser = argparse.ArgumentParser()
+    aparser.add_argument('--max-n', type=int, default=100, help='Maximum number of images to print')
+    aparser.add_argument('--black', action='store_true', help='Whether to fill mask on GT with black')
+    aparser.add_argument('config', type=str, help='Path to evaluation config (e.g. configs/eval1.yaml)')
+    aparser.add_argument('outpath', type=str, help='Where to put results')
+    aparser.add_argument('datadir', type=str,
+                         help='Path to folder with images and masks')
+    aparser.add_argument('predictdirs', type=str,
+                         nargs='+',
+                         help='Path to folders with predicts')
+
+
+    main(aparser.parse_args())

bin/split_tar.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+
+
+import tqdm
+import webdataset as wds
+
+
+def main(args):
+    input_dataset = wds.Dataset(args.infile)
+    output_dataset = wds.ShardWriter(args.outpattern)
+    for rec in tqdm.tqdm(input_dataset):
+        output_dataset.write(rec)
+
+
+if __name__ == '__main__':
+    import argparse
+
+    aparser = argparse.ArgumentParser()
+    aparser.add_argument('infile', type=str)
+    aparser.add_argument('outpattern', type=str)
+
+    main(aparser.parse_args())

bin/train.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+
+import logging
+import os
+import sys
+import traceback
+
+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 hydra
+from omegaconf import OmegaConf
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.loggers import TensorBoardLogger
+from pytorch_lightning.plugins import DDPPlugin
+
+from saicinpainting.training.trainers import make_training_model
+from saicinpainting.utils import register_debug_signal_handlers, handle_ddp_subprocess, handle_ddp_parent_process, \
+    handle_deterministic_config
+
+LOGGER = logging.getLogger(__name__)
+
+
+@handle_ddp_subprocess()
+@hydra.main(config_path='../configs/training', config_name='tiny_test.yaml')
+def main(config: OmegaConf):
+    try:
+        need_set_deterministic = handle_deterministic_config(config)
+
+        register_debug_signal_handlers()  # kill -10 <pid> will result in traceback dumped into log
+
+        is_in_ddp_subprocess = handle_ddp_parent_process()
+
+        config.visualizer.outdir = os.path.join(os.getcwd(), config.visualizer.outdir)
+        if not is_in_ddp_subprocess:
+            LOGGER.info(OmegaConf.to_yaml(config))
+            OmegaConf.save(config, os.path.join(os.getcwd(), 'config.yaml'))
+
+        checkpoints_dir = os.path.join(os.getcwd(), 'models')
+        os.makedirs(checkpoints_dir, exist_ok=True)
+
+        # there is no need to suppress this logger in ddp, because it handles rank on its own
+        metrics_logger = TensorBoardLogger(config.location.tb_dir, name=os.path.basename(os.getcwd()))
+        metrics_logger.log_hyperparams(config)
+
+        training_model = make_training_model(config)
+
+        trainer_kwargs = OmegaConf.to_container(config.trainer.kwargs, resolve=True)
+        if need_set_deterministic:
+            trainer_kwargs['deterministic'] = True
+
+        trainer = Trainer(
+            # there is no need to suppress checkpointing in ddp, because it handles rank on its own
+            callbacks=ModelCheckpoint(dirpath=checkpoints_dir, **config.trainer.checkpoint_kwargs),
+            logger=metrics_logger,
+            default_root_dir=os.getcwd(),
+            **trainer_kwargs
+        )
+        trainer.fit(training_model)
+    except KeyboardInterrupt:
+        LOGGER.warning('Interrupted by user')
+    except Exception as ex:
+        LOGGER.critical(f'Training failed due to {ex}:\n{traceback.format_exc()}')
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+    main()

conda_env.yml

@@ -0,0 +1,165 @@
+name: lama
+channels:
+  - defaults
+  - conda-forge
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=4.5=1_gnu
+  - absl-py=0.13.0=py36h06a4308_0
+  - aiohttp=3.7.4.post0=py36h7f8727e_2
+  - antlr-python-runtime=4.8=py36h9f0ad1d_2
+  - async-timeout=3.0.1=py36h06a4308_0
+  - attrs=21.2.0=pyhd3eb1b0_0
+  - blas=1.0=mkl
+  - blinker=1.4=py36h06a4308_0
+  - brotlipy=0.7.0=py36h27cfd23_1003
+  - bzip2=1.0.8=h7b6447c_0
+  - c-ares=1.17.1=h27cfd23_0
+  - ca-certificates=2021.7.5=h06a4308_1
+  - cachetools=4.2.2=pyhd3eb1b0_0
+  - certifi=2021.5.30=py36h06a4308_0
+  - cffi=1.14.6=py36h400218f_0
+  - chardet=4.0.0=py36h06a4308_1003
+  - charset-normalizer=2.0.4=pyhd3eb1b0_0
+  - click=8.0.1=pyhd3eb1b0_0
+  - cloudpickle=2.0.0=pyhd3eb1b0_0
+  - coverage=5.5=py36h27cfd23_2
+  - cryptography=3.4.7=py36hd23ed53_0
+  - cudatoolkit=10.2.89=hfd86e86_1
+  - cycler=0.10.0=py36_0
+  - cython=0.29.24=py36h295c915_0
+  - cytoolz=0.11.0=py36h7b6447c_0
+  - dask-core=1.1.4=py36_1
+  - dataclasses=0.8=pyh4f3eec9_6
+  - dbus=1.13.18=hb2f20db_0
+  - decorator=5.0.9=pyhd3eb1b0_0
+  - easydict=1.9=py_0
+  - expat=2.4.1=h2531618_2
+  - ffmpeg=4.2.2=h20bf706_0
+  - fontconfig=2.13.1=h6c09931_0
+  - freetype=2.10.4=h5ab3b9f_0
+  - fsspec=2021.8.1=pyhd3eb1b0_0
+  - future=0.18.2=py36_1
+  - glib=2.69.1=h5202010_0
+  - gmp=6.2.1=h2531618_2
+  - gnutls=3.6.15=he1e5248_0
+  - google-auth=1.33.0=pyhd3eb1b0_0
+  - google-auth-oauthlib=0.4.4=pyhd3eb1b0_0
+  - grpcio=1.36.1=py36h2157cd5_1
+  - gst-plugins-base=1.14.0=h8213a91_2
+  - gstreamer=1.14.0=h28cd5cc_2
+  - hydra-core=1.1.0=pyhd8ed1ab_0
+  - icu=58.2=he6710b0_3
+  - idna=3.2=pyhd3eb1b0_0
+  - idna_ssl=1.1.0=py36h06a4308_0
+  - imageio=2.9.0=pyhd3eb1b0_0
+  - importlib-metadata=4.8.1=py36h06a4308_0
+  - importlib_resources=5.2.0=pyhd3eb1b0_1
+  - intel-openmp=2021.3.0=h06a4308_3350
+  - joblib=1.0.1=pyhd3eb1b0_0
+  - jpeg=9b=h024ee3a_2
+  - kiwisolver=1.3.1=py36h2531618_0
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.35.1=h7274673_9
+  - libblas=3.9.0=11_linux64_mkl
+  - libcblas=3.9.0=11_linux64_mkl
+  - libffi=3.3=he6710b0_2
+  - libgcc-ng=9.3.0=h5101ec6_17
+  - libgfortran-ng=9.3.0=ha5ec8a7_17
+  - libgfortran5=9.3.0=ha5ec8a7_17
+  - libgomp=9.3.0=h5101ec6_17
+  - libidn2=2.3.2=h7f8727e_0
+  - liblapack=3.9.0=11_linux64_mkl
+  - libopus=1.3.1=h7b6447c_0
+  - libpng=1.6.37=hbc83047_0
+  - libprotobuf=3.17.2=h4ff587b_1
+  - libstdcxx-ng=9.3.0=hd4cf53a_17
+  - libtasn1=4.16.0=h27cfd23_0
+  - libtiff=4.2.0=h85742a9_0
+  - libunistring=0.9.10=h27cfd23_0
+  - libuuid=1.0.3=h1bed415_2
+  - libuv=1.40.0=h7b6447c_0
+  - libvpx=1.7.0=h439df22_0
+  - libwebp-base=1.2.0=h27cfd23_0
+  - libxcb=1.14=h7b6447c_0
+  - libxml2=2.9.12=h03d6c58_0
+  - lz4-c=1.9.3=h295c915_1
+  - markdown=3.3.4=py36h06a4308_0
+  - matplotlib=3.3.4=py36h06a4308_0
+  - matplotlib-base=3.3.4=py36h62a2d02_0
+  - mkl=2021.3.0=h06a4308_520
+  - multidict=5.1.0=py36h27cfd23_2
+  - ncurses=6.2=he6710b0_1
+  - nettle=3.7.3=hbbd107a_1
+  - networkx=2.2=py36_1
+  - ninja=1.10.2=hff7bd54_1
+  - numpy=1.19.5=py36hfc0c790_2
+  - oauthlib=3.1.1=pyhd3eb1b0_0
+  - olefile=0.46=py36_0
+  - omegaconf=2.1.1=py36h5fab9bb_0
+  - openh264=2.1.0=hd408876_0
+  - openjpeg=2.4.0=h3ad879b_0
+  - openssl=1.1.1l=h7f8727e_0
+  - packaging=21.0=pyhd3eb1b0_0
+  - pandas=1.1.5=py36h284efc9_0
+  - pcre=8.45=h295c915_0
+  - pillow=8.3.1=py36h2c7a002_0
+  - pip=21.0.1=py36h06a4308_0
+  - protobuf=3.17.2=py36h295c915_0
+  - pyasn1=0.4.8=pyhd3eb1b0_0
+  - pyasn1-modules=0.2.8=py_0
+  - pycparser=2.20=py_2
+  - pyjwt=2.1.0=py36h06a4308_0
+  - pyopenssl=20.0.1=pyhd3eb1b0_1
+  - pyparsing=2.4.7=pyhd3eb1b0_0
+  - pyqt=5.9.2=py36h05f1152_2
+  - pysocks=1.7.1=py36h06a4308_0
+  - python=3.6.13=h12debd9_1
+  - python-dateutil=2.8.2=pyhd3eb1b0_0
+  - python_abi=3.6=2_cp36m
+  - pytz=2021.1=pyhd3eb1b0_0
+  - pywavelets=1.1.1=py36h7b6447c_2
+  - pyyaml=5.4.1=py36h27cfd23_1
+  - qt=5.9.7=h5867ecd_1
+  - readline=8.1=h27cfd23_0
+  - requests=2.26.0=pyhd3eb1b0_0
+  - requests-oauthlib=1.3.0=py_0
+  - rsa=4.7.2=pyhd3eb1b0_1
+  - scikit-image=0.17.2=py36h284efc9_4
+  - scikit-learn=0.24.2=py36ha9443f7_0
+  - scipy=1.5.3=py36h9e8f40b_0
+  - setuptools=58.0.4=py36h06a4308_0
+  - sip=4.19.8=py36hf484d3e_0
+  - six=1.16.0=pyhd3eb1b0_0
+  - sqlite=3.36.0=hc218d9a_0
+  - tabulate=0.8.9=py36h06a4308_0
+  - tensorboard=2.4.0=pyhc547734_0
+  - tensorboard-plugin-wit=1.6.0=py_0
+  - threadpoolctl=2.2.0=pyh0d69192_0
+  - tifffile=2020.10.1=py36hdd07704_2
+  - tk=8.6.11=h1ccaba5_0
+  - toolz=0.11.1=pyhd3eb1b0_0
+  - tqdm=4.62.2=pyhd3eb1b0_1
+  - typing-extensions=3.10.0.2=hd3eb1b0_0
+  - typing_extensions=3.10.0.2=pyh06a4308_0
+  - urllib3=1.26.6=pyhd3eb1b0_1
+  - werkzeug=2.0.1=pyhd3eb1b0_0
+  - wheel=0.37.0=pyhd3eb1b0_1
+  - x264=1!157.20191217=h7b6447c_0
+  - xz=5.2.5=h7b6447c_0
+  - yaml=0.2.5=h7b6447c_0
+  - yarl=1.6.3=py36h27cfd23_0
+  - zipp=3.5.0=pyhd3eb1b0_0
+  - zlib=1.2.11=h7b6447c_3
+  - zstd=1.4.9=haebb681_0
+  - pip:
+    - albumentations==0.5.2
+    - braceexpand==0.1.7
+    - imgaug==0.4.0
+    - kornia==0.5.0
+    - opencv-python==4.5.3.56
+    - opencv-python-headless==4.5.3.56
+    - shapely==1.7.1
+    - webdataset==0.1.76
+    - wldhx-yadisk-direct==0.0.6

configs/analyze_mask_errors.yaml

@@ -0,0 +1,7 @@
+dataset_kwargs:
+  img_suffix: .jpg
+  inpainted_suffix: .jpg
+
+take_global_top: 30
+take_worst_best_top: 30
+take_overlapping_top: 30

configs/data_gen/gen_segm_dataset1.yaml

@@ -0,0 +1,25 @@
+generator_kind: segmentation
+
+mask_generator_kwargs:
+  confidence_threshold: 0.5
+  max_object_area: 0.5
+  min_mask_area: 0.02
+  downsample_levels: 6
+  num_variants_per_mask: 5
+  rigidness_mode: 1
+  max_foreground_coverage: 0.3
+  max_foreground_intersection: 0.7
+  max_mask_intersection: 0.1
+  max_hidden_area: 0.1
+  max_scale_change: 0.25
+  horizontal_flip: True
+  max_vertical_shift: 0.2
+  position_shuffle: True
+
+max_masks_per_image: 5
+
+cropping:
+  out_min_size: 512
+  handle_small_mode: drop
+  out_square_crop: True
+  crop_min_overlap: 0.5

configs/data_gen/gen_segm_dataset3.yaml

@@ -0,0 +1,25 @@
+generator_kind: segmentation
+
+mask_generator_kwargs:
+  confidence_threshold: 0.5
+  max_object_area: 0.5
+  min_mask_area: 0.07
+  downsample_levels: 6
+  num_variants_per_mask: 3
+  rigidness_mode: 1
+  max_foreground_coverage: 0.4
+  max_foreground_intersection: 0.8
+  max_mask_intersection: 0.2
+  max_hidden_area: 0.1
+  max_scale_change: 0.25
+  horizontal_flip: True
+  max_vertical_shift: 0.3
+  position_shuffle: True
+
+max_masks_per_image: 3
+
+cropping:
+  out_min_size: 512
+  handle_small_mode: drop
+  out_square_crop: True
+  crop_min_overlap: 0.5

configs/data_gen/random_medium_256.yaml

@@ -0,0 +1,33 @@
+generator_kind: random
+
+mask_generator_kwargs:
+  irregular_proba: 1
+  irregular_kwargs:
+    min_times: 4
+    max_times: 5
+    max_width: 50
+    max_angle: 4
+    max_len: 100
+
+  box_proba: 0.3
+  box_kwargs:
+    margin: 0
+    bbox_min_size: 10
+    bbox_max_size: 50
+    max_times: 5
+    min_times: 1
+
+  segm_proba: 0
+  squares_proba: 0
+
+  variants_n: 5
+
+max_masks_per_image: 1
+
+cropping:
+  out_min_size: 256
+  handle_small_mode: upscale
+  out_square_crop: True
+  crop_min_overlap: 1
+
+max_tamper_area: 0.5

configs/data_gen/random_medium_512.yaml

@@ -0,0 +1,33 @@
+generator_kind: random
+
+mask_generator_kwargs:
+  irregular_proba: 1
+  irregular_kwargs:
+    min_times: 4
+    max_times: 10
+    max_width: 100
+    max_angle: 4
+    max_len: 200
+
+  box_proba: 0.3
+  box_kwargs:
+    margin: 0
+    bbox_min_size: 30
+    bbox_max_size: 150
+    max_times: 5
+    min_times: 1
+
+  segm_proba: 0
+  squares_proba: 0
+
+  variants_n: 5
+
+max_masks_per_image: 1
+
+cropping:
+  out_min_size: 512
+  handle_small_mode: upscale
+  out_square_crop: True
+  crop_min_overlap: 1
+
+max_tamper_area: 0.5

configs/data_gen/random_thick_256.yaml

@@ -0,0 +1,33 @@
+generator_kind: random
+
+mask_generator_kwargs:
+  irregular_proba: 1
+  irregular_kwargs:
+    min_times: 1
+    max_times: 5
+    max_width: 100
+    max_angle: 4
+    max_len: 200
+
+  box_proba: 0.3
+  box_kwargs:
+    margin: 10
+    bbox_min_size: 30
+    bbox_max_size: 150
+    max_times: 3
+    min_times: 1
+
+  segm_proba: 0
+  squares_proba: 0
+
+  variants_n: 5
+
+max_masks_per_image: 1
+
+cropping:
+  out_min_size: 256
+  handle_small_mode: upscale
+  out_square_crop: True
+  crop_min_overlap: 1
+
+max_tamper_area: 0.5

configs/data_gen/random_thick_512.yaml

@@ -0,0 +1,33 @@
+generator_kind: random
+
+mask_generator_kwargs:
+  irregular_proba: 1
+  irregular_kwargs:
+    min_times: 1
+    max_times: 5
+    max_width: 250
+    max_angle: 4
+    max_len: 450
+
+  box_proba: 0.3
+  box_kwargs:
+    margin: 10
+    bbox_min_size: 30
+    bbox_max_size: 300
+    max_times: 4
+    min_times: 1
+
+  segm_proba: 0
+  squares_proba: 0
+
+  variants_n: 5
+
+max_masks_per_image: 1
+
+cropping:
+  out_min_size: 512
+  handle_small_mode: upscale
+  out_square_crop: True
+  crop_min_overlap: 1
+
+max_tamper_area: 0.5

configs/data_gen/random_thin_256.yaml

@@ -0,0 +1,25 @@
+generator_kind: random
+
+mask_generator_kwargs:
+  irregular_proba: 1
+  irregular_kwargs:
+    min_times: 4
+    max_times: 50
+    max_width: 10
+    max_angle: 4
+    max_len: 40
+  box_proba: 0
+  segm_proba: 0
+  squares_proba: 0
+
+  variants_n: 5
+
+max_masks_per_image: 1
+
+cropping:
+  out_min_size: 256
+  handle_small_mode: upscale
+  out_square_crop: True
+  crop_min_overlap: 1
+
+max_tamper_area: 0.5

configs/data_gen/random_thin_512.yaml

@@ -0,0 +1,25 @@
+generator_kind: random
+
+mask_generator_kwargs:
+  irregular_proba: 1
+  irregular_kwargs:
+    min_times: 4
+    max_times: 70
+    max_width: 20
+    max_angle: 4
+    max_len: 100
+  box_proba: 0
+  segm_proba: 0
+  squares_proba: 0
+
+  variants_n: 5
+
+max_masks_per_image: 1
+
+cropping:
+  out_min_size: 512
+  handle_small_mode: upscale
+  out_square_crop: True
+  crop_min_overlap: 1
+
+max_tamper_area: 0.5

configs/data_gen/segm_256.yaml

@@ -0,0 +1,27 @@
+generator_kind: segmentation
+
+mask_generator_kwargs:
+  confidence_threshold: 0.5
+  max_object_area: 0.5
+  min_mask_area: 0.05
+  downsample_levels: 6
+  num_variants_per_mask: 3
+  rigidness_mode: 1
+  max_foreground_coverage: 1  # turn off filtering by overlap
+  max_foreground_intersection: 1  # turn off filtering by overlap
+  max_mask_intersection: 0.2  # the lower this value the higher diversity
+  max_hidden_area: 0.5
+  max_scale_change: 0.25
+  horizontal_flip: True
+  max_vertical_shift: 0.3
+  position_shuffle: True
+
+max_masks_per_image: 1
+
+cropping:
+  out_min_size: 256
+  handle_small_mode: upscale
+  out_square_crop: True
+  crop_min_overlap: 1
+
+max_tamper_area: 0.5

configs/data_gen/segm_512.yaml

@@ -0,0 +1,27 @@
+generator_kind: segmentation
+
+mask_generator_kwargs:
+  confidence_threshold: 0.5
+  max_object_area: 0.5
+  min_mask_area: 0.05
+  downsample_levels: 6
+  num_variants_per_mask: 3
+  rigidness_mode: 1
+  max_foreground_coverage: 1  # turn off filtering by overlap
+  max_foreground_intersection: 1  # turn off filtering by overlap
+  max_mask_intersection: 0.2  # the lower this value the higher diversity
+  max_hidden_area: 0.5
+  max_scale_change: 0.25
+  horizontal_flip: True
+  max_vertical_shift: 0.3
+  position_shuffle: True
+
+max_masks_per_image: 1
+
+cropping:
+  out_min_size: 512
+  handle_small_mode: upscale
+  out_square_crop: True
+  crop_min_overlap: 1
+
+max_tamper_area: 0.5

configs/data_gen/sr_256.yaml

@@ -0,0 +1,25 @@
+generator_kind: random
+
+mask_generator_kwargs:
+  irregular_proba: 0
+  box_proba: 0
+  segm_proba: 0
+  squares_proba: 0
+  superres_proba: 1
+  superres_kwargs:
+    min_step: 2
+    max_step: 4
+    min_width: 1
+    max_width: 3
+
+  variants_n: 5
+
+max_masks_per_image: 1
+
+cropping:
+  out_min_size: 256
+  handle_small_mode: upscale
+  out_square_crop: True
+  crop_min_overlap: 1
+
+max_tamper_area: 1

configs/data_gen/whydra/location/mml-ws01-celeba-hq.yaml

@@ -0,0 +1,5 @@
+# @package _group_
+
+root_dir: /media/inpainting/CelebA-HQ
+out_dir: /media/inpainting/paper_data/CelebA-HQ_val_test
+extension: jpg

configs/data_gen/whydra/location/mml-ws01-ffhq.yaml

@@ -0,0 +1,5 @@
+# @package _group_
+
+root_dir: /media/inpainting/FFHQ/
+out_dir: /media/inpainting/paper_data/FFHQ_val
+extension: png

configs/data_gen/whydra/location/mml-ws01-paris.yaml

@@ -0,0 +1,5 @@
+# @package _group_
+
+root_dir: /media/inpainting/Paris_StreetView_Dataset
+out_dir: /media/inpainting/paper_data/Paris_StreetView_Dataset_val
+extension: png

configs/data_gen/whydra/location/mml7-places.yaml

@@ -0,0 +1,5 @@
+# @package _group_
+
+root_dir: /data/inpainting/Places365
+out_dir: /data/inpainting/paper_data/Places365_val_test
+extension: jpg

configs/data_gen/whydra/random_medium_256.yaml

@@ -0,0 +1,42 @@
+datadir: val_large
+indir: ${location.root_dir}/${datadir}
+outdir: ${location.out_dir}/${datadir}/random_medium_256
+
+n_jobs: 8
+
+generator_kind: random
+
+mask_generator_kwargs:
+  irregular_proba: 1
+  irregular_kwargs:
+    min_times: 4
+    max_times: 5
+    max_width: 50
+    max_angle: 4
+    max_len: 100
+
+  box_proba: 0.3
+  box_kwargs:
+    margin: 0
+    bbox_min_size: 10
+    bbox_max_size: 50
+    max_times: 5
+    min_times: 1
+
+  segm_proba: 0
+  squares_proba: 0
+
+  variants_n: 5
+
+max_masks_per_image: 1
+
+cropping:
+  out_min_size: 256
+  handle_small_mode: upscale
+  out_square_crop: True
+  crop_min_overlap: 1
+
+max_tamper_area: 0.5
+
+defaults:
+  - location: mml7-places

configs/data_gen/whydra/random_medium_512.yaml

@@ -0,0 +1,42 @@
+datadir: val_large
+indir: ${location.root_dir}/${datadir}
+outdir: ${location.out_dir}/${datadir}/random_medium_512
+
+n_jobs: 8
+
+generator_kind: random
+
+mask_generator_kwargs:
+  irregular_proba: 1
+  irregular_kwargs:
+    min_times: 4
+    max_times: 10
+    max_width: 100
+    max_angle: 4
+    max_len: 200
+
+  box_proba: 0.3
+  box_kwargs:
+    margin: 0
+    bbox_min_size: 30
+    bbox_max_size: 150
+    max_times: 5
+    min_times: 1
+
+  segm_proba: 0
+  squares_proba: 0
+
+  variants_n: 5
+
+max_masks_per_image: 1
+
+cropping:
+  out_min_size: 512
+  handle_small_mode: upscale
+  out_square_crop: True
+  crop_min_overlap: 1
+
+max_tamper_area: 0.5
+
+defaults:
+  - location: mml7-places