451 lines
13 KiB
Python
451 lines
13 KiB
Python
from __future__ import annotations
|
||
|
||
"""
|
||
统一的分割模型加载入口。
|
||
|
||
当前支持：
|
||
- SAM (segment-anything)
|
||
- Mask2Former（使用 HuggingFace transformers 的语义分割实现）
|
||
"""
|
||
|
||
from dataclasses import dataclass
|
||
from enum import Enum
|
||
from typing import Callable, List
|
||
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
|
||
# -----------------------------
|
||
# SAM 交互式（SamPredictor），与 AutomaticMaskGenerator 分流缓存
|
||
# -----------------------------
|
||
|
||
_sam_prompt_predictor = None
|
||
|
||
|
||
def get_sam_prompt_predictor():
|
||
"""懒加载 SamPredictor（vit_h），用于点/框提示分割。"""
|
||
global _sam_prompt_predictor
|
||
if _sam_prompt_predictor is not None:
|
||
return _sam_prompt_predictor
|
||
|
||
sam_root = _ensure_sam_on_path()
|
||
ckpt_path = _download_sam_checkpoint_if_needed(sam_root)
|
||
|
||
from segment_anything import sam_model_registry, SamPredictor # type: ignore[import]
|
||
import torch
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
sam = sam_model_registry["vit_h"](checkpoint=str(ckpt_path)).to(device)
|
||
_sam_prompt_predictor = SamPredictor(sam)
|
||
return _sam_prompt_predictor
|
||
|
||
|
||
def run_sam_prompt(
|
||
image_rgb: np.ndarray,
|
||
point_coords: np.ndarray,
|
||
point_labels: np.ndarray,
|
||
box_xyxy: np.ndarray | None = None,
|
||
) -> np.ndarray:
|
||
"""
|
||
使用点提示（必选）与可选矩形框在 RGB 图上分割。
|
||
|
||
参数均为原图像素坐标：
|
||
- point_coords: (N, 2) float
|
||
- point_labels: (N,) int，1=前景点，0=背景
|
||
- box_xyxy: (4,) [x1, y1, x2, y2] 或 None
|
||
|
||
返回：bool 掩膜 (H, W)
|
||
"""
|
||
if image_rgb.dtype != np.uint8:
|
||
image_rgb = np.ascontiguousarray(image_rgb.astype(np.uint8))
|
||
else:
|
||
image_rgb = np.ascontiguousarray(image_rgb)
|
||
|
||
if image_rgb.ndim != 3 or image_rgb.shape[2] != 3:
|
||
raise ValueError(f"image_rgb 期望 HWC RGB uint8，当前 shape={image_rgb.shape}")
|
||
|
||
predictor = get_sam_prompt_predictor()
|
||
predictor.set_image(image_rgb)
|
||
|
||
pc = np.asarray(point_coords, dtype=np.float32)
|
||
pl = np.asarray(point_labels, dtype=np.int64)
|
||
if pc.ndim != 2 or pc.shape[1] != 2:
|
||
raise ValueError("point_coords 应为 Nx2")
|
||
if pl.ndim != 1 or pl.shape[0] != pc.shape[0]:
|
||
raise ValueError("point_labels 长度须与 point_coords 行数一致")
|
||
|
||
box_arg = None
|
||
if box_xyxy is not None:
|
||
b = np.asarray(box_xyxy, dtype=np.float32).reshape(4)
|
||
box_arg = b
|
||
|
||
masks3, scores3, _low3 = predictor.predict(
|
||
point_coords=pc,
|
||
point_labels=pl,
|
||
box=box_arg,
|
||
multimask_output=True,
|
||
)
|
||
m_multi = _union_sam_masks_covering_foreground(masks3, scores3, pc, pl)
|
||
|
||
masks1, scores1, _low1 = predictor.predict(
|
||
point_coords=pc,
|
||
point_labels=pl,
|
||
box=box_arg,
|
||
multimask_output=False,
|
||
)
|
||
m_single = _to_bool_mask(masks1[0])
|
||
|
||
m = _prefer_larger_foreground_mask(m_multi, m_single, pc, pl)
|
||
m = _refine_sam_instance_mask(m)
|
||
return m
|
||
|
||
|
||
def _mask_contains_foreground_points(
|
||
mask_bool: np.ndarray,
|
||
point_coords: np.ndarray,
|
||
point_labels: np.ndarray,
|
||
*,
|
||
dilate_for_hit: int = 0,
|
||
) -> bool:
|
||
"""dilate_for_hit>0 时：前景点落在膨胀后的掩膜内即算覆盖（缓解边缘亚像素偏差）。"""
|
||
m = np.asarray(mask_bool, dtype=bool)
|
||
if dilate_for_hit > 0:
|
||
try:
|
||
from scipy import ndimage as ndi # type: ignore[import]
|
||
|
||
m = ndi.binary_dilation(m, structure=np.ones((3, 3), dtype=bool), iterations=dilate_for_hit)
|
||
except Exception:
|
||
pass
|
||
for (x, y), lab in zip(point_coords, point_labels):
|
||
if int(lab) != 1:
|
||
continue
|
||
xi = int(round(float(x)))
|
||
yi = int(round(float(y)))
|
||
if yi < 0 or xi < 0 or yi >= m.shape[0] or xi >= m.shape[1]:
|
||
return False
|
||
if not bool(m[yi, xi]):
|
||
return False
|
||
return True
|
||
|
||
|
||
def _to_bool_mask(m: np.ndarray) -> np.ndarray:
|
||
if m.dtype == np.bool_:
|
||
return m
|
||
return m > 0.5
|
||
|
||
|
||
def _union_sam_masks_covering_foreground(
|
||
masks: np.ndarray,
|
||
scores: np.ndarray,
|
||
point_coords: np.ndarray,
|
||
point_labels: np.ndarray,
|
||
) -> np.ndarray:
|
||
"""
|
||
对「包含全部前景点」的候选掩膜做逻辑并，通常比单选一片更接近完整物体（船身+棚+人物等）。
|
||
若并集过大（误吞整图）则退回「有效候选中面积最大」的一片。
|
||
"""
|
||
n = int(masks.shape[0])
|
||
valid: List[np.ndarray] = []
|
||
for i in range(n):
|
||
mb = _to_bool_mask(masks[i])
|
||
if _mask_contains_foreground_points(mb, point_coords, point_labels, dilate_for_hit=2):
|
||
valid.append(mb)
|
||
if not valid:
|
||
best = int(np.argmax(scores))
|
||
return _to_bool_mask(masks[best])
|
||
if len(valid) == 1:
|
||
return valid[0]
|
||
u = np.zeros_like(valid[0], dtype=bool)
|
||
for v in valid:
|
||
u |= v
|
||
area_u = int(np.count_nonzero(u))
|
||
h, w = u.shape[:2]
|
||
if area_u > int(0.62 * h * w):
|
||
return max(valid, key=lambda x: int(np.count_nonzero(x)))
|
||
return u
|
||
|
||
|
||
def _prefer_larger_foreground_mask(
|
||
a: np.ndarray,
|
||
b: np.ndarray,
|
||
point_coords: np.ndarray,
|
||
point_labels: np.ndarray,
|
||
) -> np.ndarray:
|
||
"""在「多掩膜并集」与「单掩膜」之间选：均含前景点时优先面积明显更大者（更接近整物体）。"""
|
||
a = np.asarray(a, dtype=bool)
|
||
b = np.asarray(b, dtype=bool)
|
||
ca = _mask_contains_foreground_points(a, point_coords, point_labels, dilate_for_hit=2)
|
||
cb = _mask_contains_foreground_points(b, point_coords, point_labels, dilate_for_hit=2)
|
||
aa = int(np.count_nonzero(a))
|
||
ab = int(np.count_nonzero(b))
|
||
|
||
if ca and cb:
|
||
if ab > aa * 1.08:
|
||
return b
|
||
if aa > ab * 1.08:
|
||
return a
|
||
return a if aa >= ab else b
|
||
if ca:
|
||
return a
|
||
if cb:
|
||
return b
|
||
return a if aa >= ab else b
|
||
|
||
|
||
def _refine_sam_instance_mask(mask_bool: np.ndarray) -> np.ndarray:
|
||
"""填孔 + 闭运算连接断裂，再轻微膨胀以包全外轮廓。"""
|
||
m = np.asarray(mask_bool, dtype=bool)
|
||
try:
|
||
from scipy import ndimage as ndi # type: ignore[import]
|
||
|
||
m = ndi.binary_fill_holes(m)
|
||
m = ndi.binary_closing(m, structure=np.ones((5, 5), dtype=bool))
|
||
m = ndi.binary_closing(m, structure=np.ones((3, 3), dtype=bool))
|
||
m = ndi.binary_dilation(m, structure=np.ones((3, 3), dtype=bool), iterations=1)
|
||
return m
|
||
except Exception:
|
||
pass
|
||
try:
|
||
import cv2 # type: ignore[import]
|
||
|
||
u8 = (m.astype(np.uint8) * 255)
|
||
k5 = np.ones((5, 5), np.uint8)
|
||
u8 = cv2.morphologyEx(u8, cv2.MORPH_CLOSE, k5)
|
||
u8 = cv2.dilate(u8, np.ones((3, 3), np.uint8), iterations=1)
|
||
return u8 > 127
|
||
except Exception:
|
||
return m
|
||
|
||
|
||
def expand_mask(mask_bool: np.ndarray, expand_px: int = 0) -> np.ndarray:
|
||
"""
|
||
轻微扩大二值 mask（用于消除边界过软/过细带来的“漏边”观感）。
|
||
|
||
- expand_px <= 0: 原样返回
|
||
- expand_px > 0: 做 dilation（膨胀）expand_px 像素级别
|
||
"""
|
||
mask_bool = np.asarray(mask_bool, dtype=bool)
|
||
if expand_px <= 0:
|
||
return mask_bool
|
||
|
||
# 优先用 OpenCV（通常更快）
|
||
try:
|
||
import cv2 # type: ignore[import]
|
||
|
||
u8 = (mask_bool.astype(np.uint8) * 255)
|
||
k = int(expand_px) * 2 + 1
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
|
||
out = cv2.dilate(u8, kernel, iterations=1)
|
||
return out > 127
|
||
except Exception:
|
||
pass
|
||
|
||
# 回退到 scipy（若存在）
|
||
try:
|
||
from scipy.ndimage import binary_dilation # type: ignore[import]
|
||
|
||
# 使用近似圆形结构元素
|
||
r = int(expand_px)
|
||
yy, xx = np.ogrid[-r : r + 1, -r : r + 1]
|
||
selem = (xx * xx + yy * yy) <= (r * r)
|
||
return binary_dilation(mask_bool, structure=selem)
|
||
except Exception:
|
||
# 再回退：无依赖时就不扩大（保持可用性）
|
||
return mask_bool
|
||
|
||
|
||
def mask_to_contour_xy(
|
||
mask_bool: np.ndarray,
|
||
epsilon_px: float = 2.0,
|
||
) -> List[List[float]]:
|
||
"""
|
||
从二值掩膜提取最大外轮廓，并用 Douglas-Peucker 简化。
|
||
返回 [[x, y], ...]（裁剪图坐标系）。
|
||
"""
|
||
u8 = (np.asarray(mask_bool, dtype=np.uint8) * 255).astype(np.uint8)
|
||
try:
|
||
import cv2 # type: ignore[import]
|
||
|
||
contours, _h = cv2.findContours(u8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
except Exception:
|
||
contours = None
|
||
|
||
if not contours:
|
||
# 无 OpenCV 时的极简回退：外接矩形
|
||
ys, xs = np.nonzero(mask_bool)
|
||
if ys.size == 0:
|
||
return []
|
||
x0, x1 = int(xs.min()), int(xs.max())
|
||
y0, y1 = int(ys.min()), int(ys.max())
|
||
return [
|
||
[float(x0), float(y0)],
|
||
[float(x1), float(y0)],
|
||
[float(x1), float(y1)],
|
||
[float(x0), float(y1)],
|
||
]
|
||
|
||
cnt = max(contours, key=cv2.contourArea)
|
||
if cv2.contourArea(cnt) < 1.0:
|
||
return []
|
||
peri = cv2.arcLength(cnt, True)
|
||
eps = max(epsilon_px, 0.0025 * peri)
|
||
approx = cv2.approxPolyDP(cnt, eps, True)
|
||
out: List[List[float]] = []
|
||
for p in approx:
|
||
out.append([float(p[0][0]), float(p[0][1])])
|
||
if len(out) >= 3 and (out[0][0] != out[-1][0] or out[0][1] != out[-1][1]):
|
||
pass # 保持开放折线，前端可自行闭合
|
||
return out
|
||
|
||
_THIS_DIR = Path(__file__).resolve().parent
|
||
|
||
|
||
class SegBackend(str, Enum):
|
||
SAM = "sam"
|
||
MASK2FORMER = "mask2former"
|
||
|
||
|
||
@dataclass
|
||
class UnifiedSegConfig:
|
||
backend: SegBackend = SegBackend.SAM
|
||
|
||
|
||
# -----------------------------
|
||
# SAM (Segment Anything)
|
||
# -----------------------------
|
||
|
||
|
||
def _ensure_sam_on_path() -> Path:
|
||
sam_root = _THIS_DIR / "segment-anything"
|
||
if not sam_root.is_dir():
|
||
raise FileNotFoundError(f"未找到 segment-anything 仓库目录: {sam_root}")
|
||
sam_path = str(sam_root)
|
||
if sam_path not in sys.path:
|
||
sys.path.insert(0, sam_path)
|
||
return sam_root
|
||
|
||
|
||
def _download_sam_checkpoint_if_needed(sam_root: Path) -> Path:
|
||
import requests
|
||
|
||
ckpt_dir = sam_root / "checkpoints"
|
||
ckpt_dir.mkdir(parents=True, exist_ok=True)
|
||
ckpt_path = ckpt_dir / "sam_vit_h_4b8939.pth"
|
||
|
||
if ckpt_path.is_file():
|
||
return ckpt_path
|
||
|
||
url = (
|
||
"https://dl.fbaipublicfiles.com/segment_anything/"
|
||
"sam_vit_h_4b8939.pth"
|
||
)
|
||
print(f"自动下载 SAM 权重:\n {url}\n -> {ckpt_path}")
|
||
|
||
resp = requests.get(url, stream=True)
|
||
resp.raise_for_status()
|
||
|
||
total = int(resp.headers.get("content-length", "0") or "0")
|
||
downloaded = 0
|
||
chunk_size = 1024 * 1024
|
||
|
||
with ckpt_path.open("wb") as f:
|
||
for chunk in resp.iter_content(chunk_size=chunk_size):
|
||
if not chunk:
|
||
continue
|
||
f.write(chunk)
|
||
downloaded += len(chunk)
|
||
if total > 0:
|
||
done = int(50 * downloaded / total)
|
||
print(
|
||
"\r[{}{}] {:.1f}%".format(
|
||
"#" * done,
|
||
"." * (50 - done),
|
||
downloaded * 100 / total,
|
||
),
|
||
end="",
|
||
)
|
||
print("\nSAM 权重下载完成。")
|
||
return ckpt_path
|
||
|
||
|
||
def _make_sam_predictor() -> Callable[[np.ndarray], np.ndarray]:
|
||
"""
|
||
返回一个分割函数：
|
||
- 输入：RGB uint8 图像 (H, W, 3)
|
||
- 输出：语义标签图 (H, W)，每个目标一个 int id（从 1 开始）
|
||
"""
|
||
sam_root = _ensure_sam_on_path()
|
||
ckpt_path = _download_sam_checkpoint_if_needed(sam_root)
|
||
|
||
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator # type: ignore[import]
|
||
import torch
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
|
||
sam = sam_model_registry["vit_h"](
|
||
checkpoint=str(ckpt_path),
|
||
).to(device)
|
||
|
||
mask_generator = SamAutomaticMaskGenerator(sam)
|
||
|
||
def _predict(image_rgb: np.ndarray) -> np.ndarray:
|
||
if image_rgb.dtype != np.uint8:
|
||
image_rgb_u8 = image_rgb.astype("uint8")
|
||
else:
|
||
image_rgb_u8 = image_rgb
|
||
|
||
masks = mask_generator.generate(image_rgb_u8)
|
||
h, w, _ = image_rgb_u8.shape
|
||
label_map = np.zeros((h, w), dtype="int32")
|
||
|
||
for idx, m in enumerate(masks, start=1):
|
||
seg = m.get("segmentation")
|
||
if seg is None:
|
||
continue
|
||
label_map[seg.astype(bool)] = idx
|
||
|
||
return label_map
|
||
|
||
return _predict
|
||
|
||
|
||
# -----------------------------
|
||
# Mask2Former (占位)
|
||
# -----------------------------
|
||
|
||
|
||
def _make_mask2former_predictor() -> Callable[[np.ndarray], np.ndarray]:
|
||
from .mask2former_loader import build_mask2former_hf_predictor
|
||
|
||
predictor, _ = build_mask2former_hf_predictor()
|
||
return predictor
|
||
|
||
|
||
# -----------------------------
|
||
# 统一构建函数
|
||
# -----------------------------
|
||
|
||
|
||
def build_seg_predictor(
|
||
cfg: UnifiedSegConfig | None = None,
|
||
) -> tuple[Callable[[np.ndarray], np.ndarray], SegBackend]:
|
||
"""
|
||
统一构建分割预测函数。
|
||
|
||
返回：
|
||
- predictor(image_rgb: np.ndarray[H, W, 3], uint8) -> np.ndarray[H, W], int32
|
||
- 实际使用的 backend
|
||
"""
|
||
cfg = cfg or UnifiedSegConfig()
|
||
|
||
if cfg.backend == SegBackend.SAM:
|
||
return _make_sam_predictor(), SegBackend.SAM
|
||
|
||
if cfg.backend == SegBackend.MASK2FORMER:
|
||
return _make_mask2former_predictor(), SegBackend.MASK2FORMER
|
||
|
||
raise ValueError(f"不支持的分割后端: {cfg.backend}")
|
||
|