Create app.py

app.py

@@ -0,0 +1,120 @@
+import gradio as gr
+import numpy as np
+from PIL import Image
+import torch
+from transformers import pipeline
+from functools import lru_cache
+import cv2
+import logging
+
+# 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},
+    {"name": "U2Net", "repo": "silks-road/u2net", "weight": 0.8},
+    {"name": "U2Net-Human", "repo": "mattmdjaga/u2net-human-seg", "weight": 0.7},
+    {"name": "ISNet-General", "repo": "xuebinqin/ISNet-general-use", "weight": 0.6},
+    {"name": "ISNet-Anime", "repo": "skytnt/anime-seg", "weight": 0.5}
+]
+
+def process_single_model(image, model):
+    """Process image with a single model"""
+    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']
+    except Exception as e:
+        logger.warning(f"{model['name']} failed: {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)
+    
+    for mask, weight in zip(masks, weights):
+        if mask is not None:
+            combined += (mask.astype(np.float32) * weight
+    
+    return (combined / total_weight).astype(np.uint8)
+
+def remove_background(image):
+    """Main processing pipeline"""
+    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)
+        
+        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))
+        
+        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")
+    
+    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()
+    )
+
+if __name__ == "__main__":
+    app.launch(server_name="0.0.0.0", server_port=7860)