Update app.py

app.py

@@ -1,181 +1,207 @@
-# ... keep all imports and processing functions the same until the interface part ...
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
+import gradio as gr  # Make sure this import is present
+from briarmbg import BriaRMBG
+import PIL
+from PIL import Image
+import tempfile
+import os
+import time
+import uuid
+import shutil
 
-# Update the Gradio interface setup
-title = "Background Removal Tool"
-description = """
-<style>
-    /* ... previous styles ... */
+# 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"""
+    width, height = image.size
+    aspect_ratio = width / height
+    
+    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...")
+    orig_image = Image.fromarray(image)
+    original_size = orig_image.size
+    
+    progress(0.1, desc="Preparing image...")
+    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()
+    
+    with torch.no_grad():
+        result = net(im_tensor)
+    
+    progress(0.6, desc="Post-processing...")
+    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)
+    
+    result_array = (result * 255).cpu().data.numpy().astype(np.uint8)
+    pil_mask = Image.fromarray(np.squeeze(result_array))
+    
+    if pil_mask.size != original_size:
+        pil_mask = pil_mask.resize(original_size, Image.LANCZOS)
+    
+    new_im = orig_image.copy()
+    new_im.putalpha(pil_mask)
+    
+    progress(0.8, desc="Preparing download...")
+    unique_id = str(uuid.uuid4())[:8]
+    filename = f"background_removed_{unique_id}.png"
+    filepath = os.path.join(OUTPUT_DIR, filename)
+    
+    new_im.save(filepath, format='PNG', quality=100)
+    output_array = np.array(new_im.convert("RGBA"))
+    
+    progress(1.0, desc="Done!")
+    return output_array, gr.update(visible=True, value=filepath, interactive=True)
+
+def update_comparison(image, result):
+    if result is None:
+        return None
+    orig_image = Image.fromarray(image)
+    result_image = Image.fromarray(result)
+    
+    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)
+
+# Gradio interface
+with gr.Blocks(css="""
+    @import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@400;500;700&family=Roboto+Mono:wght@300;400;700&display=swap');
+    
+    :root {
+        --neon-cyan: #00ffff;
+        --neon-pink: #ff00de;
+        --dark-background: #0a0a1e;
+    }
+
+    body {
+        font-family: 'Roboto Mono', monospace;
+        background: linear-gradient(135deg, var(--dark-background) 0%, #121238 100%);
+        color: #ffffff;
+    }
+
+    .container {
+        max-width: 800px;
+        margin: 0 auto;
+        padding: 20px;
+    }
 
-    /* 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;
+        overflow: hidden !important;
     }
 
-    .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;
+    .title-text {
+        color: var(--neon-pink);
+        font-family: 'Orbitron', sans-serif;
+        font-size: 2em;
+        text-align: center;
+        margin: 20px 0;
+        text-shadow: 0 0 10px rgba(255, 0, 222, 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%);
+    .subtitle-text {
+        color: var(--neon-cyan);
+        text-align: center;
+        margin-bottom: 20px;
+        text-shadow: 0 0 10px rgba(0, 255, 255, 0.5);
     }
 
-    .slider-handle::before {
-        top: 25%;
-        transform: translateX(-50%) rotate(45deg);
+    .download-button {
+        background: linear-gradient(45deg, var(--neon-cyan), var(--neon-pink)) !important;
+        border: none !important;
+        padding: 10px 20px !important;
+        border-radius: 8px !important;
+        color: white !important;
+        font-family: 'Orbitron', sans-serif !important;
+        cursor: pointer !important;
+        transition: all 0.3s ease !important;
     }
 
-    .slider-handle::after {
-        top: 25%;
-        transform: translateX(-50%) rotate(-45deg);
+    .download-button:hover {
+        transform: translateY(-2px) !important;
+        box-shadow: 0 5px 15px rgba(0, 255, 255, 0.4) !important;
     }
-</style>
-<div class="custom-container">
-    <h1 class="title-text">AI Background Removal</h1>
-    <p class="subtitle-text">
-        Remove backgrounds instantly using advanced AI technology
-    </p>
-</div>
-"""
-
-# Create the Gradio interface
-with gr.Blocks(css="""
-    /* ... previous CSS styles ... */
 """) as demo:
-    gr.Markdown(description)
+    gr.Markdown("""
+        <h1 class="title-text">AI Background Removal</h1>
+        <p class="subtitle-text">Remove backgrounds instantly using advanced AI technology</p>
+    """)
     
-    with gr.Column(scale=1):
+    with gr.Column():
         input_image = gr.Image(
             type="numpy",
             label="Upload Your Image",
-            elem_id="input-image",
-            elem_classes="input-image",
-            container=True
+            elem_classes="image-comparison"
         )
         
-        # 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="Before / After Comparison",
-                elem_id="image-comparison",
-                elem_classes="image-comparison",
-                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;
-            
-            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))
+        image_comparison = gr.Image(
+            type="numpy",
+            label="Before / After Comparison",
+            elem_classes="image-comparison"
+        )
         
-        return np.array(comparison)
+        download_file = gr.File(
+            label="",
+            file_count="single",
+            interactive=True,
+            visible=False,
+            elem_classes="download-button"
+        )
 
-    # Process automatically when image is uploaded
     input_image.change(
         fn=process,
         inputs=input_image,