1import os
2import torch
3import torch.nn as nn
4from PIL import Image
5from torchvision.transforms import ToTensor
6import numpy as np
7from concurrent.futures import ThreadPoolExecutor
8
9class DenoisingModel(nn.Module):
10 def __init__(self):
11 super(DenoisingModel, self).__init__()
12 self.enc1 = nn.Sequential(
13 nn.Conv2d(3, 64, 3, padding=1),
14 nn.ReLU(),
15 nn.Conv2d(64, 64, 3, padding=1),
16 nn.ReLU()
17 )
18 self.pool1 = nn.MaxPool2d(2, 2)
19
20 self.up1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
21 self.dec1 = nn.Sequential(
22 nn.Conv2d(64, 64, 3, padding=1),
23 nn.ReLU(),
24 nn.Conv2d(64, 3, 3, padding=1)
25 )
26
27 def forward(self, x):
28 e1 = self.enc1(x)
29 p1 = self.pool1(e1)
30 u1 = self.up1(p1)
31 d1 = self.dec1(u1)
32 return d1
33
34def denoise_patch(model, patch):
35 transform = ToTensor()
36 input_patch = transform(patch).unsqueeze(0)
37
38 with torch.no_grad():
39 output_patch = model(input_patch)
40
41 denoised_patch = output_patch.squeeze(0).permute(1, 2, 0).numpy() * 255
42 denoised_patch = np.clip(denoised_patch, 0, 255).astype(np.uint8)
43
44 original_patch = np.array(patch)
45 very_bright_mask = original_patch > 240
46 bright_mask = (original_patch > 220) & (original_patch <= 240)
47
48 denoised_patch[very_bright_mask] = original_patch[very_bright_mask]
49
50 blend_factor = 0.7
51 denoised_patch[bright_mask] = (
52 blend_factor * original_patch[bright_mask] +
53 (1 - blend_factor) * denoised_patch[bright_mask]
54 )
55
56 return denoised_patch
57
58def denoise_image(image_path, model_path, patch_size=256, num_threads=4, overlap=32):
59 model = DenoisingModel()
60 checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
61 model.load_state_dict(checkpoint['model_state_dict'])
62 model.eval()
63
64 # Load and get original image dimensions
65 image = Image.open(image_path).convert("RGB")
66 width, height = image.size
67
68 # Calculate padding needed
69 pad_right = patch_size - (width % patch_size) if width % patch_size != 0 else 0
70 pad_bottom = patch_size - (height % patch_size) if height % patch_size != 0 else 0
71
72 # Add padding with reflection instead of zeros
73 padded_width = width + pad_right
74 padded_height = height + pad_bottom
75
76 # Create padded image using reflection padding
77 padded_image = Image.new("RGB", (padded_width, padded_height))
78 padded_image.paste(image, (0, 0))
79
80 # Fill right border with reflected content
81 if pad_right > 0:
82 right_border = image.crop((width - pad_right, 0, width, height))
83 padded_image.paste(right_border.transpose(Image.FLIP_LEFT_RIGHT), (width, 0))
84
85 # Fill bottom border with reflected content
86 if pad_bottom > 0:
87 bottom_border = image.crop((0, height - pad_bottom, width, height))
88 padded_image.paste(bottom_border.transpose(Image.FLIP_TOP_BOTTOM), (0, height))
89
90 # Fill corner if needed
91 if pad_right > 0 and pad_bottom > 0:
92 corner = image.crop((width - pad_right, height - pad_bottom, width, height))
93 padded_image.paste(corner.transpose(Image.FLIP_LEFT_RIGHT).transpose(Image.FLIP_TOP_BOTTOM),
94 (width, height))
95
96 # Generate patches with positions
97 patches = []
98 positions = []
99 for i in range(0, padded_height, patch_size - overlap):
100 for j in range(0, padded_width, patch_size - overlap):
101 patch = padded_image.crop((j, i, min(j + patch_size, padded_width), min(i + patch_size, padded_height)))
102 patches.append(patch)
103 positions.append((i, j))
104
105 # Process patches in parallel
106 with ThreadPoolExecutor(max_workers=num_threads) as executor:
107 denoised_patches = list(executor.map(lambda p: denoise_patch(model, p), patches))
108
109 # Initialize output arrays
110 denoised_image = np.zeros((padded_height, padded_width, 3), dtype=np.float32)
111 weight_map = np.zeros((padded_height, padded_width), dtype=np.float32)
112
113 # Create smooth blending weights
114 for (i, j), denoised_patch in zip(positions, denoised_patches):
115 patch_height, patch_width, _ = denoised_patch.shape
116 patch_weights = np.ones((patch_height, patch_width), dtype=np.float32)
117 if i > 0:
118 patch_weights[:overlap, :] *= np.linspace(0, 1, overlap)[:, np.newaxis]
119 if j > 0:
120 patch_weights[:, :overlap] *= np.linspace(0, 1, overlap)[np.newaxis, :]
121 if i + patch_height < padded_height:
122 patch_weights[-overlap:, :] *= np.linspace(1, 0, overlap)[:, np.newaxis]
123 if j + patch_width < padded_width:
124 patch_weights[:, -overlap:] *= np.linspace(1, 0, overlap)[np.newaxis, :]
125
126 # Clip the patch values to prevent very bright pixels
127 denoised_patch = np.clip(denoised_patch, 0, 255)
128
129 denoised_image[i:i + patch_height, j:j + patch_width] += (
130 denoised_patch * patch_weights[:, :, np.newaxis]
131 )
132 weight_map[i:i + patch_height, j:j + patch_width] += patch_weights
133
134 # Normalize by weights
135 mask = weight_map > 0
136 denoised_image[mask] = denoised_image[mask] / weight_map[mask, np.newaxis]
137
138 # Crop to original size
139 denoised_image = denoised_image[:height, :width]
140 denoised_image = np.clip(denoised_image, 0, 255).astype(np.uint8)
141
142 # Save the result
143 denoised_image_path = os.path.splitext(image_path)[0] + "_denoised.png"
144 print(f"Saving denoised image to {denoised_image_path}")
145
146 Image.fromarray(denoised_image).save(denoised_image_path)
147
148if __name__ == "__main__":
149 image_path = input("Enter the path of the image: ")
150 model_path = r"path/to/model.pkl"
151 denoise_image(image_path, model_path, num_threads=12)
152 print("Denoising completed.") # Use the number of threads your processor has.)
