Skip to content

Commit

Permalink
Added the model
Browse files Browse the repository at this point in the history
  • Loading branch information
Peart-Guy committed Nov 5, 2024
1 parent 1444b70 commit 7f2a922
Show file tree
Hide file tree
Showing 2 changed files with 113 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
# 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

# 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'))
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")

0 comments on commit 7f2a922

Please sign in to comment.