|
| 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