From 7f2a922d877e203cde1859fbd03198f7b50701d6 Mon Sep 17 00:00:00 2001
From: Ankan <ankanmukhopadhyay0512@gmail.com>
Date: Tue, 5 Nov 2024 14:54:35 +0000
Subject: [PATCH] Added the model

---
 .../README.md                                 | 22 +++++
 .../main.py                                   | 91 +++++++++++++++++++
 2 files changed, 113 insertions(+)
 create mode 100644 Image processing/AI Super-Resolution and Image Restoration/README.md
 create mode 100644 Image processing/AI Super-Resolution and Image Restoration/main.py

diff --git a/Image processing/AI Super-Resolution and Image Restoration/README.md b/Image processing/AI Super-Resolution and Image Restoration/README.md
new file mode 100644
index 0000000000..80a580f743
--- /dev/null
+++ b/Image processing/AI Super-Resolution and Image Restoration/README.md	
@@ -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/)
+
diff --git a/Image processing/AI Super-Resolution and Image Restoration/main.py b/Image processing/AI Super-Resolution and Image Restoration/main.py
new file mode 100644
index 0000000000..f58d0e5a24
--- /dev/null
+++ b/Image processing/AI Super-Resolution and Image Restoration/main.py	
@@ -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")
+