Update app.py

app.py

@@ -1,18 +1,10 @@
-from fastapi import FastAPI, UploadFile, File
-from fastapi.responses import Response
+import gradio as gr
 import numpy as np
 from PIL import Image
-import io
-import cv2
 from transformers import pipeline
-import logging
-
-app = FastAPI(title="Advanced Background Remover")
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+import cv2
 
+# Model sequence with weights
 MODELS = [
     {"name": "BRIA", "repo": "BRIA-AI/bria-rmbg", "weight": 1.0},
     {"name": "INSPyReNet", "repo": "mattmdjaga/INSPyReNet", "weight": 0.9},
@@ -23,63 +15,55 @@ MODELS = [
 ]
 
 def load_model(model_repo):
-    try:
-        return pipeline("image-segmentation", model_repo)
-    except Exception as e:
-        logger.error(f"Failed to load {model_repo}: {e}")
-        return None
+    return pipeline("image-segmentation", model_repo)
 
-def process_image(image: np.ndarray):
+def process_image(input_image):
+    # Convert Gradio input to PIL Image
+    if isinstance(input_image, np.ndarray):
+        input_image = Image.fromarray(input_image)
+    
     masks = []
     weights = []
     
     for model in MODELS:
-        pipe = load_model(model["repo"])
-        if pipe:
-            try:
-                result = pipe(image)
-                mask = result[0]['mask'] if isinstance(result, list) else result['mask']
-                masks.append(mask)
-                weights.append(model["weight"])
-            except Exception as e:
-                logger.warning(f"{model['name']} failed: {e}")
+        try:
+            pipe = load_model(model["repo"])
+            result = pipe(np.array(input_image))
+            mask = result[0]['mask'] if isinstance(result, list) else result['mask']
+            masks.append(mask)
+            weights.append(model["weight"])
+            print(f"{model['name']} completed successfully")  # Debug print
+        except Exception as e:
+            print(f"{model['name']} failed: {str(e)}")  # Debug print
+            continue
     
     if not masks:
-        raise ValueError("All models failed")
+        return None
     
     # Weighted average of masks
-    total_weight = sum(weights)
     combined = np.zeros_like(masks[0], dtype=np.float32)
     for mask, weight in zip(masks, weights):
         combined += mask.astype(np.float32) * weight
-    final_mask = (combined / total_weight).astype(np.uint8)
+    final_mask = (combined / sum(weights)).astype(np.uint8)
     
-    return final_mask
-
-@app.post("/remove-background")
-async def remove_background(file: UploadFile = File(...)):
-    try:
-        # Read and convert image
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
-        image_np = np.array(image)
-        
-        # Process image
-        mask = process_image(image_np)
-        
-        # Apply mask
-        background = Image.new('RGB', image.size, (0, 0, 0))
-        result = Image.composite(image, background, Image.fromarray(mask))
-        
-        # Return result
-        img_byte_arr = io.BytesIO()
-        result.save(img_byte_arr, format='PNG')
-        return Response(content=img_byte_arr.getvalue(), media_type="image/png")
+    # Create transparent background
+    result = input_image.copy()
+    result.putalpha(Image.fromarray(final_mask))
     
-    except Exception as e:
-        logger.error(f"Error: {e}")
-        return {"error": str(e)}, 500
+    return result
+
+# Gradio interface
+demo = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(label="Input Image"),
+    outputs=gr.Image(label="Result (PNG with Transparency)"),
+    title="🎨 Advanced Background Remover",
+    description="Combines 6 AI models for perfect background removal",
+    examples=[
+        ["example1.jpg"],
+        ["example2.jpg"],
+        ["example3.png"]
+    ]
+)
 
-@app.get("/")
-def health_check():
-    return {"status": "healthy", "models": [m["name"] for m in MODELS]}
+demo.launch(server_name="0.0.0.0", server_port=7860)