Add files

app.py

@@ -0,0 +1,259 @@
+#!/usr/bin/env python
+"""Gradio app for ClearView image deraining model.
+
+A user-friendly interface for removing rain from images using deep learning.
+"""
+
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from PIL import Image
+import numpy as np
+import time
+from huggingface_hub import hf_hub_download
+
+# Import your model (adjust paths as needed)
+from clearview.models import UNet
+from clearview.utils.image import numpy_to_tensor, tensor_to_numpy
+
+
+class DerainModel:
+    """Wrapper for the deraining model with preprocessing/postprocessing."""
+    
+    def __init__(self, checkpoint_path: str, device: str = "cuda"):
+        """Initialize the model.
+        
+        Args:
+            checkpoint_path: Path to model checkpoint (.pth file)
+            device: Device to run inference on ('cuda' or 'cpu')
+        """
+        self.device = device if torch.cuda.is_available() else "cpu"
+        
+        # Load model
+        self.model = UNet(in_channels=3, out_channels=3)
+        
+        # Load checkpoint
+        checkpoint = torch.load(checkpoint_path, map_location=self.device)
+        if 'model_state_dict' in checkpoint:
+            self.model.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            self.model.load_state_dict(checkpoint)
+        
+        self.model = self.model.to(self.device)
+        self.model.eval()
+        
+        print(f"Model loaded on {self.device}")
+    
+    def preprocess(self, image: Image.Image) -> torch.Tensor:
+        """Preprocess PIL image to tensor.
+        
+        Args:
+            image: PIL Image
+            
+        Returns:
+            Preprocessed tensor (1, 3, H, W)
+        """
+        # Convert to numpy array
+        img_np = np.array(image).astype(np.float32) / 255.0
+        
+        # Convert to tensor (C, H, W)
+        img_tensor = numpy_to_tensor(img_np, normalize=False)
+        
+        # Add batch dimension (1, C, H, W)
+        img_tensor = img_tensor.unsqueeze(0).to(self.device)
+        
+        return img_tensor
+    
+    def postprocess(self, tensor: torch.Tensor) -> Image.Image:
+        """Postprocess tensor to PIL image.
+        
+        Args:
+            tensor: Model output (1, 3, H, W)
+            
+        Returns:
+            PIL Image
+        """
+        # Remove batch dimension and move to CPU
+        tensor = tensor.squeeze(0).cpu()
+        
+        # Clamp to [0, 1]
+        tensor = torch.clamp(tensor, 0, 1)
+        
+        # Convert to numpy (H, W, C)
+        img_np = tensor_to_numpy(tensor)
+        
+        # Convert to uint8
+        img_np = (img_np * 255).astype(np.uint8)
+        
+        # Convert to PIL
+        return Image.fromarray(img_np)
+    
+    @torch.no_grad()
+    def derain(self, image: Image.Image) -> tuple[Image.Image, float]:
+        """Remove rain from image.
+        
+        Args:
+            image: Input PIL Image with rain
+            
+        Returns:
+            Tuple of (derained PIL Image, inference time in seconds)
+        """
+        start_time = time.time()
+        
+        # Store original size
+        original_size = image.size  # (width, height)
+        
+        # Preprocess
+        input_tensor = self.preprocess(image)
+        
+        # Inference
+        output_tensor = self.model(input_tensor)
+        
+        # Postprocess
+        output_image = self.postprocess(output_tensor)
+        
+        # Ensure output matches input size exactly
+        if output_image.size != original_size:
+            output_image = output_image.resize(original_size, Image.LANCZOS)
+        
+        inference_time = time.time() - start_time
+        
+        return output_image, inference_time
+
+
+# Initialize model (will be loaded when app starts)
+MODEL = None
+CHECKPOINT_PATH = hf_hub_download(
+    repo_id="dronefreak/clearview-derain-unet",  # Your model repo
+    filename="clearview-derain-unet.pth"
+)
+
+
+def load_model():
+    """Load model on startup."""
+    global MODEL
+    if MODEL is None:
+        MODEL = DerainModel(CHECKPOINT_PATH)
+    return MODEL
+
+
+def derain_interface(image: Image.Image) -> tuple[tuple[Image.Image, Image.Image], str]:
+    """Gradio interface function.
+    
+    Args:
+        image: Input PIL Image
+        
+    Returns:
+        Tuple of ((input_image, output_image) for slider, info text)
+    """
+    if image is None:
+        return None, "Please upload an image first."
+    
+    # Load model if not already loaded
+    model = load_model()
+
+    image = image.convert("RGB")
+    
+    # Run inference
+    output_image, inference_time = model.derain(image)
+
+    output_image = output_image.convert("RGB")
+    output_image = output_image.resize(image.size, Image.NEAREST)
+    
+    # Create info text
+    info = f"✅ Rain removed successfully!\n"
+    info += f"⚡ Inference time: {inference_time:.3f}s\n"
+    info += f"📐 Image size: {image.size[0]}x{image.size[1]}"
+    
+    # Return (input, output) tuple for ImageSlider
+    return (image, output_image), info
+
+
+# Create Gradio interface
+with gr.Blocks(title="ClearView: Image Deraining") as demo:
+    gr.Markdown("""
+    # 🌧️ ClearView: Image Deraining
+    
+    Remove rain streaks from images using deep learning. Upload a rainy image and see it transform!
+    
+    **Model:** UNet with L1 loss, trained on Rain1400 dataset (12,600 images)
+    **Performance:** 30.9 PSNR / 0.914 SSIM on test set
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                type="pil",
+                label="📤 Upload Rainy Image",
+                # height=400
+            )
+            
+            derain_button = gr.Button("✨ Remove Rain", variant="primary", size="lg")
+            
+            info_text = gr.Textbox(
+                label="ℹ️ Info",
+                lines=3,
+                interactive=False
+            )
+        
+        with gr.Column():
+            output_slider = gr.ImageSlider(
+                label="🔄 Before & After (Drag slider to compare)",
+                # height=400,
+                interactive=False,  # Don't allow uploading to the slider
+                show_download_button=True
+            )
+    
+    # Examples
+    gr.Markdown("### 📸 Try these examples:")
+    
+    gr.Examples(
+        examples=[
+            # Add paths to your example images here
+            ["examples/heavy_rain.jpg"],
+            ["examples/light_rain.jpg"],
+        ],
+        inputs=input_image,
+        outputs=[output_slider, info_text],
+        fn=derain_interface,
+        cache_examples=False,
+    )
+    
+    gr.Markdown("""
+    ---
+    ### 📝 Notes:
+    - Works best on synthetic rain patterns (trained on Rain1400)
+    - May smooth fine textures slightly
+    - For best results, use images with clear rain streaks
+    - Real-world rain may have mixed results
+    
+    ### 🔗 Links:
+    - [GitHub Repository](https://github.com/dronefreak/clearview)
+    - [Model Card](https://huggingface.co/YOUR_USERNAME/clearview-unet)
+    - [Paper/Documentation](#)
+    
+    ### ⚠️ Limitations:
+    - Trained on synthetic data (may not generalize to all real-world scenarios)
+    - No temporal consistency for video (processes frames independently)
+    - May lose fine details in heavily textured areas
+    """)
+    
+    # Connect button to function
+    derain_button.click(
+        fn=derain_interface,
+        inputs=input_image,
+        outputs=[output_slider, info_text]
+    )
+
+
+if __name__ == "__main__":
+    print("Starting ClearView Gradio app...")
+    print(f"Model checkpoint: {CHECKPOINT_PATH}")
+    print(f"CUDA available: {torch.cuda.is_available()}")
+    
+    # Launch app
+    demo.launch(
+        share=False,  # Set to True to create public link
+        server_name="0.0.0.0",  # Allow external connections
+        server_port=7860,
+    )

requirements.txt

@@ -0,0 +1,6 @@
+gradio>=4.0.0
+torch>=2.0.0
+torchvision>=0.15.0
+Pillow>=9.0.0
+numpy>=1.21.0
+huggingface-hub>=0.19.0