Update app.py

app.py

@@ -4,20 +4,21 @@ from PIL import Image
 import numpy as np
 import torch
 import os
+import sys
 import requests
 from tqdm import tqdm
 import subprocess
 
-# Clone InSPyReNet repository if not present
+# Clone and set up repository
 if not os.path.exists('InSPyReNet'):
     print("Cloning InSPyReNet repository...")
     subprocess.run(['git', 'clone', 'https://github.com/plemeri/InSPyReNet.git'])
 
-# Add to Python path
-import sys
-sys.path.insert(0, 'InSPyReNet/lib')
+# Add correct paths to system
+sys.path.insert(0, os.path.abspath('InSPyReNet'))
+sys.path.insert(0, os.path.abspath('InSPyReNet/lib'))
 
-# Download model weights if not present
+# Download model weights
 def download_file(url, filename):
     response = requests.get(url, stream=True)
     total_size = int(response.headers.get('content-length', 0))
@@ -40,7 +41,7 @@ if not os.path.exists('InSPyReNet.pth'):
         "InSPyReNet.pth"
     )
 
-# Import after setup
+# Import after setting paths
 from InSPyReNet import InSPyReNet
 from modules.layers import load_model
 from utils.misc import load_config
@@ -53,32 +54,36 @@ model = InSPyReNet(cfg)
 model = load_model(model, 'InSPyReNet.pth', device)
 model.eval()
 
-def process_with_inspyrenet(image):
-    # Convert to numpy and normalize
+def preprocess(image):
     image = np.array(image).astype(np.float32)
     image -= np.array([104.00699, 116.66877, 122.67892])
     image = image.transpose((2, 0, 1))
-    image = torch.from_numpy(image).unsqueeze(0).to(device)
-    
-    # Predict
+    return torch.from_numpy(image).unsqueeze(0)
+
+def process_with_inspyrenet(image):
+    image_tensor = preprocess(image).to(device)
     with torch.no_grad():
-        pred = model(image)
-    
-    # Create mask
+        pred = model(image_tensor)
     mask = (pred.squeeze().cpu().numpy() > 0.5).astype(np.uint8) * 255
     return mask
 
-def remove_background(input_image):
+def remove_background(input_image, model_choice="Rembg (U²-Net)"):
     try:
         if isinstance(input_image, np.ndarray):
             input_image = Image.fromarray(input_image)
         
-        # Process with InSPyReNet
-        mask = process_with_inspyrenet(input_image)
-        output = input_image.copy()
-        output.putalpha(Image.fromarray(mask))
+        if model_choice == "InSPyReNet":
+            mask = process_with_inspyrenet(input_image)
+            output = input_image.copy()
+            output.putalpha(Image.fromarray(mask))
+        else:
+            output = remove(input_image)
+            if output.mode == 'RGBA':
+                mask = output.split()[-1]
+            else:
+                mask = Image.new('L', output.size, 255)
         
-        return output, Image.fromarray(mask)
+        return output, mask
     
     except Exception as e:
         print(f"Error: {str(e)}")
@@ -86,13 +91,20 @@ def remove_background(input_image):
 
 iface = gr.Interface(
     fn=remove_background,
-    inputs=gr.Image(type="pil", label="Input Image"),
+    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"), 
         gr.Image(type="pil", label="Mask")
     ],
     title="Professional Background Remover",
-    description="Using InSPyReNet for high-quality background removal"
+    description="Choose between Rembg (faster) or InSPyReNet (higher quality)"
 )
 
 if __name__ == "__main__":