Update app.py

app.py

@@ -1,155 +1,65 @@
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torchvision.transforms.functional import normalize
-import gradio as gr
-from briarmbg import BriaRMBG
-import PIL
-from PIL import Image
-import tempfile
-import os
-import time
-import uuid
-import shutil
+# ... keep all imports and processing functions the same until the interface part ...
 
-# Load the pre-trained model
-print("Loading model...")
-net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-net.to(device)
-net.eval()
-print(f"Model loaded on {device}")
-
-# Create output directory if it doesn't exist
-OUTPUT_DIR = "output_images"
-os.makedirs(OUTPUT_DIR, exist_ok=True)
-
-def resize_image(image, max_size=1024):
-    """Resize image while maintaining aspect ratio and quality"""
-    # Get original size
-    width, height = image.size
-    
-    # Calculate aspect ratio
-    aspect_ratio = width / height
-    
-    # Only resize if the image is larger than max_size in either dimension
-    if width > max_size or height > max_size:
-        if width > height:
-            new_width = max_size
-            new_height = int(max_size / aspect_ratio)
-        else:
-            new_height = max_size
-            new_width = int(max_size * aspect_ratio)
-        image = image.resize((new_width, new_height), Image.LANCZOS)
-    
-    return image
-
-def process(image, progress=gr.Progress()):
-    if image is None:
-        return None, gr.update(visible=False)
-    
-    progress(0, desc="Starting processing...")
-    
-    # Prepare the input
-    progress(0.1, desc="Preparing image...")
-    orig_image = Image.fromarray(image)
-    original_size = orig_image.size
-    
-    # Resize only if needed for processing
-    process_image = resize_image(orig_image)
-    w, h = process_image.size
-    
-    im_np = np.array(process_image)
-    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
-    im_tensor = torch.unsqueeze(im_tensor, 0)
-    im_tensor = torch.divide(im_tensor, 255.0)
-    im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
-    
-    progress(0.3, desc="Processing with AI model...")
-    if torch.cuda.is_available():
-        im_tensor = im_tensor.cuda()
-    
-    # Inference with the model
-    with torch.no_grad():
-        result = net(im_tensor)
-    
-    progress(0.6, desc="Post-processing...")
-    # Post-process the result
-    result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
-    ma = torch.max(result)
-    mi = torch.min(result)
-    result = (result - mi) / (ma - mi)
-    
-    # Convert the result to an image
-    result_array = (result * 255).cpu().data.numpy().astype(np.uint8)
-    pil_mask = Image.fromarray(np.squeeze(result_array))
-    
-    # Resize mask back to original size if needed
-    if pil_mask.size != original_size:
-        pil_mask = pil_mask.resize(original_size, Image.LANCZOS)
-    
-    # Add the mask as alpha channel to the original image
-    new_im = orig_image.copy()
-    new_im.putalpha(pil_mask)
-    
-    progress(0.8, desc="Preparing download...")
-    # Generate a unique filename
-    unique_id = str(uuid.uuid4())[:8]
-    filename = f"background_removed_{unique_id}.png"
-    filepath = os.path.join(OUTPUT_DIR, filename)
-    
-    # Save the processed image in original resolution
-    new_im.save(filepath, format='PNG', quality=100)
-    
-    # Convert to numpy array for display
-    output_array = np.array(new_im.convert("RGBA"))
-    
-    progress(1.0, desc="Done!")
-    
-    return output_array, gr.update(visible=True, value=filepath, interactive=True)
-
-# Gradio interface setup
+# Update the Gradio interface setup
 title = "Background Removal Tool"
 description = """
 <style>
-    .custom-container {
-        text-align: center;
-        max-width: 100%;
-        margin: 0 auto;
-        padding: 20px;
-        background: rgba(10, 10, 30, 0.6);
-        border-radius: 15px;
-        box-shadow: 0 0 20px rgba(0, 255, 255, 0.2);
+    /* ... previous styles ... */
+
+    /* Image comparison slider styles */
+    .image-comparison {
+        position: relative;
+        width: 100% !important;
+        max-width: 800px !important;
+        margin: 0 auto !important;
+        overflow: hidden !important;
+        border-radius: 12px !important;
+        border: 2px solid var(--neon-cyan) !important;
+        background: rgba(18, 18, 56, 0.7) !important;
     }
-    .title-text {
-        color: #ff00de;
-        font-family: 'Orbitron', sans-serif;
-        font-size: 2em;
-        margin: 20px 0;
-        text-shadow: 0 0 10px rgba(255, 0, 222, 0.5);
-        animation: title-pulse 2s infinite alternate;
+
+    .comparison-slider {
+        position: absolute !important;
+        width: 4px !important;
+        height: 100% !important;
+        background: var(--neon-cyan) !important;
+        box-shadow: 0 0 10px rgba(0, 255, 255, 0.5) !important;
+        z-index: 10 !important;
+        cursor: ew-resize !important;
+    }
+
+    .slider-handle {
+        position: absolute !important;
+        width: 40px !important;
+        height: 40px !important;
+        background: var(--neon-cyan) !important;
+        border-radius: 50% !important;
+        top: 50% !important;
+        left: 50% !important;
+        transform: translate(-50%, -50%) !important;
+        box-shadow: 0 0 15px rgba(0, 255, 255, 0.8) !important;
+        cursor: ew-resize !important;
     }
-    .subtitle-text {
-        color: #00ffff;
-        font-family: 'Roboto Mono', monospace;
-        font-size: 1em;
-        margin-top: 10px;
-        line-height: 1.5;
-        text-shadow: 0 0 10px rgba(0, 255, 255, 0.5);
+
+    .slider-handle::before,
+    .slider-handle::after {
+        content: '';
+        position: absolute;
+        width: 2px;
+        height: 50%;
+        background: rgba(0, 0, 0, 0.8);
+        left: 50%;
+        transform: translateX(-50%);
     }
-    @keyframes title-pulse {
-        0% { text-shadow: 0 0 5px rgba(255, 0, 222, 0.5); }
-        100% { text-shadow: 0 0 15px rgba(255, 0, 222, 0.8), 0 0 25px rgba(255, 0, 222, 0.5); }
+
+    .slider-handle::before {
+        top: 25%;
+        transform: translateX(-50%) rotate(45deg);
     }
 
-    /* Responsive text sizes */
-    @media (max-width: 768px) {
-        .title-text {
-            font-size: 1.5em;
-        }
-        .subtitle-text {
-            font-size: 0.9em;
-        }
+    .slider-handle::after {
+        top: 25%;
+        transform: translateX(-50%) rotate(-45deg);
     }
 </style>
 <div class="custom-container">
@@ -162,154 +72,118 @@ description = """
 
 # Create the Gradio interface
 with gr.Blocks(css="""
-    /* Import fonts */
-    @import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@400;500;700&family=Roboto+Mono:wght@300;400;700&display=swap');
-    
-    /* Variables */
-    :root {
-        --neon-cyan: #00ffff;
-        --neon-pink: #ff00de;
-        --neon-yellow: #ffdd00;
-        --dark-background: #0a0a1e;
-        --deep-blue: #121238;
-    }
-
-    /* Global styles */
-    body {
-        font-family: 'Roboto Mono', monospace;
-        background: linear-gradient(135deg, var(--dark-background) 0%, var(--deep-blue) 100%);
-        color: #ffffff;
-        min-height: 100vh;
-    }
-
-    /* Responsive container */
-    .container {
-        width: 100%;
-        max-width: 1200px;
-        margin: 0 auto;
-        padding: 10px;
-    }
-
-    /* Input/Output areas with responsive sizing */
-    .input-image, .output-image {
-        width: 100% !important;
-        max-width: 800px !important;
-        height: auto !important;
-        min-height: 300px !important;
-        object-fit: contain !important;
-        background: rgba(18, 18, 56, 0.7) !important;
-        border: 2px solid var(--neon-cyan) !important;
-        border-radius: 12px !important;
-        transition: all 0.3s ease !important;
-        overflow: hidden !important;
-        margin: 0 auto !important;
-    }
-
-    .input-image img, .output-image img {
-        max-width: 100% !important;
-        max-height: 800px !important;
-        object-fit: contain !important;
-        margin: auto !important;
-    }
-
-    /* Responsive columns */
-    .contain-center {
-        display: flex;
-        flex-direction: column;
-        align-items: center;
-        gap: 20px;
-    }
-
-    /* Download button styling */
-    .download-container [data-testid="file"] button {
-        background: linear-gradient(45deg, var(--neon-cyan), var(--neon-pink)) !important;
-        color: white !important;
-        border: none !important;
-        padding: 12px 28px !important;
-        font-family: 'Orbitron', sans-serif !important;
-        font-size: 16px !important;
-        font-weight: 600 !important;
-        text-transform: uppercase !important;
-        letter-spacing: 1px !important;
-        border-radius: 8px !important;
-        cursor: pointer !important;
-        transition: all 0.3s ease !important;
-        animation: button-glow 2s infinite alternate !important;
-        width: 100% !important;
-        max-width: 300px !important;
-    }
-
-    /* Labels */
-    label {
-        color: var(--neon-cyan) !important;
-        font-family: 'Orbitron', sans-serif !important;
-        font-size: 1.1em !important;
-        text-shadow: 0 0 5px rgba(0, 255, 255, 0.5) !important;
-        margin-bottom: 8px !important;
-        text-align: center !important;
-    }
-
-    /* Responsive layout */
-    @media (max-width: 768px) {
-        .input-image, .output-image {
-            min-height: 200px !important;
-        }
-        
-        .input-image img, .output-image img {
-            max-height: 500px !important;
-        }
-        
-        label {
-            font-size: 0.9em !important;
-        }
-        
-        .download-container [data-testid="file"] button {
-            padding: 10px 20px !important;
-            font-size: 14px !important;
-        }
-    }
-
-    /* Additional Animations */
-    @keyframes button-glow {
-        0% { box-shadow: 0 0 5px rgba(0, 255, 255, 0.5); }
-        100% { box-shadow: 0 0 15px rgba(0, 255, 255, 0.8), 0 0 25px rgba(255, 0, 222, 0.5); }
-    }
+    /* ... previous CSS styles ... */
 """) as demo:
     gr.Markdown(description)
     
