Update app.py

app.py

@@ -1,207 +1,121 @@
-# import gradio as gr
-# import torch
-# import uuid
-# from PIL import Image
-# from torchvision import transforms
-# from transformers import AutoModelForImageSegmentation
-# from typing import Union, List
-# from loadimg import load_img  # Your helper to load from URL or file
-
-# torch.set_float32_matmul_precision("high")
-
-# # Load BiRefNet model
-# birefnet = AutoModelForImageSegmentation.from_pretrained(
-#     "ZhengPeng7/BiRefNet", trust_remote_code=True
-# )
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# birefnet.to(device)
-
-# # Image transformation
-# transform_image = transforms.Compose([
-#     transforms.Resize((1024, 1024)),
-#     transforms.ToTensor(),
-#     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-# ])
-
-# def process(image: Image.Image) -> Image.Image:
-#     image_size = image.size
-#     input_tensor = transform_image(image).unsqueeze(0).to(device)
-
-#     with torch.no_grad():
-#         preds = birefnet(input_tensor)[-1].sigmoid().cpu()
-
-#     pred = preds[0].squeeze()
-#     mask = transforms.ToPILImage()(pred).resize(image_size).convert("L")
-#     binary_mask = mask.point(lambda p: 255 if p > 127 else 0)
-
-#     white_bg = Image.new("RGB", image_size, (255, 255, 255))
-#     result = Image.composite(image, white_bg, binary_mask)
-#     return result
-
-# def handler(image=None, image_url=None, batch_urls=None) -> Union[str, List[str], None]:
-#     results = []
-
-#     try:
-#         # Single image upload
-#         if image is not None:
-#             image = image.convert("RGB")
-#             processed = process(image)
-#             filename = f"output_{uuid.uuid4().hex[:8]}.png"
-#             processed.save(filename)
-#             return filename
-
-#         # Single image from URL
-#         if image_url:
-#             im = load_img(image_url, output_type="pil").convert("RGB")
-#             processed = process(im)
-#             filename = f"output_{uuid.uuid4().hex[:8]}.png"
-#             processed.save(filename)
-#             return filename
-
-#         # Batch of URLs
-#         if batch_urls:
-#             urls = [u.strip() for u in batch_urls.split(",") if u.strip()]
-#             for url in urls:
-#                 try:
-#                     im = load_img(url, output_type="pil").convert("RGB")
-#                     processed = process(im)
-#                     filename = f"output_{uuid.uuid4().hex[:8]}.png"
-#                     processed.save(filename)
-#                     results.append(filename)
-#                 except Exception as e:
-#                     print(f"Error with {url}: {e}")
-#             return results if results else None
-
-#     except Exception as e:
-#         print("General error:", e)
-
-#     return None
-
-# # Interface
-# demo = gr.Interface(
-#     fn=handler,
-#     inputs=[
-#         gr.Image(label="Upload Image", type="pil"),
-#         gr.Textbox(label="Paste Image URL"),
-#         gr.Textbox(label="Comma-separated Image URLs (Batch)"),
-#     ],
-#     outputs=gr.File(label="Output File(s)", file_count="multiple"),
-#     title="Background Remover (White Fill)",
-#     description="Upload an image, paste a URL, or send a batch of URLs to remove the background and replace it with white.",
-# )
-
-# if __name__ == "__main__":
-#     demo.launch(show_error=True, mcp_server=True)
-
-
-
 import gradio as gr
 import torch
 import uuid
 import base64
+import numpy as np
+import onnxruntime as ort
+import cv2
 from PIL import Image
-from torchvision import transforms
-from transformers import AutoModelForImageSegmentation
+from torchvision.transforms.functional import normalize
+import torch.nn.functional as F
 from typing import Union, List
-from loadimg import load_img  # Your helper to load from URL or file
 from io import BytesIO
