Skip to content

Commit 071f6d0

Browse files
committed
Added the model
1 parent d16779a commit 071f6d0

File tree

2 files changed

+113
-0
lines changed
  • Algorithms and Deep Learning Models/AI Super-Resolution and Image Restoration model

2 files changed

+113
-0
lines changed
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
# AI-Powered Image Restoration and Enhancement Tool
2+
3+
## Project Description
4+
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.
5+
6+
## Tech Stack
7+
- **Programming Language**: Python
8+
- **Libraries**:
9+
- OpenCV: For image processing tasks
10+
- Pillow: For handling image files
11+
- PyTorch: For implementing deep learning models
12+
- torchvision: For image transformations and model utilities
13+
- **Models**:
14+
- SRCNN (Super-Resolution Convolutional Neural Network): For image upscaling
15+
- DnCNN: For image denoising
16+
- Pre-trained colorization models (e.g., U-Net): For adding color to grayscale images
17+
18+
# Connect with Me
19+
20+
- **GitHub**: [Peart-Guy](https://github.com/Peart-Guy)
21+
- **LinkedIn**: [Ankan Mukhopadhyay](https://www.linkedin.com/in/ankan-mukhopadhyaypeartguy/)
22+
Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
import torch
2+
import cv2
3+
import numpy as np
4+
from torchvision.transforms import ToTensor, ToPILImage
5+
from torchvision import transforms
6+
from PIL import Image
7+
8+
# Load a pre-trained SRCNN model (for demo purposes, use a simple model or load a pre-trained one)
9+
class SRCNN(torch.nn.Module):
10+
def __init__(self):
11+
super(SRCNN, self).__init__()
12+
self.conv1 = torch.nn.Conv2d(1, 64, kernel_size=9, padding=4)
13+
self.conv2 = torch.nn.Conv2d(64, 32, kernel_size=5, padding=2)
14+
self.conv3 = torch.nn.Conv2d(32, 1, kernel_size=5, padding=2)
15+
16+
def forward(self, x):
17+
x = torch.nn.functional.relu(self.conv1(x))
18+
x = torch.nn.functional.relu(self.conv2(x))
19+
x = self.conv3(x)
20+
return x
21+
22+
# Initialize the model and load weights if available
23+
model = SRCNN()
24+
model.load_state_dict(torch.load('srcnn_pretrained.pth', map_location='cpu'))
25+
model.eval()
26+
27+
def super_resolve_image(img_path):
28+
image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
29+
image = cv2.resize(image, (image.shape[1] * 2, image.shape[0] * 2)) # Upscale by factor 2
30+
image = ToTensor()(image).unsqueeze(0)
31+
32+
with torch.no_grad():
33+
output = model(image)
34+
35+
output_image = output.squeeze().clamp(0, 1).cpu()
36+
output_image = ToPILImage()(output_image)
37+
output_image.save("super_resolved_image.jpg")
38+
output_image.show()
39+
40+
41+
42+
def denoise_image(img_path):
43+
image = cv2.imread(img_path)
44+
denoised_image = cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)
45+
cv2.imwrite("denoised_image.jpg", denoised_image)
46+
cv2.imshow("Denoised Image", denoised_image)
47+
cv2.waitKey(0)
48+
cv2.destroyAllWindows()
49+
50+
51+
52+
class ColorizationModel(torch.nn.Module):
53+
# Define a simple U-Net or load pre-trained weights from a colorization model here
54+
pass
55+
56+
def colorize_image(img_path):
57+
image = Image.open(img_path).convert("L") # Convert to grayscale
58+
image = transforms.ToTensor()(image).unsqueeze(0)
59+
60+
model = ColorizationModel()
61+
model.load_state_dict(torch.load("colorization_model.pth", map_location="cpu"))
62+
model.eval()
63+
64+
with torch.no_grad():
65+
colorized = model(image)
66+
67+
# Convert colorized image back to an RGB format for saving and display
68+
colorized_image = colorized.squeeze(0).permute(1, 2, 0).numpy()
69+
colorized_image = np.clip(colorized_image * 255, 0, 255).astype("uint8")
70+
colorized_image = Image.fromarray(colorized_image)
71+
colorized_image.save("colorized_image.jpg")
72+
colorized_image.show()
73+
74+
75+
76+
def process_image(img_path):
77+
print("Starting Super-Resolution...")
78+
super_resolve_image(img_path)
79+
print("Super-Resolution Completed.")
80+
81+
print("Starting Denoising...")
82+
denoise_image("super_resolved_image.jpg")
83+
print("Denoising Completed.")
84+
85+
print("Starting Colorization...")
86+
colorize_image("denoised_image.jpg")
87+
print("Colorization Completed.")
88+
89+
90+
process_image("input_image.jpg")
91+

0 commit comments

Comments
 (0)