Image Restoration using U-Net GAN with Self-Attention

Table of Contents

• Introduction
• Requirements
• How to use
• Preview
• Team
• Contribution
• Improvements

Abstract

This project implements a U-Net-based GAN model for image restoration tasks such as deblurring, denoising, and artifact removal. The model enhances traditional U-Net architecture using Residual Blocks and a Self-Attention mechanism to capture both local and global image features. A user-friendly web app is also developed to allow users to interactively upload and restore corrupted images.

Requirements

Package	Version	Link
Python	3.10.0	Python 3.10.0
PyTorch	2.1.0+cu118	PyTorch
torchvision	0.16.0	torchvision
streamlit	1.25.0	Streamlit
Pillow	10.0.0	Pillow
NumPy	1.26.2	NumPy
OpenCV	4.8.0	OpenCV
matplotlib	3.8.0	matplotlib

Installation and usage

1. Clone the repository:

git clone https://github.com/yourusername/your-repo-name.git
cd your-repo-name

2. (Optional) Create and activate a virtual environment:

python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

3. Install the dependencies:

pip install -r requirements.txt

4. Make sure the model checkpoint best_gopro_unet_gan.pth is placed inside an outputs/ folder:

your-repo-name/
├── app.py
├── model.py
├── outputs/
│   └── best_gopro_unet_gan.pth

5. Run the Streamlit app:

streamlit run app.py

Preview

Screenshots of the project:

Team details

Team Number:

14

Senior Mentor:

Ekanth Sai

Junior Mentor:

Bhuvan

Team Member 1:

Aditya Pillai

Team Member 2:

Vishal

Contribution

This section provides instructions and details on how to submit a contribution via a pull request. It is important to follow these guidelines to make sure your pull request is accepted.

1. Before choosing to propose changes to this project, it is advisable to go through the README file of the project to understand the philosophy and the motive behind this project. The pull request should align with the philosophy and the motive of the original poster.
2. To add your changes, make sure the programming language you are proposing the changes in is the same as used in the project. The versions of the programming language and the libraries (if any) should also match the original code.
3. Write documentation on the changes you are proposing. This should include problems you have noticed (if any), the changes you want to propose, reasons for the changes, and sample test cases.
4. Submit a pull request via Git etiquettes.

Improvements

• Add dataset uploader to enable batch restoration.
• Integrate image quality metrics like PSNR/SSIM in the app.
• Provide multiple model choices: baseline U-Net, ResNet U-Net, Transformer U-Net.
• Add GPU support toggle in Streamlit UI.
• Save inference logs or upload outputs directly to cloud.
• Include a mobile-responsive version of the app.