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