142 lines
4.5 KiB
Python
142 lines
4.5 KiB
Python
"""Shared helpers for the SAM-family adapters."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
def mask_to_polygon(
|
|
mask: np.ndarray,
|
|
*,
|
|
min_points: int = 4,
|
|
max_points: int = 10,
|
|
smooth: bool = True,
|
|
) -> list[tuple[float, float]]:
|
|
"""Largest external contour, simplified to a vertex count in
|
|
[min_points, max_points].
|
|
|
|
Pipeline:
|
|
1. Squeeze + binarise the mask.
|
|
2. Optional morphological close+open + Gaussian smoothing — this is
|
|
what turns staircase-jagged SAM mask edges into curves the polygon
|
|
can actually trace cleanly.
|
|
3. Pull the full contour (`CHAIN_APPROX_NONE`) so the simplifier has
|
|
every pixel-edge sample to choose from.
|
|
4. Binary-search Douglas-Peucker epsilon to hit a vertex count in
|
|
[min_points, max_points], biased toward more points (more detail).
|
|
|
|
Why this matters: the user has to drag every vertex by hand in the
|
|
annotation UI. Raw `findContours` output is often hundreds of points
|
|
along a single edge — unusable. Without step 2, even a 15-point output
|
|
looks choppy because the chosen vertices sit on stair-step pixels.
|
|
"""
|
|
m = np.asarray(mask)
|
|
# Tolerate stray leading singleton dims (some processors return
|
|
# [1, H, W] or [1, 1, H, W]).
|
|
m = np.squeeze(m)
|
|
while m.ndim > 2:
|
|
m = m[0]
|
|
if m.ndim != 2 or m.size == 0:
|
|
return []
|
|
|
|
m = m.astype(np.uint8)
|
|
if m.max() == 1:
|
|
m = m * 255
|
|
|
|
if smooth:
|
|
m = _smooth_mask(m)
|
|
|
|
contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
|
if not contours:
|
|
return []
|
|
cnt = max(contours, key=cv2.contourArea)
|
|
if len(cnt) < 3:
|
|
return []
|
|
|
|
simplified = _simplify_to_band(cnt, min_points=min_points, max_points=max_points)
|
|
if simplified is None or len(simplified) < 3:
|
|
return []
|
|
pts = simplified.reshape(-1, 2)
|
|
return [(float(x), float(y)) for x, y in pts]
|
|
|
|
|
|
def _smooth_mask(mask_u8: np.ndarray) -> np.ndarray:
|
|
"""Clean up a binary mask before contour extraction.
|
|
|
|
Kernel size scales mildly with the mask's bounding box so a tiny object
|
|
doesn't get erased and a huge object gets enough smoothing to matter.
|
|
"""
|
|
h, w = mask_u8.shape[:2]
|
|
# Find the actual blob size (not the canvas size) for kernel sizing.
|
|
ys, xs = np.where(mask_u8 > 0)
|
|
if xs.size == 0:
|
|
return mask_u8
|
|
bw = xs.max() - xs.min() + 1
|
|
bh = ys.max() - ys.min() + 1
|
|
blob_diag = max(1.0, float(np.hypot(bw, bh)))
|
|
# Kernel ~1.5% of blob diagonal, clamped to a usable range.
|
|
k = int(round(blob_diag * 0.015))
|
|
k = max(3, min(k | 1, 11)) # odd, between 3 and 11
|
|
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
|
|
# Close fills tiny holes / single-pixel notches; open removes specks.
|
|
m = cv2.morphologyEx(mask_u8, cv2.MORPH_CLOSE, kernel)
|
|
m = cv2.morphologyEx(m, cv2.MORPH_OPEN, kernel)
|
|
|
|
# A light Gaussian + re-threshold smooths the edge by a sub-pixel amount,
|
|
# which makes the contour follow curves cleanly instead of staircasing.
|
|
blur_k = max(3, min(k | 1, 9))
|
|
m = cv2.GaussianBlur(m, (blur_k, blur_k), 0)
|
|
_, m = cv2.threshold(m, 127, 255, cv2.THRESH_BINARY)
|
|
return m
|
|
|
|
|
|
def _simplify_to_band(
|
|
contour: np.ndarray,
|
|
*,
|
|
min_points: int,
|
|
max_points: int,
|
|
iterations: int = 24,
|
|
) -> np.ndarray | None:
|
|
"""Binary-search epsilon for `cv2.approxPolyDP` so vertex count lands in
|
|
[min_points, max_points], biased toward the larger end of the band.
|
|
|
|
Strategy:
|
|
- epsilon=0 → original (many) points
|
|
- epsilon=arc → degenerates to a single point
|
|
- smaller epsilon ⇒ more points; larger ⇒ fewer.
|
|
|
|
We track the in-band result with the highest vertex count seen
|
|
(most-detailed pick within the cap) and return it.
|
|
"""
|
|
arc = cv2.arcLength(contour, True)
|
|
if arc == 0:
|
|
return contour
|
|
if len(contour) <= max_points:
|
|
return contour if len(contour) >= min_points else None
|
|
|
|
lo, hi = 0.0, arc
|
|
best: np.ndarray | None = None
|
|
best_n = -1
|
|
|
|
for _ in range(iterations):
|
|
mid = 0.5 * (lo + hi)
|
|
approx = cv2.approxPolyDP(contour, mid, True)
|
|
n = len(approx)
|
|
if n > max_points:
|
|
lo = mid
|
|
elif n < min_points:
|
|
hi = mid
|
|
else:
|
|
if n > best_n:
|
|
best = approx
|
|
best_n = n
|
|
hi = mid
|
|
if hi - lo < 1e-6 * arc:
|
|
break
|
|
|
|
if best is not None:
|
|
return best
|
|
return cv2.approxPolyDP(contour, lo, True)
|