Update app.py

app.py

@@ -2,102 +2,93 @@ import gradio as gr
 from rembg import remove
 from PIL import Image
 import numpy as np
+import torch
 import cv2
-from skimage import filters
-import time
+import os
 
-def enhance_mask(mask):
-    """Refine the mask edges for better quality"""
-    if len(mask.shape) == 3:
-        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+# Initialize InSPyReNet if available
+try:
+    from InSPyReNet.models.InSPyReNet import InSPyReNet
+    from InSPyReNet.utils.dataloader import test_dataset
     
-    mask = cv2.GaussianBlur(mask, (5, 5), 0)
-    _, binary_mask = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-    edges = filters.sobel(binary_mask)
-    refined_mask = np.where(edges > 0.1, 255, binary_mask)
-    return refined_mask.astype(np.uint8)
-
-def resize_image(img, max_size=1024):
-    """Resize large images while maintaining aspect ratio"""
-    width, height = img.size
-    if max(width, height) > max_size:
-        ratio = max_size / max(width, height)
-        new_size = (int(width * ratio), int(height * ratio))
-        img = img.resize(new_size, Image.LANCZOS)
-    return img
+    # Download InSPyReNet weights
+    if not os.path.exists('InSPyReNet.pth'):
+        os.system('wget https://github.com/plemeri/InSPyReNet/releases/download/v1.0/InSPyReNet.pth')
+        
+    # Load InSPyReNet model
+    inspyrenet = InSPyReNet()
+    inspyrenet.load_state_dict(torch.load('InSPyReNet.pth', map_location='cpu'))
+    inspyrenet.eval()
+    HAS_INSPYRE = True
+except:
+    HAS_INSPYRE = False
 
-def remove_background(input_image, post_process=True, alpha_matting=False):
-    start_time = time.time()
+def process_with_inspyrenet(image):
+    # Preprocess
+    image = test_dataset.preprocess(np.array(image))
+    image = torch.from_numpy(image).unsqueeze(0)
     
+    # Predict
+    with torch.no_grad():
+        pred = inspyrenet(image)
+    
+    # Post-process
+    mask = (pred.squeeze().cpu().numpy() > 0.5).astype(np.uint8) * 255
+    return mask
+
+def remove_background(input_image, model_choice="Rembg (U²-Net)"):
     try:
+        # Convert to PIL Image if it's a numpy array
         if isinstance(input_image, np.ndarray):
             input_image = Image.fromarray(input_image)
         
-        input_image = resize_image(input_image)
-        
-        output = remove(
-            input_image, 
-            post_process=post_process,
-            alpha_matting=alpha_matting,
-            alpha_matting_foreground_threshold=240,
-            alpha_matting_background_threshold=10,
-            alpha_matting_erode_size=10
-        )
-        
-        if output.mode == 'RGBA':
-            mask = output.split()[-1]
-            mask_np = np.array(mask)
-            if post_process:
-                mask_np = enhance_mask(mask_np)
+        # Process with selected model
+        if model_choice == "InSPyReNet" and HAS_INSPYRE:
+            mask = process_with_inspyrenet(input_image)
+            mask_img = Image.fromarray(mask)
+            
+            # Apply mask to original image
+            output = input_image.copy()
+            output.putalpha(mask_img)
         else:
-            mask_np = np.ones(output.size[::-1], dtype=np.uint8) * 255
-        
-        if post_process:
-            output.putalpha(Image.fromarray(mask_np))
+            # Default to Rembg
+            output = remove(input_image)
+            if output.mode == 'RGBA':
+                mask = output.split()[-1]
+                mask_np = np.array(mask)
+            else:
+                mask_np = np.ones(output.size[::-1], dtype=np.uint8) * 255
+            mask_img = Image.fromarray(mask_np)
         
-        proc_time = time.time() - start_time
-        return output, Image.fromarray(mask_np), f"Processed in {proc_time:.2f} seconds"
+        return output, mask_img
     
     except Exception as e:
         print(f"Error processing image: {str(e)}")
-        return None, None, "Error processing image"
-
-# Custom CSS for better UI
-custom_css = """
-.gradio-container { max-width: 900px !important; }
-.output-image { border: 1px solid #e2e8f0 !important; border-radius: 8px !important; }
-.processing-time { font-size: 0.9em; color: #64748b; margin-top: 8px; }
-"""
+        return None, None
 
 # Create interface
-with gr.Blocks(css=custom_css) as demo:
-    gr.Markdown("""
-    # 🖼️ Professional Background Remover
-    *Powered by U²-Net with enhanced post-processing*
-    """)
-    
-    with gr.Row():
-        with gr.Column():
-            input_img = gr.Image(label="Upload Image", type="pil", elem_id="input-image")
-            with gr.Accordion("Advanced Options", open=False):
-                post_process = gr.Checkbox(label="Enhanced Post-Processing", value=True)
-                alpha_matting = gr.Checkbox(label="Use Alpha Matting (for fine details)", value=False)
-            submit_btn = gr.Button("Remove Background", variant="primary")
-            
-        with gr.Column():
-            output_img = gr.Image(label="Result", type="pil", elem_id="output-image")
-            output_mask = gr.Image(label="Segmentation Mask", type="pil")
-            time_text = gr.Markdown(elem_classes=["processing-time"])
-    
-    submit_btn.click(
-        fn=remove_background,
-        inputs=[input_img, post_process, alpha_matting],
-        outputs=[output_img, output_mask, time_text]
-    )
+iface = gr.Interface(
+    fn=remove_background,
+    inputs=[
+        gr.Image(type="pil", label="Input Image"),
+        gr.Radio(
+            choices=["Rembg (U²-Net)", "InSPyReNet"],
+            value="Rembg (U²-Net)",
+            label="Model Selection"
+        )
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Result with Transparent Background"), 
+        gr.Image(type="pil", label="Segmentation Mask")
+    ],
+    title="Hybrid Background Remover (CPU)",
+    description="""
+    Upload an image to remove the background. Choose between:
+    - Rembg (U²-Net): Faster (5-15 sec)
+    - InSPyReNet: More accurate but slower (15-30 sec)
+    """
+)
 
+# Launch with minimal configuration
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True
-    )
+    iface.launch(server_name="0.0.0.0", server_port=7860)