Update app.py

app.py

@@ -5,33 +5,36 @@ 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
+        # Initialize the pipeline with correct parameters and dependencies
         segmentor = pipeline(
-            task="image-segmentation",
+            "image-segmentation", 
             model="briaai/RMBG-1.4",
-            trust_remote_code=True
+            trust_remote_code=True,
+            device="cpu",
+            framework="pt"
         )
         
-        # Convert input to PIL Image if it's not already
-        if not isinstance(input_image, Image.Image):
-            input_image = Image.fromarray(input_image)
-            
         # Process the image
-        result = segmentor(input_image, return_mask=True)
-        
-        # Create transparent background image
-        output_image = Image.new('RGBA', input_image.size, (0, 0, 0, 0))
-        output_image.paste(input_image, mask=result)
-        
-        return output_image
-
+        result = segmentor(input_image)
+        return result['output_image']
     except Exception as e:
         raise gr.Error(f"Error processing image: {str(e)}")
 
 # Create Gradio interface
-with gr.Blocks(theme=gr.themes.Default()) as demo:
+with gr.Blocks() as demo:
     gr.HTML(
         """
         <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">