Create app.py

app.py

@@ -0,0 +1,261 @@
+# app.py
+
+import os
+import gradio as gr
+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
+from skimage import exposure
+
+# --- Model Definition ---
+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
+
+# --- Denoising Patch Function ---
+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
+
+# --- Main Denoise Image Function (Dynamically uses all CPU cores) ---
+def denoise_image(image: Image.Image, model_path: str, patch_size: int = 256, overlap: int = 32) -> Image.Image:
+    # Dynamically set the number of threads based on available CPU cores
+    num_threads = os.cpu_count()
+    if num_threads is None: # Fallback in case os.cpu_count() returns None
+        num_threads = 2 # Default to 2 if cannot detect
+    print(f"Utilizing {num_threads} CPU cores for parallel processing.")
+
+    # Load the model
+    model = DenoisingModel()
+    # Ensure model is loaded on CPU, crucial for Hugging Face Spaces free tier
+    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval() # Set model to evaluation mode
+    
+    # Process image (convert to RGB, get dimensions)
+    image = image.convert("RGB")
+    width, height = image.size
+    
+    # Calculate padding needed for full patches
+    pad_right = (patch_size - (width % patch_size)) % patch_size if width % patch_size != 0 else 0
+    pad_bottom = (patch_size - (height % patch_size)) % patch_size if height % patch_size != 0 else 0
+    
+    padded_width = width + pad_right
+    padded_height = height + pad_bottom
+    
+    # Create padded image with reflection padding
+    padded_image = Image.new("RGB", (padded_width, padded_height))
+    padded_image.paste(image, (0, 0)) # Paste original image
+    
+    # Fill borders 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))
+    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))
+    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 and their positions for processing
+    patches = []
+    positions = []
+    # Adjust ranges to ensure full patch sizes near borders using min()
+    for i in range(0, padded_height, patch_size - overlap):
+        for j in range(0, padded_width, patch_size - overlap):
+            # Ensure patch doesn't go out of bounds for the last patches
+            actual_i = min(i, padded_height - patch_size)
+            actual_j = min(j, padded_width - patch_size)
+            patch = padded_image.crop((actual_j, actual_i, actual_j + patch_size, actual_i + patch_size))
+            patches.append(patch)
+            positions.append((actual_i, actual_j))
+            
+    # Process patches in parallel using ThreadPoolExecutor
+    with ThreadPoolExecutor(max_workers=num_threads) as executor:
+        denoised_patches = list(executor.map(lambda p: denoise_patch(model, p), patches))
+        
+    # Reconstruct the image from denoised patches using blending
+    denoised_image_np = np.zeros((padded_height, padded_width, 3), dtype=np.float32)
+    weight_map = np.zeros((padded_height, padded_width), dtype=np.float32)
+    
+    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)
+        
+        # Apply smooth blending weights at overlaps
+        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: # Check if it's not the last row
+            patch_weights[-overlap:, :] *= np.linspace(1, 0, overlap)[:, np.newaxis]
+        if j + patch_width < padded_width: # Check if it's not the last column
+            patch_weights[:, -overlap:] *= np.linspace(1, 0, overlap)[np.newaxis, :]
+        
+        # Clip values and apply gamma correction before blending
+        denoised_patch_processed = exposure.adjust_gamma(np.clip(denoised_patch, 0, 255), gamma=1.0)
+        
+        denoised_image_np[i:i + patch_height, j:j + patch_width] += (
+            denoised_patch_processed * patch_weights[:, :, np.newaxis]
+        )
+        weight_map[i:i + patch_height, j:j + patch_width] += patch_weights
+    
+    # Normalize by weights to get the final blended image
+    mask = weight_map > 0
+    denoised_image_np[mask] /= weight_map[mask, np.newaxis]
+    
+    # Crop back to original dimensions and finalize
+    final_image_np = denoised_image_np[:height, :width]
+    final_image_np = np.clip(final_image_np, 0, 255).astype(np.uint8)
+    
+    return Image.fromarray(final_image_np)
+
+# --- Gradio Interface Setup ---
+
+# Function to find available models in the 'models' directory
+def get_available_models():
+    model_dir = "models"
+    if not os.path.exists(model_dir):
+        print(f"Warning: '{model_dir}' directory not found. No models will be available.")
+        return []
+    
+    # Filter for .pth or .pt files
+    models = [f for f in os.listdir(model_dir) if f.endswith(".pth") or f.endswith(".pt")]
+    if not models:
+        print(f"Warning: No .pth or .pt model files found in '{model_dir}'.")
+    return models
+
+# The main Gradio function that orchestrates the denoising process
+def gradio_interface(input_image: np.ndarray, model_name: str, progress=gr.Progress(track_tqdm=True)) -> Image.Image:
+    if input_image is None:
+        raise gr.Error("Please upload an image to denoise.")
+    if not model_name:
+        raise gr.Error("Please select a model from the dropdown.")
+        
+    # Convert numpy input image from Gradio to PIL Image for processing
+    pil_image = Image.fromarray(input_image)
+    
+    # Construct full model path
+    model_path = os.path.join("models", model_name)
+    
+    print(f"Starting denoising process with model: '{model_name}'")
+    progress(0, desc=f"Loading model: {model_name}...")
+    
+    # Call the core denoising function
+    denoised_pil_image = denoise_image(pil_image, model_path)
+    
+    print("Denoising completed successfully.")
+    progress(1, desc="Done!")
+    
+    return denoised_pil_image
+
+# Get initial list of models for the dropdown
+available_models = get_available_models()
+
+# Define the Gradio interface using gr.Blocks for a structured layout
+with gr.Blocks(theme=gr.themes.Soft(), title="Image Denoiser") as demo:
+    gr.Markdown(
+        """
+        # 🖼️ Universal Image Denoiser
+        Upload an image and select a pre-trained model to effectively remove noise.
+        """
+    )
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_img = gr.Image(type="numpy", label="Input Image", value=None) # Set initial value to None
+            
+            # Dropdown for model selection
+            model_dropdown = gr.Dropdown(
+                choices=available_models,
+                label="Select Denoising Model",
+                value=available_models[0] if available_models else None, # Pre-select first model if available
+                info="Place your .pth or .pt model files in the 'models/' directory."
+            )
+            
+            denoise_button = gr.Button("Denoise Image", variant="primary")
+            gr.Markdown(
+                """
+                **Note:** Processing large images can take time. The app utilizes all available CPU cores for faster denoising.
+                """
+            )
+        with gr.Column(scale=1):
+            output_img = gr.Image(type="pil", label="Denoised Image")
+    
+    # Examples section to demonstrate usage
+    if available_models and os.path.exists("examples"):
+        example_images = []
+        # Find some example images (e.g., .png, .jpg)
+        for fname in os.listdir("examples"):
+            if fname.lower().endswith(('.png', '.jpg', '.jpeg')):
+                example_images.append(os.path.join("examples", fname))
+        
+        # Create examples only if there are models AND example images
+        if example_images:
+            gr.Examples(
+                examples=[[img_path, available_models[0]] for img_path in example_images],
+                inputs=[input_img, model_dropdown],
+                outputs=output_img,
+                fn=gradio_interface,
+                cache_examples=True # Speeds up example loading
+            )
+        else:
+            gr.Markdown("*(No example images found in 'examples/' directory)*")
+    else:
+        gr.Markdown("*(No models or example images found to populate examples)*")
+
+
+    # Connect the button click to the denoising function
+    denoise_button.click(
+        fn=gradio_interface,
+        inputs=[input_img, model_dropdown],
+        outputs=output_img
+    )
+
+if __name__ == "__main__":
+    demo.launch()