vericudebuget commited on
Commit
c1f3b43
·
verified ·
1 Parent(s): ee50f63

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +240 -3
README.md CHANGED
@@ -1,3 +1,240 @@
1
- ---
2
- license: gpl-3.0
3
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: gpl-3.0
3
+ tags:
4
+ - img2img
5
+ - denoiser
6
+ - image
7
+ ---
8
+
9
+ # denoise_large_v1
10
+
11
+ denoise_large_v1 is an image denoiser made for images that have a high/medium amount of noise.
12
+
13
+ It performs slightly better than [denoise_medium_v1](https://huggingface.co/vericudebuget/denoise_medium_v1) on most images and car reconstruct a higher level of detail.
14
+
15
+
16
+ ## Model Details
17
+
18
+
19
+
20
+ ### Model Description
21
+
22
+ <!-- Provide a longer summary of what this model is. -->
23
+
24
+
25
+
26
+ - **Developed by:** [ConvoLite AI]
27
+ - **Funded by:** [VDB]
28
+ - **Model type:** [img2img]
29
+ - **License:** [gpl-3.0]
30
+
31
+
32
+ ## Uses
33
+
34
+ <!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
35
+ For comercial and noncomercial use.
36
+
37
+ ### Direct Use
38
+ For CPU, use the code below:
39
+ ``` python
40
+ import os
41
+ import torch
42
+ import torch.nn as nn
43
+ from PIL import Image
44
+ from torchvision.transforms import ToTensor
45
+ import numpy as np
46
+ from concurrent.futures import ThreadPoolExecutor
47
+
48
+ class DenoisingModel(nn.Module):
49
+ def __init__(self):
50
+ super(DenoisingModel, self).__init__()
51
+ self.enc1 = nn.Sequential(
52
+ nn.Conv2d(3, 64, 3, padding=1),
53
+ nn.ReLU(),
54
+ nn.Conv2d(64, 64, 3, padding=1),
55
+ nn.ReLU()
56
+ )
57
+ self.pool1 = nn.MaxPool2d(2, 2)
58
+
59
+ self.up1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
60
+ self.dec1 = nn.Sequential(
61
+ nn.Conv2d(64, 64, 3, padding=1),
62
+ nn.ReLU(),
63
+ nn.Conv2d(64, 3, 3, padding=1)
64
+ )
65
+
66
+ def forward(self, x):
67
+ e1 = self.enc1(x)
68
+ p1 = self.pool1(e1)
69
+ u1 = self.up1(p1)
70
+ d1 = self.dec1(u1)
71
+ return d1
72
+
73
+ def denoise_patch(model, patch):
74
+ transform = ToTensor()
75
+ input_patch = transform(patch).unsqueeze(0)
76
+
77
+ with torch.no_grad():
78
+ output_patch = model(input_patch)
79
+
80
+ denoised_patch = output_patch.squeeze(0).permute(1, 2, 0).numpy() * 255
81
+ denoised_patch = np.clip(denoised_patch, 0, 255).astype(np.uint8)
82
+
83
+ original_patch = np.array(patch)
84
+ very_bright_mask = original_patch > 240
85
+ bright_mask = (original_patch > 220) & (original_patch <= 240)
86
+
87
+ denoised_patch[very_bright_mask] = original_patch[very_bright_mask]
88
+
89
+ blend_factor = 0.7
90
+ denoised_patch[bright_mask] = (
91
+ blend_factor * original_patch[bright_mask] +
92
+ (1 - blend_factor) * denoised_patch[bright_mask]
93
+ )
94
+
95
+ return denoised_patch
96
+
97
+ def denoise_image(image_path, model_path, patch_size=256, num_threads=4, overlap=32):
98
+ model = DenoisingModel()
99
+ checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
100
+ model.load_state_dict(checkpoint['model_state_dict'])
101
+ model.eval()
102
+
103
+ # Load and get original image dimensions
104
+ image = Image.open(image_path).convert("RGB")
105
+ width, height = image.size
106
+
107
+ # Calculate padding needed
108
+ pad_right = patch_size - (width % patch_size) if width % patch_size != 0 else 0
109
+ pad_bottom = patch_size - (height % patch_size) if height % patch_size != 0 else 0
110
+
111
+ # Add padding with reflection instead of zeros
112
+ padded_width = width + pad_right
113
+ padded_height = height + pad_bottom
114
+
115
+ # Create padded image using reflection padding
116
+ padded_image = Image.new("RGB", (padded_width, padded_height))
117
+ padded_image.paste(image, (0, 0))
118
+
119
+ # Fill right border with reflected content
120
+ if pad_right > 0:
121
+ right_border = image.crop((width - pad_right, 0, width, height))
122
+ padded_image.paste(right_border.transpose(Image.FLIP_LEFT_RIGHT), (width, 0))
123
+
124
+ # Fill bottom border with reflected content
125
+ if pad_bottom > 0:
126
+ bottom_border = image.crop((0, height - pad_bottom, width, height))
127
+ padded_image.paste(bottom_border.transpose(Image.FLIP_TOP_BOTTOM), (0, height))
128
+
129
+ # Fill corner if needed
130
+ if pad_right > 0 and pad_bottom > 0:
131
+ corner = image.crop((width - pad_right, height - pad_bottom, width, height))
132
+ padded_image.paste(corner.transpose(Image.FLIP_LEFT_RIGHT).transpose(Image.FLIP_TOP_BOTTOM),
133
+ (width, height))
134
+
135
+ # Generate patches with positions
136
+ patches = []
137
+ positions = []
138
+ for i in range(0, padded_height, patch_size - overlap):
139
+ for j in range(0, padded_width, patch_size - overlap):
140
+ patch = padded_image.crop((j, i, min(j + patch_size, padded_width), min(i + patch_size, padded_height)))
141
+ patches.append(patch)
142
+ positions.append((i, j))
143
+
144
+ # Process patches in parallel
145
+ with ThreadPoolExecutor(max_workers=num_threads) as executor:
146
+ denoised_patches = list(executor.map(lambda p: denoise_patch(model, p), patches))
147
+
148
+ # Initialize output arrays
149
+ denoised_image = np.zeros((padded_height, padded_width, 3), dtype=np.float32)
150
+ weight_map = np.zeros((padded_height, padded_width), dtype=np.float32)
151
+
152
+ # Create smooth blending weights
153
+ for (i, j), denoised_patch in zip(positions, denoised_patches):
154
+ patch_height, patch_width, _ = denoised_patch.shape
155
+ patch_weights = np.ones((patch_height, patch_width), dtype=np.float32)
156
+ if i > 0:
157
+ patch_weights[:overlap, :] *= np.linspace(0, 1, overlap)[:, np.newaxis]
158
+ if j > 0:
159
+ patch_weights[:, :overlap] *= np.linspace(0, 1, overlap)[np.newaxis, :]
160
+ if i + patch_height < padded_height:
161
+ patch_weights[-overlap:, :] *= np.linspace(1, 0, overlap)[:, np.newaxis]
162
+ if j + patch_width < padded_width:
163
+ patch_weights[:, -overlap:] *= np.linspace(1, 0, overlap)[np.newaxis, :]
164
+
165
+ # Clip the patch values to prevent very bright pixels
166
+ denoised_patch = np.clip(denoised_patch, 0, 255)
167
+
168
+ denoised_image[i:i + patch_height, j:j + patch_width] += (
169
+ denoised_patch * patch_weights[:, :, np.newaxis]
170
+ )
171
+ weight_map[i:i + patch_height, j:j + patch_width] += patch_weights
172
+
173
+ # Normalize by weights
174
+ mask = weight_map > 0
175
+ denoised_image[mask] = denoised_image[mask] / weight_map[mask, np.newaxis]
176
+
177
+ # Crop to original size
178
+ denoised_image = denoised_image[:height, :width]
179
+ denoised_image = np.clip(denoised_image, 0, 255).astype(np.uint8)
180
+
181
+ # Save the result
182
+ denoised_image_path = os.path.splitext(image_path)[0] + "_denoised.png"
183
+ print(f"Saving denoised image to {denoised_image_path}")
184
+
185
+ Image.fromarray(denoised_image).save(denoised_image_path)
186
+
187
+ if __name__ == "__main__":
188
+ image_path = input("Enter the path of the image: ")
189
+ model_path = r"path/to/model.pkl"
190
+ denoise_image(image_path, model_path, num_threads=12)
191
+ print("Denoising completed.") # Use the number of threads your processor has.)
192
+ ```
193
+
194
+
195
+ ### Out-of-Scope Use
196
+
197
+ <!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
198
+
199
+ If the image does not have a high level of noise, it is not recommended to use this model, as it will produce less than ideal results.
200
+
201
+
202
+ ## Training Details
203
+
204
+ This model was trained on a single Nvidia T4 GPU for around 2 hours and 30 minutes.
205
+
206
+ ### Training Data
207
+
208
+ Around 13 GB of publicly available images under the Creative Commons license.
209
+
210
+ #### Speed
211
+
212
+ With an AMD Ryzen 5 5500 it can denoise a 2k image in approx. 2 seconds using multithreading. Still have not tested it out with CUDA, but it's probably faster.
213
+
214
+
215
+
216
+ #### Hardware
217
+
218
+
219
+ | Specifications | Minimum | Recommended |
220
+ |----------|----------|----------|
221
+ | CPU | Intel Core i7-2700K or something else that can run Python | AMD Ryzen 5 5500 |
222
+ | RAM | 4 GB | 16 GB |
223
+ | GPU | not needed | Nvidia GTX 1660 Ti |
224
+
225
+
226
+ #### Software
227
+
228
+ Python
229
+
230
+
231
+
232
+ ## Model Card Authors
233
+
234
+ Vericu de Buget
235
+
236
+
237
+ ## Model Card Contact
238
+
239
+ [convolite@europe.com](mailto:convolite@europe.com)
240
+ [ConvoLite](https://convolite.github.io/selector.html)