Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added the AI-Powered Super-Resolution and Image Restoration model #1104

Merged
merged 3 commits into from
Nov 8, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
# AI-Powered Image Restoration and Enhancement Tool

## Project Description
This project aims to develop an **AI-Powered Image Restoration Tool** that revitalizes low-quality historical photos by upscaling, denoising, and adding realistic color. Using advanced deep learning techniques, it transforms degraded images into vibrant, high-quality visuals while preserving their historical context. This tool is perfect for heritage conservation, family archiving, and historical documentation.

## Tech Stack
- **Programming Language**: Python
- **Libraries**:
- OpenCV: For image processing tasks
- Pillow: For handling image files
- PyTorch: For implementing deep learning models
- torchvision: For image transformations and model utilities
- **Models**:
- SRCNN (Super-Resolution Convolutional Neural Network): For image upscaling
- DnCNN: For image denoising
- Pre-trained colorization models (e.g., U-Net): For adding color to grayscale images

## Datasets
To train or fine-tune the SRCNN model, you can use the following datasets:

1. **DIV2K**: A high-quality dataset for super-resolution with 800 training images.
- [Download DIV2K Dataset](https://data.vision.ee.ethz.ch/cvl/DIV2K/)

2. **Flickr2K**: Contains 2,656 high-resolution images, useful as a complement to DIV2K.
- [Download Flickr2K Dataset](http://cv.snu.ac.kr/research/EDSR/Flickr2K.tar)

3. **BSD300** and **BSD500**: Classical image processing datasets.
- [Download BSD300 Dataset](https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/bsds/BSDS300-images.tgz)
- [Download BSD500 Dataset](https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/bsds/BSDS500.tgz)

4. **Set5 and Set14**: Small datasets often used for testing super-resolution models.
- [Download Set5 & Set14 Datasets](https://github.com/jbhuang0604/SelfExSR/tree/master/data)

## Pre-trained SRCNN Model
To skip the training process, you can use a pre-trained SRCNN model:
- [Download SRCNN Pretrained Model](https://github.com/leftthomas/SRGAN/blob/master/model/srresnet.pth)


# Connect with Me

- **GitHub**: [Peart-Guy](https://github.com/Peart-Guy)
- **LinkedIn**: [Ankan Mukhopadhyay](https://www.linkedin.com/in/ankan-mukhopadhyaypeartguy/)

91 changes: 91 additions & 0 deletions Image processing/AI Super-Resolution and Image Restoration/main.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
import torch
import cv2
import numpy as np
from torchvision.transforms import ToTensor, ToPILImage
from torchvision import transforms
from PIL import Image

# Load a pre-trained SRCNN model (for demo purposes, use a simple model or load a pre-trained one)
class SRCNN(torch.nn.Module):
def __init__(self):
super(SRCNN, self).__init__()
self.conv1 = torch.nn.Conv2d(1, 64, kernel_size=9, padding=4)
self.conv2 = torch.nn.Conv2d(64, 32, kernel_size=5, padding=2)
self.conv3 = torch.nn.Conv2d(32, 1, kernel_size=5, padding=2)

def forward(self, x):
x = torch.nn.functional.relu(self.conv1(x))
x = torch.nn.functional.relu(self.conv2(x))
x = self.conv3(x)
return x

# Initialize the model and load weights if available
model = SRCNN()
model.load_state_dict(torch.load('srcnn_pretrained.pth', map_location='cpu'))
Peart-Guy marked this conversation as resolved.
Show resolved Hide resolved
model.eval()

def super_resolve_image(img_path):
image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
image = cv2.resize(image, (image.shape[1] * 2, image.shape[0] * 2)) # Upscale by factor 2
image = ToTensor()(image).unsqueeze(0)

with torch.no_grad():
output = model(image)

output_image = output.squeeze().clamp(0, 1).cpu()
output_image = ToPILImage()(output_image)
output_image.save("super_resolved_image.jpg")
output_image.show()



def denoise_image(img_path):
image = cv2.imread(img_path)
denoised_image = cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)
cv2.imwrite("denoised_image.jpg", denoised_image)
cv2.imshow("Denoised Image", denoised_image)
cv2.waitKey(0)
cv2.destroyAllWindows()



class ColorizationModel(torch.nn.Module):
# Define a simple U-Net or load pre-trained weights from a colorization model here
pass

def colorize_image(img_path):
image = Image.open(img_path).convert("L") # Convert to grayscale
image = transforms.ToTensor()(image).unsqueeze(0)

model = ColorizationModel()
model.load_state_dict(torch.load("colorization_model.pth", map_location="cpu"))
model.eval()

with torch.no_grad():
colorized = model(image)

# Convert colorized image back to an RGB format for saving and display
colorized_image = colorized.squeeze(0).permute(1, 2, 0).numpy()
colorized_image = np.clip(colorized_image * 255, 0, 255).astype("uint8")
colorized_image = Image.fromarray(colorized_image)
colorized_image.save("colorized_image.jpg")
colorized_image.show()



def process_image(img_path):
print("Starting Super-Resolution...")
super_resolve_image(img_path)
print("Super-Resolution Completed.")

print("Starting Denoising...")
denoise_image("super_resolved_image.jpg")
print("Denoising Completed.")

print("Starting Colorization...")
colorize_image("denoised_image.jpg")
print("Colorization Completed.")


process_image("input_image.jpg")

2 changes: 2 additions & 0 deletions Project-Structure.md
Original file line number Diff line number Diff line change
Expand Up @@ -670,6 +670,8 @@
* [Model](Generative-AI/image-to-3D%20model/model.py)

## Image Processing
* Ai Super-Resolution And Image Restoration
* [Main](Image%20processing/AI%20Super-Resolution%20and%20Image%20Restoration/main.py)
* Cat Dog Cnn Classifier
* [Cat Dog Cnn Classifier](Image%20processing/Cat_Dog_CNN_classifier/Cat_Dog_CNN_classifier.ipynb)
* Lane Line Detection [Open Cv]
Expand Down