Update README.md

README.md

@@ -1,3 +1,241 @@
----
-license: gpl-3.0
----
+---
+license: gpl-3.0
+tags:
+- img2img
+- denoiser
+- image
+---
+
+# denoise_medium_v1
+
+denoise_medium_v1 is an image denoiser made for images that have low-light noise.
+
+It performs slightly better than [denoise_small_v1](https://huggingface.co/vericudebuget/denoise_small_v1) on images that have less colorfull noise and can reconstruct a higher level of detail from the original.
+
+For better results, try re-proccesing the same image a few times in a loop (2 times is recommended for really noisy images).
+
+## Model Details
+
+
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [ConvoLite AI]
+- **Funded by:** [VDB]
+- **Model type:** [img2img]
+- **License:** [gpl-3.0]
+
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+For comercial and noncomercial use.
+
+### Direct Use
+For CPU, use the code below:
+``` python
+import os
+import torch
+import torch.nn as nn
+from PIL import Image
+from torchvision.transforms import ToTensor
+import numpy as np
+from concurrent.futures import ThreadPoolExecutor
+
+class DenoisingModel(nn.Module):
+    def __init__(self):
+        super(DenoisingModel, self).__init__()
+        self.enc1 = nn.Sequential(
+            nn.Conv2d(3, 64, 3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 64, 3, padding=1),
+            nn.ReLU()
+        )
+        self.pool1 = nn.MaxPool2d(2, 2)
+        
+        self.up1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.dec1 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 3, 3, padding=1)
+        )
+
+    def forward(self, x):
+        e1 = self.enc1(x)
+        p1 = self.pool1(e1)
+        u1 = self.up1(p1)
+        d1 = self.dec1(u1)
+        return d1
+
+def denoise_patch(model, patch):
+    transform = ToTensor()
+    input_patch = transform(patch).unsqueeze(0)
+    
+    with torch.no_grad():
+        output_patch = model(input_patch)
+    
+    denoised_patch = output_patch.squeeze(0).permute(1, 2, 0).numpy() * 255
+    denoised_patch = np.clip(denoised_patch, 0, 255).astype(np.uint8)
+    
+    original_patch = np.array(patch)
+    very_bright_mask = original_patch > 240
+    bright_mask = (original_patch > 220) & (original_patch <= 240)
+    
+    denoised_patch[very_bright_mask] = original_patch[very_bright_mask]
+    
+    blend_factor = 0.7
+    denoised_patch[bright_mask] = (
+        blend_factor * original_patch[bright_mask] +
+        (1 - blend_factor) * denoised_patch[bright_mask]
+    )
+    
+    return denoised_patch
+
+def denoise_image(image_path, model_path, patch_size=256, num_threads=4, overlap=32):
+    model = DenoisingModel()
+    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+    
+    # Load and get original image dimensions
+    image = Image.open(image_path).convert("RGB")
+    width, height = image.size
+    
+    # Calculate padding needed
+    pad_right = patch_size - (width % patch_size) if width % patch_size != 0 else 0
+    pad_bottom = patch_size - (height % patch_size) if height % patch_size != 0 else 0
+    
+    # Add padding with reflection instead of zeros
+    padded_width = width + pad_right
+    padded_height = height + pad_bottom
+    
+    # Create padded image using reflection padding
+    padded_image = Image.new("RGB", (padded_width, padded_height))
+    padded_image.paste(image, (0, 0))
+    
+    # Fill right border with reflected content
+    if pad_right > 0:
+        right_border = image.crop((width - pad_right, 0, width, height))
+        padded_image.paste(right_border.transpose(Image.FLIP_LEFT_RIGHT), (width, 0))
+    
+    # Fill bottom border with reflected content
+    if pad_bottom > 0:
+        bottom_border = image.crop((0, height - pad_bottom, width, height))
+        padded_image.paste(bottom_border.transpose(Image.FLIP_TOP_BOTTOM), (0, height))
+    
+    # Fill corner if needed
+    if pad_right > 0 and pad_bottom > 0:
+        corner = image.crop((width - pad_right, height - pad_bottom, width, height))
+        padded_image.paste(corner.transpose(Image.FLIP_LEFT_RIGHT).transpose(Image.FLIP_TOP_BOTTOM), 
+                          (width, height))
+    
+    # Generate patches with positions
+    patches = []
+    positions = []
+    for i in range(0, padded_height, patch_size - overlap):
+        for j in range(0, padded_width, patch_size - overlap):
+            patch = padded_image.crop((j, i, min(j + patch_size, padded_width), min(i + patch_size, padded_height)))
+            patches.append(patch)
+            positions.append((i, j))
+    
+    # Process patches in parallel
+    with ThreadPoolExecutor(max_workers=num_threads) as executor:
+        denoised_patches = list(executor.map(lambda p: denoise_patch(model, p), patches))
+    
+    # Initialize output arrays
+    denoised_image = np.zeros((padded_height, padded_width, 3), dtype=np.float32)
+    weight_map = np.zeros((padded_height, padded_width), dtype=np.float32)
+    
+    # Create smooth blending weights
+    for (i, j), denoised_patch in zip(positions, denoised_patches):
+        patch_height, patch_width, _ = denoised_patch.shape
+        patch_weights = np.ones((patch_height, patch_width), dtype=np.float32)
+        if i > 0:
+            patch_weights[:overlap, :] *= np.linspace(0, 1, overlap)[:, np.newaxis]
+        if j > 0:
+            patch_weights[:, :overlap] *= np.linspace(0, 1, overlap)[np.newaxis, :]
+        if i + patch_height < padded_height:
+            patch_weights[-overlap:, :] *= np.linspace(1, 0, overlap)[:, np.newaxis]
+        if j + patch_width < padded_width:
+            patch_weights[:, -overlap:] *= np.linspace(1, 0, overlap)[np.newaxis, :]
+        
+        # Clip the patch values to prevent very bright pixels
+        denoised_patch = np.clip(denoised_patch, 0, 255)
+        
+        denoised_image[i:i + patch_height, j:j + patch_width] += (
+            denoised_patch * patch_weights[:, :, np.newaxis]
+        )
+        weight_map[i:i + patch_height, j:j + patch_width] += patch_weights
+    
+    # Normalize by weights
+    mask = weight_map > 0
+    denoised_image[mask] = denoised_image[mask] / weight_map[mask, np.newaxis]
+    
+    # Crop to original size
+    denoised_image = denoised_image[:height, :width]
+    denoised_image = np.clip(denoised_image, 0, 255).astype(np.uint8)
+    
+    # Save the result
+    denoised_image_path = os.path.splitext(image_path)[0] + "_denoised.png"
+    print(f"Saving denoised image to {denoised_image_path}")
+    
+    Image.fromarray(denoised_image).save(denoised_image_path)
+
+if __name__ == "__main__":
+    image_path = input("Enter the path of the image: ")
+    model_path = r"path/to/model.pkl"
+    denoise_image(image_path, model_path, num_threads=12)
+    print("Denoising completed.")  # Use the number of threads your processor has.)
+```
+
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+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.
+
+
+## Training Details
+
+This model was trained on a single Nvidia T4 GPU for around one hour.
+
+### Training Data
+
+Around 10 GB of publicly available images under the Creative Commons license.
+
+#### Speed
+
+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.
+
+
+
+#### Hardware
+
+
+| Specifications | Minimum | Recommended |
+|----------|----------|----------|
+| CPU | Intel Core i7-2700K or something else that can run Python | AMD Ryzen 5 5500 |
+| RAM | 4 GB | 16 GB |
+| GPU | not needed | Nvidia GTX 1660 Ti |
+
+
+#### Software
+
+Python
+
+
+
+## Model Card Authors 
+
+Vericu de Buget
+
+
+## Model Card Contact
+
+[convolite@europe.com](mailto:convolite@europe.com)
+[ConvoLite](https://convolite.github.io/selector.html)