Update app.py

app.py

@@ -26,7 +26,7 @@ os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 def process(image, progress=gr.Progress()):
     if image is None:
-        return None, None
+        return None, None, None
     try:
         progress(0, desc="Starting processing...")
         orig_image = Image.fromarray(image)
@@ -73,6 +73,188 @@ def process(image, progress=gr.Progress()):
         # Convert to numpy array for display
         output_array = np.array(new_im)
         
+        progress(1.0, desc="Done!")
+        return (
+            output_array, 
+            filepath, 
+            gr.update(value=f"""
+                <script>
+                    setTimeout(function() {{
+                        const link = document.createElement('a');
+                        link.href = '/file={filepath}';
+                        link.download = '{filename}';
+                        document.body.appendChild(link);
+                        link.click();
+                        document.body.removeChild(link);
+                    }}, 1000);
+                </script>
+            """, visible=True)
+        )
+    
+    except Exception as e:
+        print(f"Error processing image: {str(e)}")
+        return None, None, None
+
+css = """
+@import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@400;500;700&display=swap');
+
+.container { max-width: 850px; margin: 0 auto; padding: 20px; }
+
+.title-text {
+    color: #ff00de;
+    font-family: 'Orbitron', sans-serif;
+    font-size: 2.5em;
+    text-align: center;
+    margin: 20px 0;
+    text-shadow: 0 0 10px rgba(255, 0, 222, 0.7);
+    animation: glow 2s ease-in-out infinite alternate;
+}
+
+.subtitle-text {
+    color: #00ffff;
+    text-align: center;
+    margin-bottom: 30px;
+    font-size: 1.2em;
+    text-shadow: 0 0 8px rgba(0, 255, 255, 0.7);
+}
+
+.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);
+    transition: all 0.3s ease;
+}
+
+.image-container img {
+    max-width: 100%;
+    height: auto;
+    display: block;
+    margin: 0 auto;
+}
+
+.image-container:hover {
+    box-shadow: 0 0 20px rgba(0, 255, 255, 0.4);
+    transform: translateY(-2px);
+}
+
+.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;
+    margin-top: 10px;
+    text-align: center;
+    text-transform: uppercase;
+    letter-spacing: 1px;
+    display: block;
+    width: 100%;
+}
+
+.download-btn:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 5px 15px rgba(0, 255, 255, 0.4);
+}
+
+@keyframes glow {
+    from {
+        text-shadow: 0 0 5px #ff00de, 0 0 10px #ff00de, 0 0 15px #ff00de;
+    }
+    to {
+        text-shadow: 0 0 10px #ff00de, 0 0 20px #ff00de, 0 0 30px #ff00de;
+    }
+}
+
+@media (max-width: 768px) {
+    .title-text { font-size: 1.8em; }
+    .subtitle-text { font-size: 1em; }
+    .image-container { padding: 10px; }
+    .download-btn { padding: 10px 20px; }
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("""
+        <h1 class="title-text">AI Background Removal</h1>
+        <p class="subtitle-text">Remove backgrounds instantly using advanced AI technology</p>
+    """)
+    
+    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="numpy",
+                show_label=True,
+                elem_classes="image-container"
+            )
+            
+            download_button = gr.File(
+                label="Download Result",
+                visible=True,
+                elem_classes="download-btn"
+            )
+            
+            # Add HTML component for auto-download
+            auto_download = gr.HTML(visible=False)
+    
+    input_image.change(
+        fn=process,
+        inputs=input_image,
+        outputs=[output_image, download_button, auto_download]
+    )
+
+if __name__ == "__main__":
+    demo.launch()        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.4, 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="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)
+        
+        # Convert to numpy array for display
+        output_array = np.array(new_im)
+        
         progress(1.0, desc="Done!")
         return output_array, filepath