Create app.py

app.py

@@ -0,0 +1,121 @@
+import gradio as gr
+import numpy as np
+import cv2
+from PIL import Image
+import requests
+import torch
+from torchvision.transforms import ToTensor, ToPILImage
+
+# Initialize models (will be loaded on first use)
+models = {
+    "BRIA": None,
+    "INSPyReNet": None,
+    "U2Net": None,
+    "U2NetHumanSeg": None,
+    "ISNetGeneral": None,
+    "ISNetAnime": None
+}
+
+# Model URLs and loading functions
+def load_model(model_name):
+    if model_name == "BRIA":
+        from transformers import AutoModelForImageSegmentation
+        return AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True)
+    elif model_name == "INSPyReNet":
+        from IS2.models.inspyrenet import INSPyReNet
+        model = INSPyReNet()
+        model.load_state_dict(torch.hub.load_state_dict_from_url("https://github.com/helloyufei/INSPyReNet/releases/download/v1.0.0/inspyrenet.pth"))
+        return model
+    elif model_name == "U2Net":
+        import u2net
+        return u2net.load_model(model_name="u2net")
+    elif model_name == "U2NetHumanSeg":
+        import u2net
+        return u2net.load_model(model_name="u2net_human_seg")
+    elif model_name == "ISNetGeneral":
+        from isnet import ISNet
+        model = ISNet()
+        model.load_state_dict(torch.hub.load_state_dict_from_url("https://github.com/xuebinqin/DIS/raw/main/IS-Net/isnet-general-use.pth"))
+        return model
+    elif model_name == "ISNetAnime":
+        from isnet import ISNet
+        model = ISNet()
+        model.load_state_dict(torch.hub.load_state_dict_from_url("https://github.com/xuebinqin/DIS/raw/main/IS-Net/isnet-anime.pth"))
+        return model
+
+def apply_model(image, model_name):
+    if models[model_name] is None:
+        models[model_name] = load_model(model_name)
+    
+    model = models[model_name]
+    model.eval()
+    
+    # Preprocess image based on model requirements
+    if model_name == "BRIA":
+        from transformers import AutoImageProcessor
+        processor = AutoImageProcessor.from_pretrained("briaai/RMBG-1.4")
+        inputs = processor(images=image, return_tensors="pt")
+        with torch.no_grad():
+            outputs = model(**inputs)
+        mask = outputs.logits.squeeze().cpu().numpy()
+    elif model_name in ["U2Net", "U2NetHumanSeg"]:
+        input_img = np.array(image)
+        input_img = cv2.resize(input_img, (320, 320))
+        input_img = ToTensor()(input_img).unsqueeze(0)
+        with torch.no_grad():
+            mask = model(input_img).squeeze().cpu().numpy()
+    else: # INSPyReNet, ISNet models
+        input_img = np.array(image)
+        input_img = cv2.resize(input_img, (1024, 1024))
+        input_img = ToTensor()(input_img).unsqueeze(0)
+        with torch.no_grad():
+            mask = model(input_img).squeeze().cpu().numpy()
+    
+    # Post-process mask
+    mask = (mask - mask.min()) / (mask.max() - mask.min())
+    mask = (mask * 255).astype(np.uint8)
+    mask = cv2.resize(mask, (image.width, image.height))
+    return Image.fromarray(mask)
+
+def combine_masks(masks):
+    # Combine masks using weighted average
+    combined = np.zeros_like(masks[0], dtype=np.float32)
+    weights = [0.3, 0.2, 0.15, 0.15, 0.1, 0.1]  # Adjust weights as needed
+    
+    for mask, weight in zip(masks, weights):
+        combined += mask.astype(np.float32) * weight
+    
+    combined = np.clip(combined, 0, 255).astype(np.uint8)
+    return Image.fromarray(combined)
+
+def remove_background(image):
+    # Convert to PIL Image if needed
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    
+    # Apply all models in sequence
+    masks = []
+    for model_name in ["BRIA", "INSPyReNet", "U2Net", "U2NetHumanSeg", "ISNetGeneral", "ISNetAnime"]:
+        mask = apply_model(image, model_name)
+        masks.append(np.array(mask))
+    
+    # Combine masks
+    final_mask = combine_masks(masks)
+    
+    # Apply mask to original image
+    result = image.copy()
+    result.putalpha(final_mask)
+    
+    return result
+
+# Gradio interface
+interface = gr.Interface(
+    fn=remove_background,
+    inputs=gr.Image(label="Input Image"),
+    outputs=gr.Image(label="Background Removed", type="pil"),
+    title="Multi-Model Background Removal",
+    description="Combines BRIA, INSPyReNet, U²-Net, U²-Net Human Seg, ISNet-General-Use, and ISNet-Anime for high-quality background removal"
+)
+
+if __name__ == "__main__":
+    interface.launch()