Upload 2 files

utils/u2net_detector.py

@@ -0,0 +1,182 @@
+import os
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from shapely.geometry import Polygon
+
+# -------------------------------------------------------------------
+# U²-Netp Model Definition (lightweight 4.7MB)
+# -------------------------------------------------------------------
+# Source: https://github.com/xuebinqin/U-2-Net
+# We include only the necessary modules.
+
+class REBNCONV(torch.nn.Module):
+    def __init__(self, in_ch=3, out_ch=3, dirate=1):
+        super(REBNCONV, self).__init__()
+        self.conv_s1 = torch.nn.Conv2d(
+            in_ch, out_ch, 3, padding=1*dirate, dilation=1*dirate
+        )
+        self.relu_s1 = torch.nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        hx = x
+        hx = self.relu_s1(self.conv_s1(hx))
+        return hx
+
+
+class RSU4F(torch.nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch)
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch)
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch)
+
+    def forward(self, x):
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+        hx4 = self.rebnconv4(hx3)
+
+        return hxin + hx4
+
+
+class U2NETP(torch.nn.Module):
+    def __init__(self, in_ch=3, out_ch=1):
+        super(U2NETP, self).__init__()
+
+        self.stage1 = RSU4F(in_ch, 12, 64)
+        self.stage2 = RSU4F(64, 12, 64)
+        self.stage3 = RSU4F(64, 12, 64)
+        self.stage4 = RSU4F(64, 12, 64)
+        self.stage5 = RSU4F(64, 12, 64)
+        self.stage6 = RSU4F(64, 12, 64)
+        self.side6 = torch.nn.Conv2d(64, out_ch, 3, padding=1)
+
+    def forward(self, x):
+        hx1 = self.stage1(x)
+        hx2 = self.stage2(hx1)
+        hx3 = self.stage3(hx2)
+
+        hx6 = self.stage6(hx3)
+        d6 = self.side6(hx6)
+
+        return torch.sigmoid(d6)
+
+
+# -------------------------------------------------------------------
+# Load Model (once)
+# -------------------------------------------------------------------
+MODEL_PATH = os.path.join(os.path.dirname(__file__), "u2netp.pth")
+
+_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+_u2net_model = None
+
+
+def load_u2netp():
+    global _u2net_model
+    if _u2net_model is None:
+        print("🔄 Loading U²-Netp model…")
+        model = U2NETP()
+        model.load_state_dict(torch.load(MODEL_PATH, map_location=_device))
+        model.to(_device)
+        model.eval()
+        _u2net_model = model
+        print("✅ U²-Netp Loaded.")
+    return _u2net_model
+
+
+# -------------------------------------------------------------------
+# Preprocessing
+# -------------------------------------------------------------------
+def preprocess(img_pil, size=320):
+    img = img_pil.convert("RGB")
+    img = img.resize((size, size), Image.BILINEAR)
+    arr = np.array(img).astype(np.float32) / 255.0
+
+    tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
+    return tensor.to(_device), img_pil.size
+
+
+# -------------------------------------------------------------------
+# Postprocessing → polygon conversion
+# -------------------------------------------------------------------
+def mask_to_polygons(mask, min_area=300):
+    """
+    Convert binary mask → list of polygons (list[list[(x,y)]])
+    """
+    mask = (mask * 255).astype("uint8")
+
+    # cleanup
+    kernel = np.ones((5,5), np.uint8)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
+
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+    polys = []
+
+    for cnt in contours:
+        area = cv2.contourArea(cnt)
+        if area < min_area:
+            continue
+
+        eps = 0.01 * cv2.arcLength(cnt, True)
+        approx = cv2.approxPolyDP(cnt, eps, True)
+
+        poly = [(int(p[0][0]), int(p[0][1])) for p in approx]
+        polys.append(poly)
+
+    return polys
+
+
+def resize_polygons(polygons, orig_w, orig_h, proc_size=320):
+    """Scale polygons back to original image size"""
+    scaled = []
+    for poly in polygons:
+        scaled.append([
+            (
+                int(x * orig_w / proc_size),
+                int(y * orig_h / proc_size)
+            )
+            for (x, y) in poly
+        ])
+    return scaled
+
+
+# -------------------------------------------------------------------
+# Main Bubble Detection Function
+# -------------------------------------------------------------------
+def detect_bubbles_u2net(img_pil, min_area=300):
+    """
+    Return list of bubble polygons from U²-Net saliency segmentation.
+    """
+    model = load_u2netp()
+
+    tensor, orig_size = preprocess(img_pil)
+    orig_w, orig_h = img_pil.size
+
+    with torch.no_grad():
+        pred = model(tensor)[0, 0].cpu().numpy()
+
+    # Normalize & threshold
+    pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)
+    mask = (pred > 0.4).astype(np.uint8)
+
+    # polygons from mask
+    polys = mask_to_polygons(mask, min_area=min_area)
+
+    # rescale to original image size
+    polys = resize_polygons(polys, orig_w, orig_h)
+
+    print(f"🧠 U²-Net bubbles detected: {len(polys)}")
+    return polys

utils/u2netp.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e7567cde013fb64813973ce6e1ecc25a80c05c3ca7adbc5a54f3c3d90991b854
+size 4683258