Upload 267 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+lama/saicinpainting/evaluation/masks/countless/images/gcim.jpg filter=lfs diff=lfs merge=lfs -text

lama/.gitignore

@@ -0,0 +1,137 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# temporary files
+## IDEA
+.idea/
+## vscode
+.vscode/
+## vim
+*.sw?

lama/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [2021] Samsung Research
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

lama/README.md

@@ -0,0 +1,464 @@
+# 🦙 LaMa: Resolution-robust Large Mask Inpainting with Fourier Convolutions
+
+by Roman Suvorov, Elizaveta Logacheva, Anton Mashikhin, 
+Anastasia Remizova, Arsenii Ashukha, Aleksei Silvestrov, Naejin Kong, Harshith Goka, Kiwoong Park, Victor Lempitsky.
+
+<p align="center" "font-size:30px;">
+  🔥🔥🔥
+  <br>
+  <b>
+LaMa generalizes surprisingly well to much higher resolutions (~2k❗️) than it saw during training (256x256), and achieves the excellent performance even in challenging scenarios, e.g. completion of periodic structures.</b>
+</p>
+
+[[Project page](https://advimman.github.io/lama-project/)] [[arXiv](https://arxiv.org/abs/2109.07161)] [[Supplementary](https://ashukha.com/projects/lama_21/lama_supmat_2021.pdf)] [[BibTeX](https://senya-ashukha.github.io/projects/lama_21/paper.txt)] [[Casual GAN Papers Summary](https://www.casualganpapers.com/large-masks-fourier-convolutions-inpainting/LaMa-explained.html)]
+ 
+<p align="center">
+  <a href="https://colab.research.google.com/drive/15KTEIScUbVZtUP6w2tCDMVpE-b1r9pkZ?usp=drive_link">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg"/>
+  </a>
+      <br>
+   Try out in Google Colab
+</p>
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/senya-ashukha/senya-ashukha.github.io/master/projects/lama_21/ezgif-4-0db51df695a8.gif" />
+</p>
+
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/senya-ashukha/senya-ashukha.github.io/master/projects/lama_21/gif_for_lightning_v1_white.gif" />
+</p>
+
+# LaMa development
+(Feel free to share your paper by creating an issue)
+- https://github.com/geekyutao/Inpaint-Anything --- Inpaint Anything: Segment Anything Meets Image Inpainting
+<p align="center">
+  <img src="https://raw.githubusercontent.com/geekyutao/Inpaint-Anything/main/example/MainFramework.png" />
+</p>
+
+- [Feature Refinement to Improve High Resolution Image Inpainting](https://arxiv.org/abs/2206.13644) / [video](https://www.youtube.com/watch?v=gEukhOheWgE) / code https://github.com/advimman/lama/pull/112 / by Geomagical Labs ([geomagical.com](geomagical.com))
+<p align="center">
+  <img src="https://raw.githubusercontent.com/senya-ashukha/senya-ashukha.github.io/master/images/FeatureRefinement.png" />
+</p>
+
+# Non-official 3rd party apps:
+(Feel free to share your app/implementation/demo by creating an issue)
+
+- https://github.com/enesmsahin/simple-lama-inpainting - a simple pip package for LaMa inpainting.
+- https://github.com/mallman/CoreMLaMa - Apple's Core ML model format
+- [https://cleanup.pictures](https://cleanup.pictures/) - a simple interactive object removal tool by [@cyrildiagne](https://twitter.com/cyrildiagne)
+    - [lama-cleaner](https://github.com/Sanster/lama-cleaner) by [@Sanster](https://github.com/Sanster/lama-cleaner) is a self-host version of [https://cleanup.pictures](https://cleanup.pictures/)
+- Integrated to [Huggingface Spaces](https://huggingface.co/spaces) with [Gradio](https://github.com/gradio-app/gradio). See demo: [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/akhaliq/lama) by [@AK391](https://github.com/AK391)
+- Telegram bot [@MagicEraserBot](https://t.me/MagicEraserBot) by [@Moldoteck](https://github.com/Moldoteck), [code](https://github.com/Moldoteck/MagicEraser)
+- [Auto-LaMa](https://github.com/andy971022/auto-lama) = DE:TR object detection + LaMa inpainting by [@andy971022](https://github.com/andy971022)
+- [LAMA-Magic-Eraser-Local](https://github.com/zhaoyun0071/LAMA-Magic-Eraser-Local) = a standalone inpainting application built with PyQt5 by [@zhaoyun0071](https://github.com/zhaoyun0071)
+- [Hama](https://www.hama.app/) - object removal with a smart brush which simplifies mask drawing.
+- [ModelScope](https://www.modelscope.cn/models/damo/cv_fft_inpainting_lama/summary) = the largest Model Community in Chinese by  [@chenbinghui1](https://github.com/chenbinghui1).
+- [LaMa with MaskDINO](https://github.com/qwopqwop200/lama-with-maskdino) = MaskDINO object detection + LaMa inpainting with refinement by [@qwopqwop200](https://github.com/qwopqwop200).
+- [CoreMLaMa](https://github.com/mallman/CoreMLaMa) - a script to convert Lama Cleaner's port of LaMa to Apple's Core ML model format.
+
+# Environment setup
+
+Clone the repo:
+`git clone https://github.com/advimman/lama.git`
+
+There are three options of an environment:
+
+1. Python virtualenv:
+
+    ```
+    virtualenv inpenv --python=/usr/bin/python3
+    source inpenv/bin/activate
+    pip install torch==1.8.0 torchvision==0.9.0
+    
+    cd lama
+    pip install -r requirements.txt 
+    ```
+
+2. Conda
+    
+    ```
+    % Install conda for Linux, for other OS download miniconda at https://docs.conda.io/en/latest/miniconda.html
+    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p $HOME/miniconda
+    $HOME/miniconda/bin/conda init bash
+
+    cd lama
+    conda env create -f conda_env.yml
+    conda activate lama
+    conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch -y
+    pip install pytorch-lightning==1.2.9
+    ```
+ 
+3. Docker: No actions are needed 🎉.
+
+# Inference <a name="prediction"></a>
+
+Run
+```
+cd lama
+export TORCH_HOME=$(pwd) && export PYTHONPATH=$(pwd)
+```
+
+**1. Download pre-trained models**
+
+The best model (Places2, Places Challenge):
+    
+```    
+curl -LJO https://huggingface.co/smartywu/big-lama/resolve/main/big-lama.zip
+unzip big-lama.zip
+```
+
+All models (Places & CelebA-HQ):
+
+```
+download [https://drive.google.com/drive/folders/1B2x7eQDgecTL0oh3LSIBDGj0fTxs6Ips?usp=drive_link]
+unzip lama-models.zip
+```
+
+**2. Prepare images and masks**
+
+Download test images:
+
+```
+unzip LaMa_test_images.zip
+```
+<details>
+ <summary>OR prepare your data:</summary>
+1) Create masks named as `[images_name]_maskXXX[image_suffix]`, put images and masks in the same folder. 
+
+- You can use the [script](https://github.com/advimman/lama/blob/main/bin/gen_mask_dataset.py) for random masks generation. 
+- Check the format of the files:
+    ```    
+    image1_mask001.png
+    image1.png
+    image2_mask001.png
+    image2.png
+    ```
+
+2) Specify `image_suffix`, e.g. `.png` or `.jpg` or `_input.jpg` in `configs/prediction/default.yaml`.
+
+</details>
+
+
+**3. Predict**
+
+On the host machine:
+
+    python3 bin/predict.py model.path=$(pwd)/big-lama indir=$(pwd)/LaMa_test_images outdir=$(pwd)/output
+
+**OR** in the docker
+  
+The following command will pull the docker image from Docker Hub and execute the prediction script
+```
+bash docker/2_predict.sh $(pwd)/big-lama $(pwd)/LaMa_test_images $(pwd)/output device=cpu
+```
+Docker cuda:
+```
+bash docker/2_predict_with_gpu.sh $(pwd)/big-lama $(pwd)/LaMa_test_images $(pwd)/output
+```
+
+**4. Predict with Refinement**
+
+On the host machine:
+
+    python3 bin/predict.py refine=True model.path=$(pwd)/big-lama indir=$(pwd)/LaMa_test_images outdir=$(pwd)/output
+
+# Train and Eval
+
+Make sure you run:
+
+```
+cd lama
+export TORCH_HOME=$(pwd) && export PYTHONPATH=$(pwd)
+```
+
+Then download models for _perceptual loss_:
+
+    mkdir -p ade20k/ade20k-resnet50dilated-ppm_deepsup/
+    wget -P ade20k/ade20k-resnet50dilated-ppm_deepsup/ http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth
+
+
+## Places
+
+⚠️ NB: FID/SSIM/LPIPS metric values for Places that we see in LaMa paper are computed on 30000 images that we produce in evaluation section below.
+For more details on evaluation data check [[Section 3. Dataset splits in Supplementary](https://ashukha.com/projects/lama_21/lama_supmat_2021.pdf#subsection.3.1)]  ⚠️
+
+On the host machine:
+
+    # Download data from http://places2.csail.mit.edu/download.html
+    # Places365-Standard: Train(105GB)/Test(19GB)/Val(2.1GB) from High-resolution images section
+    wget http://data.csail.mit.edu/places/places365/train_large_places365standard.tar
+    wget http://data.csail.mit.edu/places/places365/val_large.tar
+    wget http://data.csail.mit.edu/places/places365/test_large.tar
+
+    # Unpack train/test/val data and create .yaml config for it
+    bash fetch_data/places_standard_train_prepare.sh
+    bash fetch_data/places_standard_test_val_prepare.sh
+    
+    # Sample images for test and viz at the end of epoch
+    bash fetch_data/places_standard_test_val_sample.sh
+    bash fetch_data/places_standard_test_val_gen_masks.sh
+
+    # Run training
+    python3 bin/train.py -cn lama-fourier location=places_standard
+
+    # To evaluate trained model and report metrics as in our paper
+    # we need to sample previously unseen 30k images and generate masks for them
+    bash fetch_data/places_standard_evaluation_prepare_data.sh
+    
+    # Infer model on thick/thin/medium masks in 256 and 512 and run evaluation 
+    # like this:
+    python3 bin/predict.py \
+    model.path=$(pwd)/experiments/<user>_<date:time>_lama-fourier_/ \
+    indir=$(pwd)/places_standard_dataset/evaluation/random_thick_512/ \
+    outdir=$(pwd)/inference/random_thick_512 model.checkpoint=last.ckpt
+
+    python3 bin/evaluate_predicts.py \
+    $(pwd)/configs/eval2_gpu.yaml \
+    $(pwd)/places_standard_dataset/evaluation/random_thick_512/ \
+    $(pwd)/inference/random_thick_512 \
+    $(pwd)/inference/random_thick_512_metrics.csv
+
+    
+    
+Docker: TODO
+    
+## CelebA
+On the host machine:
+
+    # Make shure you are in lama folder
+    cd lama
+    export TORCH_HOME=$(pwd) && export PYTHONPATH=$(pwd)
+
+    # Download CelebA-HQ dataset
+    # Download data256x256.zip from https://drive.google.com/drive/folders/11Vz0fqHS2rXDb5pprgTjpD7S2BAJhi1P
+    
+    # unzip & split into train/test/visualization & create config for it
+    bash fetch_data/celebahq_dataset_prepare.sh
+
+    # generate masks for test and visual_test at the end of epoch
+    bash fetch_data/celebahq_gen_masks.sh
+
+    # Run training
+    python3 bin/train.py -cn lama-fourier-celeba data.batch_size=10
+
+    # Infer model on thick/thin/medium masks in 256 and run evaluation 
+    # like this:
+    python3 bin/predict.py \
+    model.path=$(pwd)/experiments/<user>_<date:time>_lama-fourier-celeba_/ \
+    indir=$(pwd)/celeba-hq-dataset/visual_test_256/random_thick_256/ \
+    outdir=$(pwd)/inference/celeba_random_thick_256 model.checkpoint=last.ckpt
+    
+    
+Docker: TODO
+
+## Places Challenge 
+
+On the host machine:
+
+    # This script downloads multiple .tar files in parallel and unpacks them
+    # Places365-Challenge: Train(476GB) from High-resolution images (to train Big-Lama) 
+    bash places_challenge_train_download.sh
+    
+    TODO: prepare
+    TODO: train 
+    TODO: eval
+      
+Docker: TODO
+
+## Create your data
+
+Please check bash scripts for data preparation and mask generation from CelebaHQ section,
+if you stuck at one of the following steps.
+
+
+On the host machine:
+
+    # Make shure you are in lama folder
+    cd lama
+    export TORCH_HOME=$(pwd) && export PYTHONPATH=$(pwd)
+
+    # You need to prepare following image folders:
+    $ ls my_dataset
+    train
+    val_source # 2000 or more images
+    visual_test_source # 100 or more images
+    eval_source # 2000 or more images
+
+    # LaMa generates random masks for the train data on the flight,
+    # but needs fixed masks for test and visual_test for consistency of evaluation.
+
+    # Suppose, we want to evaluate and pick best models 
+    # on 512x512 val dataset  with thick/thin/medium masks 
+    # And your images have .jpg extention:
+
+    python3 bin/gen_mask_dataset.py \
+    $(pwd)/configs/data_gen/random_<size>_512.yaml \ # thick, thin, medium
+    my_dataset/val_source/ \
+    my_dataset/val/random_<size>_512.yaml \# thick, thin, medium
+    --ext jpg
+
+    # So the mask generator will: 
+    # 1. resize and crop val images and save them as .png
+    # 2. generate masks
+    
+    ls my_dataset/val/random_medium_512/
+    image1_crop000_mask000.png
+    image1_crop000.png
+    image2_crop000_mask000.png
+    image2_crop000.png
+    ...
+
+    # Generate thick, thin, medium masks for visual_test folder:
+
+    python3 bin/gen_mask_dataset.py \
+    $(pwd)/configs/data_gen/random_<size>_512.yaml \  #thick, thin, medium
+    my_dataset/visual_test_source/ \
+    my_dataset/visual_test/random_<size>_512/ \ #thick, thin, medium
+    --ext jpg
+    
+
+    ls my_dataset/visual_test/random_thick_512/
+    image1_crop000_mask000.png
+    image1_crop000.png
+    image2_crop000_mask000.png
+    image2_crop000.png
+    ...
+
+    # Same process for eval_source image folder:
+    
+    python3 bin/gen_mask_dataset.py \
+    $(pwd)/configs/data_gen/random_<size>_512.yaml \  #thick, thin, medium
+    my_dataset/eval_source/ \
+    my_dataset/eval/random_<size>_512/ \ #thick, thin, medium
+    --ext jpg
+    
+
+
+    # Generate location config file which locate these folders:
+    
+    touch my_dataset.yaml
+    echo "data_root_dir: $(pwd)/my_dataset/" >> my_dataset.yaml
+    echo "out_root_dir: $(pwd)/experiments/" >> my_dataset.yaml
+    echo "tb_dir: $(pwd)/tb_logs/" >> my_dataset.yaml
+    mv my_dataset.yaml ${PWD}/configs/training/location/
+
+
+    # Check data config for consistency with my_dataset folder structure:
+    $ cat ${PWD}/configs/training/data/abl-04-256-mh-dist
+    ...
+    train:
+      indir: ${location.data_root_dir}/train
+      ...
+    val:
+      indir: ${location.data_root_dir}/val
+      img_suffix: .png
+    visual_test:
+      indir: ${location.data_root_dir}/visual_test
+      img_suffix: .png
+
+
+    # Run training
+    python3 bin/train.py -cn lama-fourier location=my_dataset data.batch_size=10
+
+    # Evaluation: LaMa training procedure picks best few models according to 
+    # scores on my_dataset/val/ 
+
+    # To evaluate one of your best models (i.e. at epoch=32) 
+    # on previously unseen my_dataset/eval do the following 
+    # for thin, thick and medium:
+
+    # infer:
+    python3 bin/predict.py \
+    model.path=$(pwd)/experiments/<user>_<date:time>_lama-fourier_/ \
+    indir=$(pwd)/my_dataset/eval/random_<size>_512/ \
+    outdir=$(pwd)/inference/my_dataset/random_<size>_512 \
+    model.checkpoint=epoch32.ckpt
+
+    # metrics calculation:
+    python3 bin/evaluate_predicts.py \
+    $(pwd)/configs/eval2_gpu.yaml \
+    $(pwd)/my_dataset/eval/random_<size>_512/ \
+    $(pwd)/inference/my_dataset/random_<size>_512 \
+    $(pwd)/inference/my_dataset/random_<size>_512_metrics.csv
+
+    
+**OR** in the docker:
+
+    TODO: train
+    TODO: eval
+    
+# Hints
+
+### Generate different kinds of masks
+The following command will execute a script that generates random masks.
+
+    bash docker/1_generate_masks_from_raw_images.sh \
+        configs/data_gen/random_medium_512.yaml \
+        /directory_with_input_images \
+        /directory_where_to_store_images_and_masks \
+        --ext png
+
+The test data generation command stores images in the format,
+which is suitable for [prediction](#prediction).
+
+The table below describes which configs we used to generate different test sets from the paper.
+Note that we *do not fix a random seed*, so the results will be slightly different each time.
+
+|        | Places 512x512         | CelebA 256x256         |
+|--------|------------------------|------------------------|
+| Narrow | random_thin_512.yaml   | random_thin_256.yaml   |
+| Medium | random_medium_512.yaml | random_medium_256.yaml |
+| Wide   | random_thick_512.yaml  | random_thick_256.yaml  |
+
+Feel free to change the config path (argument #1) to any other config in `configs/data_gen` 
+or adjust config files themselves.
+
+### Override parameters in configs
+Also you can override parameters in config like this:
+
+    python3 bin/train.py -cn <config> data.batch_size=10 run_title=my-title
+
+Where .yaml file extension is omitted
+
+### Models options 
+Config names for models from paper (substitude into the training command): 
+
+    * big-lama
+    * big-lama-regular
+    * lama-fourier
+    * lama-regular
+    * lama_small_train_masks
+
+Which are seated in configs/training/folder
+
+### Links
+- All the data (models, test images, etc.) https://disk.yandex.ru/d/AmdeG-bIjmvSug
+- Test images from the paper https://disk.yandex.ru/d/xKQJZeVRk5vLlQ
+- The pre-trained models https://disk.yandex.ru/d/EgqaSnLohjuzAg
+- The models for perceptual loss https://disk.yandex.ru/d/ncVmQlmT_kTemQ
+- Our training logs are available at https://disk.yandex.ru/d/9Bt1wNSDS4jDkQ
+
+
+### Training time & resources
+
+TODO
+
+## Acknowledgments
+
+* Segmentation code and models if form [CSAILVision](https://github.com/CSAILVision/semantic-segmentation-pytorch).
+* LPIPS metric is from [richzhang](https://github.com/richzhang/PerceptualSimilarity)
+* SSIM is from [Po-Hsun-Su](https://github.com/Po-Hsun-Su/pytorch-ssim)
+* FID is from [mseitzer](https://github.com/mseitzer/pytorch-fid)
+
+## Citation
+If you found this code helpful, please consider citing: 
+```
+@article{suvorov2021resolution,
+  title={Resolution-robust Large Mask Inpainting with Fourier Convolutions},
+  author={Suvorov, Roman and Logacheva, Elizaveta and Mashikhin, Anton and Remizova, Anastasia and Ashukha, Arsenii and Silvestrov, Aleksei and Kong, Naejin and Goka, Harshith and Park, Kiwoong and Lempitsky, Victor},
+  journal={arXiv preprint arXiv:2109.07161},
+  year={2021}
+}
+```

lama/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())

lama/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())

lama/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())

lama/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"

lama/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())

lama/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)

lama/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())

lama/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())

lama/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())

lama/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())

lama/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()

lama/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())

lama/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())

lama/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)

lama/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"
+#

lama/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"

lama/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)

lama/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

lama/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

lama/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

lama/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

lama/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

lama/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

lama/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

lama/bin/predict.py

@@ -0,0 +1,104 @@
+#!/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
+from saicinpainting.evaluation.refinement import refine_predict
+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:
+        if sys.platform != 'win32':
+            register_debug_signal_handlers()  # kill -10 <pid> will result in traceback dumped into log
+
+        device = torch.device("cpu")
+
+        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
+        train_config.visualizer.kind = 'noop'
+
+        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()
+        if not predict_config.get('refine', False):
+            model.to(device)
+
+        if not predict_config.indir.endswith('/'):
+            predict_config.indir += '/'
+
+        dataset = make_default_val_dataset(predict_config.indir, **predict_config.dataset)
+        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 = default_collate([dataset[img_i]])
+            if predict_config.get('refine', False):
+                assert 'unpad_to_size' in batch, "Unpadded size is required for the refinement"
+                # image unpadding is taken care of in the refiner, so that output image
+                # is same size as the input image
+                cur_res = refine_predict(batch, model, **predict_config.refiner)
+                cur_res = cur_res[0].permute(1,2,0).detach().cpu().numpy()
+            else:
+                with torch.no_grad():
+                    batch = move_to_device(batch, 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()
+                    unpad_to_size = batch.get('unpad_to_size', None)
+                    if unpad_to_size is not None:
+                        orig_height, orig_width = unpad_to_size
+                        cur_res = cur_res[:orig_height, :orig_width]
+
+            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()

lama/bin/predict_inner_features.py

@@ -0,0 +1,120 @@
+#!/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:
+        if sys.platform != 'win32':
+            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()

lama/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())

lama/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())

lama/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())

lama/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())

lama/bin/to_jit.py

@@ -0,0 +1,76 @@
+import os
+from pathlib import Path
+
+import hydra
+import torch
+import yaml
+from omegaconf import OmegaConf
+from torch import nn
+
+from saicinpainting.training.trainers import load_checkpoint
+from saicinpainting.utils import register_debug_signal_handlers
+
+
+class JITWrapper(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, image, mask):
+        batch = {
+            "image": image,
+            "mask": mask
+        }
+        out = self.model(batch)
+        return out["inpainted"]
+
+
+@hydra.main(config_path="../configs/prediction", config_name="default.yaml")
+def main(predict_config: OmegaConf):
+    if sys.platform != 'win32':
+        register_debug_signal_handlers()  # kill -10 <pid> will result in traceback dumped into log
+
+    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
+    train_config.visualizer.kind = "noop"
+
+    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.eval()
+    jit_model_wrapper = JITWrapper(model)
+
+    image = torch.rand(1, 3, 120, 120)
+    mask = torch.rand(1, 1, 120, 120)
+    output = jit_model_wrapper(image, mask)
+
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+
+    image = image.to(device)
+    mask = mask.to(device)
+    traced_model = torch.jit.trace(jit_model_wrapper, (image, mask), strict=False).to(device)
+
+    save_path = Path(predict_config.save_path)
+    save_path.parent.mkdir(parents=True, exist_ok=True)
+
+    print(f"Saving big-lama.pt model to {save_path}")
+    traced_model.save(save_path)
+
+    print(f"Checking jit model output...")
+    jit_model = torch.jit.load(str(save_path))
+    jit_output = jit_model(image, mask)
+    diff = (output - jit_output).abs().sum()
+    print(f"diff: {diff}")
+
+
+if __name__ == "__main__":
+    main()

lama/bin/train.py

@@ -0,0 +1,73 @@
+#!/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)
+
+        if sys.platform != 'win32':
+            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()

lama/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

lama/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

lama/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

lama/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

lama/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

lama/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

lama/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

lama/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

lama/configs/debug_mask_gen.yaml

@@ -0,0 +1,5 @@
+img_ext: .jpg
+
+gen_kwargs:
+  mask_size: 200
+  step: 0.5

lama/configs/eval1.yaml

@@ -0,0 +1,6 @@
+evaluator_kwargs:
+  batch_size: 8
+
+dataset_kwargs:
+  img_suffix: .png
+  inpainted_suffix: .jpg

lama/configs/eval2.yaml

@@ -0,0 +1,7 @@
+evaluator_kwargs:
+  batch_size: 8
+  device: cuda
+
+dataset_kwargs:
+  img_suffix: .png
+  inpainted_suffix: .png

lama/configs/eval2_cpu.yaml

@@ -0,0 +1,7 @@
+evaluator_kwargs:
+  batch_size: 8
+  device: cpu
+
+dataset_kwargs:
+  img_suffix: .png
+  inpainted_suffix: .png

lama/configs/eval2_gpu.yaml

@@ -0,0 +1,6 @@
+evaluator_kwargs:
+  batch_size: 8
+  
+dataset_kwargs:
+  img_suffix: .png
+  inpainted_suffix: .png