app.py

# 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()