-    with gr.Row(equal_height=True):
-        with gr.Column(scale=1):
-            input_image = gr.Image(
+    with gr.Column(scale=1):
+        input_image = gr.Image(
+            type="numpy",
+            label="Upload Your Image",
+            elem_id="input-image",
+            elem_classes="input-image",
+            container=True
+        )
+        
+        # Hidden output image for processing
+        output_image = gr.Image(
+            type="numpy",
+            visible=False
+        )
+        
+        # Image comparison component
+        with gr.Row(elem_classes="image-comparison-container"):
+            image_comparison = gr.Image(
                 type="numpy",
-                label="Upload Your Image",
-                elem_id="input-image",
-                elem_classes="input-image",
+                label="Before / After Comparison",
+                elem_id="image-comparison",
+                elem_classes="image-comparison",
                 container=True
             )
         
-        with gr.Column(scale=1):
-            output_image = gr.Image(
-                type="numpy",
-                label="Result",
-                elem_id="output-image",
-                elem_classes="output-image",
-                container=True
+        with gr.Row(elem_classes="download-container"):
+            download_file = gr.File(
+                label="",
+                file_count="single",
+                interactive=True,
+                visible=False
             )
+
+    # Custom JavaScript for image comparison slider
+    demo.load(js="""
+        function initComparison() {
+            const container = document.querySelector('.image-comparison');
+            if (!container) return;
             
-            with gr.Row(elem_classes="download-container"):
-                download_file = gr.File(
-                    label="",
-                    file_count="single",
-                    interactive=True,
-                    visible=False
-                )
-    
+            const slider = document.createElement('div');
+            slider.className = 'comparison-slider';
+            const handle = document.createElement('div');
+            handle.className = 'slider-handle';
+            slider.appendChild(handle);
+            container.appendChild(slider);
+            
+            let isDown = false;
+            let startX;
+            let sliderLeft;
+            
+            slider.addEventListener('mousedown', (e) => {
+                isDown = true;
+                startX = e.pageX - slider.offsetLeft;
+            });
+            
+            document.addEventListener('mouseup', () => {
+                isDown = false;
+            });
+            
+            document.addEventListener('mousemove', (e) => {
+                if (!isDown) return;
+                e.preventDefault();
+                
+                const x = e.pageX - container.offsetLeft;
+                const walk = x - startX;
+                
+                const containerWidth = container.offsetWidth;
+                let newLeft = (x / containerWidth) * 100;
+                newLeft = Math.max(0, Math.min(100, newLeft));
+                
+                slider.style.left = `${newLeft}%`;
+                container.style.setProperty('--slider-position', `${newLeft}%`);
+            });
+        }
+        
+        // Initialize comparison slider when images are loaded
+        document.addEventListener('DOMContentLoaded', initComparison);
+        // Reinitialize when new images are loaded
+        const observer = new MutationObserver(initComparison);
+        observer.observe(document.body, { childList: true, subtree: true });
+    """)
+
+    def update_comparison(image, result):
+        if result is None:
+            return None
+        # Create side-by-side comparison
+        orig_image = Image.fromarray(image)
+        result_image = Image.fromarray(result)
+        
+        # Ensure both images are the same size
+        width = max(orig_image.width, result_image.width)
+        height = max(orig_image.height, result_image.height)
+        
+        comparison = Image.new('RGBA', (width * 2, height))
+        comparison.paste(orig_image, (0, 0))
+        comparison.paste(result_image, (width, 0))
+        
+        return np.array(comparison)
+
     # Process automatically when image is uploaded
     input_image.change(
         fn=process,
         inputs=input_image,
-        outputs=[output_image, download_file],
-        show_progress="full"
+        outputs=[output_image, download_file]
+    ).then(
+        fn=update_comparison,
+        inputs=[input_image, output_image],
+        outputs=image_comparison
     )
 
 if __name__ == "__main__":