Create app.py

app.py

@@ -0,0 +1,99 @@
+import numpy as np
+from PIL import Image
+import requests
+from io import BytesIO
+import cv2
+
+# Load models from Hugging Face Hub
+from huggingface_hub import from_pretrained_keras, hf_hub_download
+import tensorflow as tf
+
+# Load U²-Net (we'll use a lightweight version suitable for CPU)
+def load_u2net():
+    model_path = hf_hub_download(repo_id="skytnt/anime-remove-background", filename="u2netp.onnx")
+    net = cv2.dnn.readNetFromONNX(model_path)
+    return net
+
+# Load BRIA model (using a CPU-compatible version)
+def load_bria():
+    model = from_pretrained_keras("briaai/RMBG-1.4", compile=False)
+    return model
+
+# Preprocess image for U²-Net
+def preprocess_u2net(image):
+    image = image.resize((320, 320))
+    image = np.array(image)
+    image = image / 255.0
+    image = image.transpose(2, 0, 1)
+    image = np.expand_dims(image, axis=0).astype('float32')
+    return image
+
+# Preprocess image for BRIA
+def preprocess_bria(image):
+    image = image.resize((1024, 1024))
+    image = np.array(image)
+    image = image / 255.0
+    image = image.astype('float32')
+    return np.expand_dims(image, axis=0)
+
+# Postprocess mask
+def postprocess_mask(mask):
+    mask = mask.squeeze()
+    mask = (mask * 255).astype('uint8')
+    mask = Image.fromarray(mask).resize((original_width, original_height))
+    return mask
+
+# Compare masks and select the better one
+def select_better_mask(mask1, mask2):
+    # Simple heuristic: select the mask with more defined edges
+    # You can implement more sophisticated comparison if needed
+    edge1 = cv2.Canny(np.array(mask1), 100, 200)
+    edge2 = cv2.Canny(np.array(mask2), 100, 200)
+    return mask1 if np.sum(edge1) > np.sum(edge2) else mask2
+
+# Load models (we'll do this once when the Space starts)
+u2net = load_u2net()
+bria = load_bria()
+
+def remove_background(image):
+    global original_width, original_height
+    original_width, original_height = image.size
+    
+    # Process with U²-Net
+    u2net_input = preprocess_u2net(image)
+    u2net.setInput(u2net_input)
+    u2net_mask = u2net.forward()
+    u2net_mask = postprocess_mask(u2net_mask[0][0])
+    
+    # Process with BRIA
+    bria_input = preprocess_bria(image)
+    bria_mask = bria.predict(bria_input)
+    bria_mask = postprocess_mask(bria_mask[0][:, :, 0])
+    
+    # Select better mask
+    final_mask = select_better_mask(u2net_mask, bria_mask)
+    
+    # Apply mask to original image
+    image = image.convert("RGBA")
+    final_mask = final_mask.convert("L")
+    image.putalpha(final_mask)
+    
+    return image
+
+# Gradio interface
+import gradio as gr
+
+def process_image(input_image):
+    image = Image.fromarray(input_image)
+    result = remove_background(image)
+    return result
+
+iface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(),
+    outputs=gr.Image(type="pil"),
+    title="Background Removal Pipeline (BRIA + U²-Net)",
+    description="Combines BRIA and U²-Net models for better background removal (CPU-only version)"
+)
+
+iface.launch()