-
-torch.set_float32_matmul_precision("high")
-
-# Load BiRefNet model
-birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet", trust_remote_code=True
-)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-birefnet.to(device)
-
-# Image transformation
-transform_image = transforms.Compose([
-    transforms.Resize((1024, 1024)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-])
-
-def load_image_from_data_url(data_url: str) -> Image.Image:
-    """Load image from base64 data URL"""
-    if data_url.startswith("data:image/"):
-        # Extract base64 data after the comma
-        if "," in data_url:
-            header, encoded = data_url.split(",", 1)
-            image_data = base64.b64decode(encoded)
-            return Image.open(BytesIO(image_data))
-        else:
-            raise ValueError(f"Invalid data URL format: {data_url[:50]}...")
-    else:
-        # Regular URL, use existing load_img function
-        return load_img(data_url, output_type="pil")
-
+from huggingface_hub import hf_hub_download
+
+# ---- Config ----
+INPUT_SIZE = [1200, 1800]  # (H, W)
+
+# ---- Load ONNX model ----
+model_path = hf_hub_download(repo_id="Trendyol/background-removal", filename="model.onnx")
+providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+try:
+    ort_sess = ort.InferenceSession(model_path, providers=providers)
+except Exception:
+    ort_sess = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
+
+
+# ---- Utils from Trendyol ----
+def keep_large_components(a: np.ndarray) -> np.ndarray:
+    dilate_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9))
+    a_mask = (a > 25).astype(np.uint8) * 255
+    analysis = cv2.connectedComponentsWithStats(a_mask, 4, cv2.CV_32S)
+    (totalLabels, label_ids, values, _) = analysis
+
+    h, w = a.shape[:2]
+    area_limit = 50000 * (h * w) / (INPUT_SIZE[1] * INPUT_SIZE[0])
+    i_to_keep = []
+    for i in range(1, totalLabels):
+        area = values[i, cv2.CC_STAT_AREA]
+        if area > area_limit:
+            i_to_keep.append(i)
+
+    if len(i_to_keep) > 0:
+        final_mask = np.zeros_like(a, dtype=np.uint8)
+        for i in i_to_keep:
+            componentMask = (label_ids == i).astype("uint8") * 255
+            final_mask = cv2.bitwise_or(final_mask, componentMask)
+        final_mask = cv2.dilate(final_mask, dilate_kernel, iterations=2)
+        a = cv2.bitwise_and(a, final_mask)
+
+    a = a.reshape((a.shape[0], a.shape[1], 1))
+    return a
+
+
+def preprocess_input(im: np.ndarray) -> torch.Tensor:
+    if len(im.shape) < 3:
+        im = im[:, :, np.newaxis]
+    if im.shape[2] == 4:
+        im = im[:, :, :3]
+    im_tensor = torch.tensor(im, dtype=torch.float32).permute(2, 0, 1)
+    im_tensor = F.upsample(torch.unsqueeze(im_tensor, 0), INPUT_SIZE, mode="bilinear").type(torch.uint8)
+    image = torch.divide(im_tensor, 255.0)
+    image = normalize(image, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
+    return image
+
+
+def postprocess_output(result: np.ndarray, orig_im_shape) -> np.ndarray:
+    result = torch.squeeze(
+        F.upsample(torch.from_numpy(result).unsqueeze(0), (orig_im_shape), mode="bilinear"), 0
+    )
+    ma = torch.max(result)
+    mi = torch.min(result)
+    result = (result - mi) / (ma - mi + 1e-8)
+    a = (result * 255).permute(1, 2, 0).cpu().data.numpy().astype(np.uint8)
+    a = keep_large_components(a)
+    return a
+
+
+# ---- Core processing ----
 def process(image: Image.Image) -> Image.Image:
     image_size = image.size
-    input_tensor = transform_image(image).unsqueeze(0).to(device)
+    np_img = np.array(image.convert("RGB"))
 
-    with torch.no_grad():
-        preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+    # Preprocess
+    img_tensor = preprocess_input(np_img)
 
-    pred = preds[0].squeeze()
-    mask = transforms.ToPILImage()(pred).resize(image_size).convert("L")
-    binary_mask = mask.point(lambda p: 255 if p > 127 else 0)
+    # Inference
+    inputs = {ort_sess.get_inputs()[0].name: img_tensor.numpy()}
+    result = ort_sess.run(None, inputs)[0][0]  # (1,1,H,W)
 
+    # Postprocess to mask
+    alpha = postprocess_output(result, (np_img.shape[0], np_img.shape[1]))  # (H,W,1)
+
+    # White background composite
+    mask = Image.fromarray(alpha.squeeze(-1)).convert("L")
+    binary_mask = mask.point(lambda p: 255 if p > 25 else 0)
     white_bg = Image.new("RGB", image_size, (255, 255, 255))
-    result = Image.composite(image, white_bg, binary_mask)
+    result = Image.composite(image.convert("RGB"), white_bg, binary_mask)
     return result
 
-def handler(image=None, image_url=None, batch_urls=None) -> Union[str, List[str], None]:
-    results = []
-
-    try:
-        # Single image upload
-        if image is not None:
-            image = image.convert("RGB")
-            processed = process(image)
-            filename = f"output_{uuid.uuid4().hex[:8]}.png"
-            processed.save(filename)
-            return filename
-
-        # Single image from URL (supports both regular URLs and base64 data URLs)
-        if image_url:
-            im = load_image_from_data_url(image_url).convert("RGB")
-            processed = process(im)
-            filename = f"output_{uuid.uuid4().hex[:8]}.png"
-            processed.save(filename)
-            return filename
-
-        # Batch of URLs (supports both regular URLs and base64 data URLs)
-        if batch_urls:
-            urls = [u.strip() for u in batch_urls.split(",") if u.strip()]
-            for url in urls:
-                try:
-                    im = load_image_from_data_url(url).convert("RGB")
-                    processed = process(im)
-                    filename = f"output_{uuid.uuid4().hex[:8]}.png"
-                    processed.save(filename)
-                    results.append(filename)
-                except Exception as e:
-                    print(f"Error with {url}: {e}")
-            return results if results else None
-
-    except Exception as e:
-        print("General error:", e)
 
+# ---- Gradio handler ----
+def handler(image=None) -> Union[str, None]:
+    if image is not None:
+        processed = process(image)
+        filename = f"output_{uuid.uuid4().hex[:8]}.png"
+        processed.save(filename)
+        return filename
     return None
 
-# Interface
+
+# ---- Gradio UI ----
 demo = gr.Interface(
     fn=handler,
-    inputs=[
-        gr.Image(label="Upload Image", type="pil"),
-        gr.Textbox(label="Paste Image URL"),
-        gr.Textbox(label="Comma-separated Image URLs (Batch)"),
-    ],
-    outputs=gr.File(label="Output File(s)", file_count="multiple"),
-    title="Background Remover (White Fill)",
-    description="Upload an image, paste a URL, or send a batch of URLs to remove the background and replace it with white.",
+    inputs=gr.Image(label="Upload Image", type="pil"),
+    outputs=gr.File(label="Output File"),
+    title="Background Remover (Trendyol)",
+    description="Upload an image to remove the background with the Trendyol ONNX model. Background is replaced with white.",
 )
 
 if __name__ == "__main__":
-    demo.launch(show_error=True, mcp_server=True)
+    demo.launch(show_error=True)