Update app.py

app.py

@@ -1,26 +1,18 @@
-import gradio as gr
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import Response
 import numpy as np
 from PIL import Image
-import torch
-from transformers import pipeline
-from functools import lru_cache
+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__)
 
-# Cache models to avoid reloading on every request
-@lru_cache(maxsize=1)
-def load_model(model_name):
-    try:
-        return pipeline("image-segmentation", model_name)
-    except Exception as e:
-        logger.error(f"Failed to load {model_name}: {e}")
-        return None
-
-# Model sequence configuration
 MODELS = [
     {"name": "BRIA", "repo": "BRIA-AI/bria-rmbg", "weight": 1.0},
     {"name": "INSPyReNet", "repo": "mattmdjaga/INSPyReNet", "weight": 0.9},
@@ -30,91 +22,64 @@ MODELS = [
     {"name": "ISNet-Anime", "repo": "skytnt/anime-seg", "weight": 0.5}
 ]
 
-def process_single_model(image, model):
-    """Process image with a single model"""
+def load_model(model_repo):
     try:
-        pipe = load_model(model["repo"])
-        if pipe is None:
-            return None
-            
-        # Convert image to numpy array if needed
-        if isinstance(image, Image.Image):
-            image_np = np.array(image)
-        else:
-            image_np = image
-            
-        result = pipe(image_np)
-        return result['mask'] if isinstance(result, dict) else result[0]['mask']
+        return pipeline("image-segmentation", model_repo)
     except Exception as e:
-        logger.warning(f"{model['name']} failed: {e}")
+        logger.error(f"Failed to load {model_repo}: {e}")
         return None
 
-def combine_masks(masks, weights):
-    """Combine masks with weighted averaging"""
-    valid_masks = [m for m in masks if m is not None]
-    if not valid_masks:
-        return None
-        
-    total_weight = sum(w for w, m in zip(weights, masks) if m is not None)
-    combined = np.zeros_like(valid_masks[0], dtype=np.float32)
+def process_image(image: np.ndarray):
+    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}")
+    
+    if not masks:
+        raise ValueError("All models failed")
+    
+    # Weighted average of masks
+    total_weight = sum(weights)
+    combined = np.zeros_like(masks[0], dtype=np.float32)
     for mask, weight in zip(masks, weights):
-        if mask is not None:
-            combined += (mask.astype(np.float32) * weight
+        combined += mask.astype(np.float32) * weight
+    final_mask = (combined / total_weight).astype(np.uint8)
     
-    return (combined / total_weight).astype(np.uint8)
+    return final_mask
 
-def remove_background(image):
-    """Main processing pipeline"""
+@app.post("/remove-background")
+async def remove_background(file: UploadFile = File(...)):
     try:
-        # Convert input to PIL Image
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image)
-            
-        # Process through all models
-        masks = []
-        for model in MODELS:
-            mask = process_single_model(image, model)
-            masks.append(mask)
-            
-        # Combine results
-        weights = [m["weight"] for m in MODELS]
-        final_mask = combine_masks(masks, weights)
+        # 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)
         
-        if final_mask is None:
-            raise ValueError("All models failed")
-            
         # Apply mask
         background = Image.new('RGB', image.size, (0, 0, 0))
-        final_image = Image.composite(image, background, Image.fromarray(final_mask))
+        result = Image.composite(image, background, Image.fromarray(mask))
         
-        return final_image
-    except Exception as e:
-        logger.error(f"Processing failed: {e}")
-        return None
-
-# Gradio interface with API endpoint
-with gr.Blocks() as app:
-    gr.Markdown("## 🖼️ Advanced Background Remover")
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(label="Upload Image")
-            submit_btn = gr.Button("Remove Background")
-        with gr.Column():
-            output_image = gr.Image(label="Result")
+        # 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")
     
-    submit_btn.click(
-        fn=remove_background,
-        inputs=input_image,
-        outputs=output_image
-    )
-    
-    # API endpoint for mobile apps
-    app.api_app = gr.routes.App.create_app(
-        fn=remove_background,
-        inputs=gr.Image(),
-        outputs=gr.Image()
-    )
+    except Exception as e:
+        logger.error(f"Error: {e}")
+        return {"error": str(e)}, 500
 
-if __name__ == "__main__":
-    app.launch(server_name="0.0.0.0", server_port=7860)
+@app.get("/")
+def health_check():
+    return {"status": "healthy", "models": [m["name"] for m in MODELS]}