Update app.py

app.py

@@ -25,6 +25,7 @@ OUTPUT_DIR = "output_images"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 def resize_image(image, max_size=1024):
+    """Resize image while maintaining aspect ratio"""
     width, height = image.size
     aspect_ratio = width / height
     
@@ -40,53 +41,60 @@ def resize_image(image, max_size=1024):
     return image
 
 def process(image, progress=gr.Progress()):
+    """Process the image and remove background"""
     if image is None:
-        return None
-    
-    progress(0, desc="Starting processing...")
-    orig_image = Image.fromarray(image)
-    original_size = orig_image.size
-    
-    progress(0.2, 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.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)
-    
-    progress(1.0, desc="Done!")
-    return gr.Image.update(value=filepath, visible=True)
+        return None, None
+    
+    try:
+        progress(0, desc="Starting processing...")
+        orig_image = Image.fromarray(image)
+        original_size = orig_image.size
+        
+        progress(0.2, desc="Preparing image...")
+        process_image = resize_image(orig_image)
+        w, h = process_image.size
+        
+        # Convert image to tensor
+        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)
+        
+        progress(1.0, desc="Done!")
+        return gr.Image.update(value=filepath, visible=True), gr.File.update(value=filepath, visible=True)
+        
+    except Exception as e:
+        print(f"Error processing image: {str(e)}")
+        return None, None
 
 # Gradio interface
 with gr.Blocks(css="""
@@ -101,6 +109,7 @@ with gr.Blocks(css="""
         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 {
@@ -118,6 +127,12 @@ with gr.Blocks(css="""
         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:hover {
+        box-shadow: 0 0 20px rgba(0, 255, 255, 0.4);
+        transform: translateY(-2px);
     }
     
     .download-btn {
@@ -129,6 +144,10 @@ with gr.Blocks(css="""
         font-family: 'Orbitron', sans-serif;
         cursor: pointer;
         transition: all 0.3s ease;
+        margin-top: 10px;
+        text-align: center;
+        text-transform: uppercase;
+        letter-spacing: 1px;
     }
     
     .download-btn:hover {
@@ -136,10 +155,20 @@ with gr.Blocks(css="""
         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; }
     }
 """) as demo:
     gr.Markdown("""
@@ -158,15 +187,22 @@ with gr.Blocks(css="""
             output_image = gr.Image(
                 label="Result",
                 type="filepath",
-                elem_classes="image-container"
+                elem_classes="image-container",
+                visible=True
             )
             
-            process_btn = gr.Button("Remove Background", variant="primary")
+            download_file = gr.File(
+                label="Download Processed Image",
+                visible=False,
+                elem_classes="download-btn"
+            )
     
-    process_btn.click(
+    # Automatic processing when image is uploaded
+    input_image.change(
         fn=process,
         inputs=input_image,
-        outputs=output_image
+        outputs=[output_image, download_file],
+        api_name="process_image"
     )
 
 if __name__ == "__main__":