Update app.py

app.py

@@ -24,13 +24,77 @@ print(f"Model loaded on {device}")
 OUTPUT_DIR = "output_images"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# ... keep all imports and model loading code the same ...
-
 # Define different sizes for different devices
 DESKTOP_SIZE = (512, 512)
-MOBILE_SIZE = (300, 300)  # Smaller size for mobile devices
+MOBILE_SIZE = (300, 300)
+
+# Resize the input image for model compatibility
+def resize_image(image, size=DESKTOP_SIZE):
+    image = image.convert('RGB')
+    image = image.resize(size, Image.LANCZOS)
+    return image
 
-# Update the description and CSS part with responsive design:
+# Background removal process
+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)
+    # Resize input image to fixed size
+    orig_image = resize_image(orig_image, DESKTOP_SIZE)
+    w, h = DESKTOP_SIZE
+    image = orig_image
+    im_np = np.array(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))
+    
+    # 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
+    new_im.save(filepath, format='PNG')
+    
+    # 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
+title = "Background Removal Tool"
 description = """
 <style>
     .custom-container {
@@ -81,7 +145,7 @@ description = """
 </div>
 """
 
-# Update the Gradio interface with responsive CSS
+# 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');
@@ -184,6 +248,12 @@ with gr.Blocks(css="""
             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); }
+    }
 """) as demo:
     gr.Markdown(description)