vericudebuget
commited on
Update README.md
Browse files
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)
|