An image reconstruction project using concepts of partial convolution

Explore the docs »

View Demo · Report Bug · Request Feature

About The Project

This is a very unique project in itself as very recently the concept of partial convolutions was introduced to the world [december 2018], before partial convolution we used to use algorithms such as PatchMatch, Iizuka et al, Yu et al and even the KNN for image reconstruction but there were two huge setbacks to these algorithms, which were :

1. Since the algorithms were ignorant of the different objects in the image, they often tended to smoothen the whole reconstruction, which was good to human eyes but lacked the actual information in terms of object segregation.

2. Another limitation of many recent approaches is the focus on rectangular-shaped holes, often assumed to be the center of the image. We find these limitations may lead to overfitting to the rectangular holes and ultimately limit the utility of these models in application

Architecture: The Image depicts flowchart for working of segmentation aware convolutions.

Image taken from Research Paper »

So to overcome the above 1st issue, a segmentation aware approach was used wherein the distorted image as well as the binary mask is inputted to the model as a result of which the model becomes aware of segmentations and different edges of the many objects in the image. Since the model is now aware of the segments it can more accurately segregate between the two objects. To overcome the 2nd issue, various methods have been documented by the authors of the Pconv research paper, but the Random Walk algorithm was chosen to generate a random mask for the model to train upon so that it does not overfit a particular hole point.

gif for iterations over many epochs

Coming on to the aim of FaceInpainting, We saw that if we train the model on large datasets of images containing a single person then maybe we can teach the model to specifically reconstruct the distorted photo of a person. So to test the extent or limit to which we can reconstruct the broken image we trained the model with 50,000 images as training data, 5000 images as test data, and 5000 images as validation data for 1 million iterations. To download the datasets and weights click here

final results shows reconstruction of 8 images , were row 1 shows the distorted image ,row 2 shows the mask thats been applied for distortions , row 4 shows the final model reconstruction , row 5 shows the ground truth image.

On our path to train model , we were able to implement one more splendid feature . Turns out that given enough data and iteration cycles the model has also aquired understanding to reconstruct some important features of human face such as eyes , ears and lips.

Built With

1. Pytorch
2. Opencv
3. Nvidia cuda toolkit

Getting Started

Step 1. Clone the repository.

Step 2. Download the dataset and pretrained weights from Here and place it in the same directory.

Installation

• Python 3.6+

• Install Pytorch

  (for cude 10.2 – GPU)

  pip install torch==1.9.0+cu102 torchvision==0.10.0+cu102 torchaudio===0.9.0 -f https://download.pytorch.org/whl/torch_stable.html

  (for CPU)

  pip install torch torchvision torchaudio

• Install python libraries

  pip install -r requirements.txt

Training

• Start from scratch

  python train.py

• Resume training

  python train.py --resume <weights_path>

Testing

• Run the command line

  python run.py --photo <test_image_path>

• Draw Mask
• Press "s"

Output will be saved in the root directory in result.jpg format.

References

1. Partial convolution research paper »
2. Segmentation aware convolution research paper »
3. reference code Implementation »
4. Base code Implementation »
5. Manual mask generator code reference »

License

Distributed under the MIT License. See LICENSE for more information.

Contact

• Contact Yatharth Kapadia @[email protected]
• Contact Poojan Panchal @ [email protected]