Update app.py

app.py

@@ -2,7 +2,7 @@ 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
+import gradio as gr
 from briarmbg import BriaRMBG
 import PIL
 from PIL import Image
@@ -25,7 +25,6 @@ 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
     
@@ -42,13 +41,13 @@ def resize_image(image, max_size=1024):
 
 def process(image, progress=gr.Progress()):
     if image is None:
-        return None, gr.update(visible=False)
+        return None
     
     progress(0, desc="Starting processing...")
     orig_image = Image.fromarray(image)
     original_size = orig_image.size
     
-    progress(0.1, desc="Preparing image...")
+    progress(0.2, desc="Preparing image...")
     process_image = resize_image(orig_image)
     w, h = process_image.size
     
@@ -58,7 +57,7 @@ def process(image, progress=gr.Progress()):
     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...")
+    progress(0.4, desc="Processing with AI model...")
     if torch.cuda.is_available():
         im_tensor = im_tensor.cuda()
     
@@ -80,95 +79,67 @@ def process(image, progress=gr.Progress()):
     new_im = orig_image.copy()
     new_im.putalpha(pil_mask)
     
-    progress(0.8, desc="Preparing download...")
+    progress(0.8, desc="Saving result...")
     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)
+    return gr.Image.update(value=filepath, visible=True)
 
 # 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');
+    @import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@400;500;700&display=swap');
+    
+    .container { max-width: 850px; margin: 0 auto; padding: 20px; }
     
-    :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 {
-        position: relative;
-        width: 100% !important;
-        max-width: 800px !important;
-        margin: 0 auto !important;
-        border-radius: 12px !important;
-        border: 2px solid var(--neon-cyan) !important;
-        background: rgba(18, 18, 56, 0.7) !important;
-        overflow: hidden !important;
-    }
-
     .title-text {
-        color: var(--neon-pink);
+        color: #ff00de;
         font-family: 'Orbitron', sans-serif;
-        font-size: 2em;
+        font-size: 2.5em;
         text-align: center;
         margin: 20px 0;
-        text-shadow: 0 0 10px rgba(255, 0, 222, 0.5);
+        text-shadow: 0 0 10px rgba(255, 0, 222, 0.7);
     }
-
+    
     .subtitle-text {
-        color: var(--neon-cyan);
+        color: #00ffff;
         text-align: center;
-        margin-bottom: 20px;
-        text-shadow: 0 0 10px rgba(0, 255, 255, 0.5);
+        margin-bottom: 30px;
+        font-size: 1.2em;
+        text-shadow: 0 0 8px rgba(0, 255, 255, 0.7);
     }
-
-    .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;
+    
+    .image-container {
+        background: rgba(10, 10, 30, 0.3);
+        border-radius: 15px;
+        padding: 20px;
+        margin: 10px 0;
+        border: 2px solid #00ffff;
+        box-shadow: 0 0 15px rgba(0, 255, 255, 0.2);
     }
-
-    .download-button:hover {
-        transform: translateY(-2px) !important;
-        box-shadow: 0 5px 15px rgba(0, 255, 255, 0.4) !important;
+    
+    .download-btn {
+        background: linear-gradient(45deg, #00ffff, #ff00de);
+        border: none;
+        padding: 12px 25px;
+        border-radius: 8px;
+        color: white;
+        font-family: 'Orbitron', sans-serif;
+        cursor: pointer;
+        transition: all 0.3s ease;
+    }
+    
+    .download-btn:hover {
+        transform: translateY(-2px);
+        box-shadow: 0 5px 15px rgba(0, 255, 255, 0.4);
+    }
+    
+    @media (max-width: 768px) {
+        .title-text { font-size: 1.8em; }
+        .subtitle-text { font-size: 1em; }
+        .image-container { padding: 10px; }
     }
 """) as demo:
     gr.Markdown("""
@@ -176,41 +147,27 @@ with gr.Blocks(css="""
         <p class="subtitle-text">Remove backgrounds instantly using advanced AI technology</p>
     """)
     
-    with gr.Column():
-        input_image = gr.Image(
-            type="numpy",
-            label="Upload Your Image",
-            elem_classes="image-comparison"
-        )
-        
-        output_image = gr.Image(
-            type="numpy",
-            visible=False
-        )
-        
-        image_comparison = gr.Image(
-            type="numpy",
-            label="Before / After Comparison",
-            elem_classes="image-comparison"
-        )
-        
-        download_file = gr.File(
-            label="",
-            file_count="single",
-            interactive=True,
-            visible=False,
-            elem_classes="download-button"
-        )
-
-    input_image.change(
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                label="Upload Image",
+                type="numpy",
+                elem_classes="image-container"
+            )
+            
+            output_image = gr.Image(
+                label="Result",
+                type="filepath",
+                elem_classes="image-container"
+            )
+            
+            process_btn = gr.Button("Remove Background", variant="primary")
+    
+    process_btn.click(
         fn=process,
         inputs=input_image,
-        outputs=[output_image, download_file]
-    ).then(
-        fn=update_comparison,
-        inputs=[input_image, output_image],
-        outputs=image_comparison
+        outputs=output_image
     )
 
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch()