Update app.py

app.py

@@ -5,31 +5,38 @@ import numpy as np
 from PIL import Image
 import io
 
-# First, ensure all required dependencies are installed
-try:
-    import torchvision
-    import skimage
-except ImportError:
-    print("Installing required packages...")
-    import subprocess
-    subprocess.check_call(["pip", "install", "torchvision", "scikit-image"])
-    import torchvision
-    import skimage
-
 def remove_background(input_image):
     try:
-        # Initialize the pipeline with correct parameters and dependencies
+        # Convert input to PIL Image if it's not already
+        if not isinstance(input_image, Image.Image):
+            input_image = Image.fromarray(input_image)
+            
+        # Initialize the pipeline
         segmentor = pipeline(
-            "image-segmentation", 
+            task="image-segmentation",
             model="briaai/RMBG-1.4",
-            trust_remote_code=True,
-            device="cpu",
-            framework="pt"
+            trust_remote_code=True
+        )
+        
+        # Process the image and get mask
+        result = segmentor(
+            input_image, 
+            return_mask=True
         )
         
-        # Process the image
-        result = segmentor(input_image)
-        return result['output_image']
+        # Create output image with transparent background
+        output_image = Image.new('RGBA', input_image.size, (0, 0, 0, 0))
+        
+        # Convert input to RGBA if it's not already
+        if input_image.mode != 'RGBA':
+            input_image = input_image.convert('RGBA')
+            
+        # Apply mask to create transparent background
+        mask = result['mask'] if isinstance(result, dict) else result
+        output_image.paste(input_image, mask=mask)
+        
+        return output_image
+        
     except Exception as e:
         raise gr.Error(f"Error processing image: {str(e)}")