ONNX MST Image Inpainting

Python scripts for performing Image Inpainting using the MST model in ONNX

Original image: https://es.wikipedia.org/wiki/Archivo:Perro_Pipper_en_Merida_-_Pipperontour.jpg

Requirements

• Check the requirements.txt file.
• For ONNX, if you have a NVIDIA GPU, then install the onnxruntime-gpu, otherwise use the onnxruntime library.

Installation

⚠️ Make sure to add the recursive when cloning

git clone https://github.com/ibaiGorordo/ONNX-MST-Image-Inpainting.git --recursive
cd ONNX-MST-Image-Inpainting
pip install -r requirements.txt

ONNX Runtime

For Nvidia GPU computers: pip install onnxruntime-gpu

Otherwise: pip install onnxruntime

ONNX model

Convert the model to ONNX using the Colab repository below and save it into the models folder

• Convert MST Inpainting to ONNX
• The License of the models is MIT: License

Examples

• Image Inpainting (mask already available):

python image_inpaint.py.py

• Sketch mask and Inpaint:

python sketch_and_inpaint.py.py

Mask Sketcher Usage

• Left Mouse Button: Draw Mask
• Middle/Wheel Mouse Button: Remove Mask
• Bottom slider: Control mask drawing tool size
• Press any key to stop sketching

inpaint.mp4

References:

• MST_inpainting: https://github.com/ewrfcas/MST_inpainting
• HAWP: https://github.com/cherubicXN/hawp
• PINTO0309's model zoo: https://github.com/PINTO0309/PINTO_model_zoo
• PINTO0309's model conversion tool: https://github.com/PINTO0309/openvino2tensorflow
• MST Original paper: https://arxiv.org/abs/2103.15087