Update app.py

app.py

@@ -1,74 +1,36 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import FileResponse
-from playwright.async_api import async_playwright
-import os
-import base64
-import io
-import time
-
-app = FastAPI()
-
-# Enable CORS for Figma
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["https://*.figma.com"],
-    allow_credentials=True,
-    allow_methods=["POST", "GET"],
-    allow_headers=["*"],
+from transformers import pipeline
+import gradio as gr
+from PIL import Image
+import numpy as np
+
+# Initialize the image segmentation pipeline with the DIS-based model
+pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
+
+# Function to remove background
+def remove_background(input_image):
+    # Ensure input is a PIL Image
+    if isinstance(input_image, np.ndarray):
+        input_image = Image.fromarray(input_image)
+    
+    # Get the mask and the processed image
+    pillow_mask = pipe(input_image, return_mask=True)  # Returns a binary mask
+    no_bg_image = pipe(input_image)  # Returns the image with background removed
+    
+    # Ensure the output is a PIL Image with transparency
+    if isinstance(no_bg_image, np.ndarray):
+        no_bg_image = Image.fromarray(no_bg_image)
+    
+    return no_bg_image
+
+# Set up Gradio interface
+demo = gr.Interface(
+    fn=remove_background,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Image(type="pil", label="Image with Background Removed"),
+    title="Background Remover using DIS (IS-Net)",
+    description="Upload an image to remove its background using the briaai/RMBG-1.4 model."
 )
 
-# Serve static files (index.html, assets)
-@app.get("/")
-async def serve_index():
-    return FileResponse("index.html")
-
-# API endpoint for background removal
-@app.post("/remove-background")
-async def remove_background(file: UploadFile = File(...)):
-    try:
-        # Save uploaded image
-        contents = await file.read()
-        temp_path = "temp_image.png"
-        with open(temp_path, "wb") as f:
-            f.write(contents)
-
-        # Use Playwright to process image via index.html
-        async with async_playwright() as p:
-            browser = await p.chromium.launch()
-            page = await browser.new_page()
-
-            # Load the Space's frontend
-            await page.goto("http://localhost:7860/")
-
-            # Upload image
-            await page.set_input_files("#upload", temp_path)
-
-            # Wait for processing (adjust based on JS behavior)
-            await page.wait_for_timeout(5000)  # Adjust if needed
-
-            # Extract processed image (assumes JS sets it in #status or similar)
-            # This depends on index-E_M5nW8h.js output (e.g., canvas or img)
-            processed_image = await page.evaluate("""
-                () => {
-                    const canvas = document.querySelector('canvas'); // Adjust selector
-                    if (canvas) {
-                        return canvas.toDataURL('image/png').split(',')[1];
-                    }
-                    return null;
-                }
-            """)
-
-            await browser.close()
-
-        # Clean up
-        os.remove(temp_path)
-
-        if not processed_image:
-            raise HTTPException(status_code=500, detail="Failed to process image")
-
-        return {"image": processed_image}
-    except Exception as e:
-        if os.path.exists(temp_path):
-            os.remove(temp_path)
-        raise HTTPException(status_code=500, detail=str(e))
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()