Create masks.py

helpers/masks.py

@@ -0,0 +1,153 @@
+"""
+Mask repair and the quality gate on what actually reaches Meshy.
+
+## The failure this exists for
+
+On a real specimen — a fish in a limestone slab — SAM3 returned a mask with **512
+interior holes covering 157,000 pixels**. It segmented the fossil-bearing half
+cleanly and then shredded the plain matrix half, because bare limestone has almost
+no texture gradient for it to hold on to. The background remover's alpha matte for
+the same photograph had **zero** holes.
+
+Meshy reconstructs what it is given. A cutout riddled with holes becomes a mesh
+riddled with holes — and it does not fail, it succeeds, which is worse. Nobody sees
+it until they orbit the model.
+
+## Why hole-filling alone is not the fix
+
+Measured on that mask:
+
+    RMBG alpha (reference)        44.5% coverage,   0 holes
+    SAM3 raw                      36.7% coverage, 512 holes
+    SAM3 + fill                   38.0% coverage,   0 holes   IoU 0.854
+    SAM3 + close(61) + fill       39.1% coverage,   0 holes   IoU 0.879
+
+Filling closes the enclosed holes, but much of the speckle is connected to the
+outside — SAM3 ate into the boundary rather than perforating the interior. Even a
+61-pixel closing cannot recover ground SAM3 never claimed. Morphology repairs
+noise; it cannot invent a segmentation.
+
+## So: repair, then judge
+
+`repair_mask` fills and closes, because a fossil slab is a solid object and an
+interior hole is always an artifact. `assess_mask` then compares the result against
+the background remover's matte, which is the one independent opinion available. If
+SAM3 has dropped a large share of what the matte called specimen, SAM3 has failed
+on this view and the matte is the better input.
+
+The pipeline switches automatically and says so. This is not "trust RMBG over SAM3"
+— SAM3's negative points are what keep the cast shadow out, and it usually wins.
+It is: verify before spending, because a bad cutout produces a confident,
+plausible, wrong mesh.
+"""
+
+from __future__ import annotations
+
+import cv2
+import numpy as np
+
+# If the repaired SAM3 mask retains less than this share of the alpha matte's area,
+# treat it as a failed segmentation rather than a tight one. 0.90 was chosen against
+# the real failure above (0.86 after repair) and a good run (>0.97).
+MIN_ALPHA_RETENTION = 0.90
+
+# Interior holes worth worrying about, as a fraction of mask area.
+MAX_HOLE_FRACTION = 0.02
+
+
+def fill_holes(mask: np.ndarray) -> np.ndarray:
+    """Fill enclosed holes.
+
+    Flood from a corner: whatever the flood cannot reach is enclosed, and on a solid
+    specimen enclosed means artifact. Cheaper and more exact than a large closing,
+    and it cannot round off genuine boundary detail the way a big kernel does.
+    """
+    m = mask.astype(bool)
+    inv = (~m).astype(np.uint8) * 255
+    ff = np.zeros((m.shape[0] + 2, m.shape[1] + 2), np.uint8)
+    cv2.floodFill(inv, ff, (0, 0), 0)
+    return m | (inv > 0)
+
+
+def repair_mask(mask: np.ndarray, close_px: int = 15) -> np.ndarray:
+    """Close small speckle, then fill enclosed holes, then keep the largest blob.
+
+    Order matters: closing first bridges pinholes so the fill has something enclosed
+    to work on, and the largest-component pass last drops any island the closing
+    happened to connect up.
+    """
+    m = mask.astype(np.uint8)
+    if close_px >= 3:
+        k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (close_px | 1, close_px | 1))
+        m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, k)
+    m = fill_holes(m.astype(bool))
+
+    n, labels, stats, _ = cv2.connectedComponentsWithStats(m.astype(np.uint8), 8)
+    if n > 1:
+        biggest = 1 + int(np.argmax(stats[1:, cv2.CC_STAT_AREA]))
+        m = labels == biggest
+    return m
+
+
+def count_holes(mask: np.ndarray) -> tuple[int, int]:
+    """(number of interior holes, total hole area in pixels)."""
+    inv = (~mask.astype(bool)).astype(np.uint8)
+    n, labels, stats, _ = cv2.connectedComponentsWithStats(inv, 8)
+    border = (set(labels[0, :]) | set(labels[-1, :])
+              | set(labels[:, 0]) | set(labels[:, -1]))
+    interior = [i for i in range(1, n) if i not in border]
+    return len(interior), int(sum(stats[i, cv2.CC_STAT_AREA] for i in interior))
+
+
+def assess_mask(sam3_mask: np.ndarray, alpha_mask: np.ndarray | None) -> dict:
+    """
+    Judge a repaired SAM3 mask against the background remover's matte.
+
+    Returns {usable, retention, hole_count, reason}. `usable=False` means SAM3
+    dropped too much of the specimen and the alpha-based cutout should be sent to
+    Meshy instead.
+    """
+    holes, hole_px = count_holes(sam3_mask)
+    area = int(sam3_mask.sum())
+    result = {
+        "hole_count": holes,
+        "hole_pixels": hole_px,
+        "coverage": round(float(sam3_mask.mean()), 4),
+        "retention": None,
+        "usable": True,
+        "reason": "",
+    }
+
+    if area == 0:
+        result.update(usable=False, reason="SAM3 returned an empty mask.")
+        return result
+
+    if alpha_mask is None:
+        if hole_px > MAX_HOLE_FRACTION * area:
+            result.update(usable=False,
+                          reason=f"{holes} interior holes ({hole_px} px) survived repair "
+                                 f"and there is no alpha matte to fall back to.")
+        return result
+
+    alpha = alpha_mask.astype(bool)
+    if alpha.sum() == 0:
+        return result
+
+    # Share of the matte's specimen area that SAM3 kept. Not IoU: SAM3 legitimately
+    # trims the matte (the cast shadow, for one), so extra area outside the matte is
+    # fine. What matters is how much of the specimen it THREW AWAY.
+    retention = float(np.logical_and(sam3_mask, alpha).sum() / alpha.sum())
+    result["retention"] = round(retention, 3)
+
+    if retention < MIN_ALPHA_RETENTION:
+        result.update(
+            usable=False,
+            reason=(
+                f"SAM3 kept only {retention:.0%} of the area the background remover "
+                f"called specimen (threshold {MIN_ALPHA_RETENTION:.0%}). On a slab "
+                f"this usually means it segmented the fossil-bearing part and shed "
+                f"the plain matrix, which has little texture for it to hold. The "
+                f"background-removal cutout is the better input for this view."
+            ),
+        )
+    return result