Update app.py

app.py

@@ -3,92 +3,97 @@ from rembg import remove
 from PIL import Image
 import numpy as np
 import torch
-import cv2
 import os
+import requests
+from tqdm import tqdm
+import subprocess
 
-# Initialize InSPyReNet if available
-try:
-    from InSPyReNet.models.InSPyReNet import InSPyReNet
-    from InSPyReNet.utils.dataloader import test_dataset
+# Clone InSPyReNet repository if not present
+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')
+
+# Download model weights if not present
+def download_file(url, filename):
+    response = requests.get(url, stream=True)
+    total_size = int(response.headers.get('content-length', 0))
     
-    # 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
+    with open(filename, 'wb') as f, tqdm(
+        desc=filename,
+        total=total_size,
+        unit='iB',
+        unit_scale=True,
+        unit_divisor=1024,
+    ) as bar:
+        for data in response.iter_content(chunk_size=1024):
+            size = f.write(data)
+            bar.update(size)
+
+if not os.path.exists('InSPyReNet.pth'):
+    print("Downloading model weights...")
+    download_file(
+        "https://github.com/plemeri/InSPyReNet/releases/download/v1.0/InSPyReNet.pth",
+        "InSPyReNet.pth"
+    )
+
+# Import after setup
+from InSPyReNet import InSPyReNet
+from modules.layers import load_model
+from utils.misc import load_config
+
+# Initialize model
+print("Loading model...")
+cfg = load_config('InSPyReNet/configs/InSPyReNet_SwinB.yaml')
+device = torch.device('cpu')
+model = InSPyReNet(cfg)
+model = load_model(model, 'InSPyReNet.pth', device)
+model.eval()
 
 def process_with_inspyrenet(image):
-    # Preprocess
-    image = test_dataset.preprocess(np.array(image))
-    image = torch.from_numpy(image).unsqueeze(0)
+    # Convert to numpy and normalize
+    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
     with torch.no_grad():
-        pred = inspyrenet(image)
+        pred = model(image)
     
-    # Post-process
+    # Create mask
     mask = (pred.squeeze().cpu().numpy() > 0.5).astype(np.uint8) * 255
     return mask
 
-def remove_background(input_image, model_choice="Rembg (U²-Net)"):
+def remove_background(input_image):
     try:
-        # Convert to PIL Image if it's a numpy array
         if isinstance(input_image, np.ndarray):
             input_image = Image.fromarray(input_image)
         
-        # 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:
-            # 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)
+        # Process with InSPyReNet
+        mask = process_with_inspyrenet(input_image)
+        output = input_image.copy()
+        output.putalpha(Image.fromarray(mask))
         
-        return output, mask_img
+        return output, Image.fromarray(mask)
     
     except Exception as e:
-        print(f"Error processing image: {str(e)}")
+        print(f"Error: {str(e)}")
         return None, None
 
-# Create interface
 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"
-        )
-    ],
+    inputs=gr.Image(type="pil", label="Input Image"),
     outputs=[
-        gr.Image(type="pil", label="Result with Transparent Background"), 
-        gr.Image(type="pil", label="Segmentation Mask")
+        gr.Image(type="pil", label="Result"), 
+        gr.Image(type="pil", label="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)
-    """
+    title="Professional Background Remover",
+    description="Using InSPyReNet for high-quality background removal"
 )
 
-# Launch with minimal configuration
 if __name__ == "__main__":
     iface.launch(server_name="0.0.0.0", server_port=7860)