example files

app.py

@@ -0,0 +1,42 @@
+import os
+os.system("gdown https://drive.google.com/uc?id=1-95IOJ-2y9BtmABiffIwndPqNZD_gLnV")
+os.system("unzip big-lama.zip")
+import cv2
+import paddlehub as hub
+import gradio as gr
+import torch
+from PIL import Image, ImageOps
+import numpy as np
+os.mkdir("data")
+os.mkdir("dataout")
+model = hub.Module(name='U2Net')
+def infer(img,mask,option):
+  img = ImageOps.contain(img, (700,700))
+  width, height = img.size
+  img.save("./data/data.png")
+  if option == "automatic (U2net)":
+      result = model.Segmentation(
+          images=[cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)],
+          paths=None,
+          batch_size=1,
+          input_size=320,
+          output_dir='output',
+          visualization=True)
+      im = Image.fromarray(result[0]['mask'])
+  else:
+      mask = mask.resize((width,height))
+      im = mask
+  im.save("./data/data_mask.png")
+  os.system('python predict.py model.path=/home/user/app/big-lama/ indir=/home/user/app/data/ outdir=/home/user/app/dataout/ device=cpu')
+  return "./dataout/data_mask.png",im
+  
+inputs = [gr.inputs.Image(type='pil', label="Original Image"),gr.inputs.Image(type='pil',source="canvas", label="Mask",invert_colors=True),gr.inputs.Radio(choices=["automatic (U2net)","manual"], type="value", default="manual", label="Masking option")]
+outputs = [gr.outputs.Image(type="file",label="output"),gr.outputs.Image(type="pil",label="Mask")]
+title = "LaMa Image Inpainting"
+description = "Gradio demo for LaMa: Resolution-robust Large Mask Inpainting with Fourier Convolutions. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below. Masks are generated by U^2net"
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2109.07161' target='_blank'>Resolution-robust Large Mask Inpainting with Fourier Convolutions</a> | <a href='https://github.com/saic-mdal/lama' target='_blank'>Github Repo</a></p>"
+examples = [
+  ['person512.png',"canvas.png","automatic (U2net)"],
+  ['person512.png',"maskexam.png","manual"]
+]
+gr.Interface(infer, inputs, outputs, title=title, description=description, article=article, examples=examples).launch(enable_queue=True,cache_examples=True)

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