Files
Pre_Seg_Server/backend/tests/test_sam2_engine.py
admin 29a1a87e52 feat: 完善 SAM2.1 模型选择与标注工作流
- 后端 SAM2 引擎新增 sam2.1_hiera_tiny、sam2.1_hiera_small、sam2.1_hiera_base_plus、sam2.1_hiera_large 四个变体定义,并按变体维护 checkpoint/config、image predictor、video predictor、加载状态、错误信息和真实状态回报。

- 后端 SAM registry 仅暴露当前产品启用的 SAM2.1 变体,保留 sam2 作为 tiny 兼容别名,拒绝 sam3 产品入口,并把 point、box、interactive、auto、propagate 都分发到所选 SAM2.1 变体。

- 后端默认配置和下载脚本切换到 SAM2.1 checkpoint 命名,支持 legacy SAM2 checkpoint fallback,并在状态消息中标出 fallback 使用情况。

- 前端全局 AI 模型状态新增 SAM2.1 tiny/small/base+/large 类型和默认 tiny,API 请求默认携带 sam2.1_hiera_tiny,AI 页面提供模型变体选择和所选模型状态展示。

- AI 智能分割页移除当前产品不使用的 SAM3/文本提示入口,保留正向点、反向点、框选和参数开关;AI 页只展示本页生成的候选 mask,并支持遮罩清晰度调节、候选 mask 上继续加正/反点、清空本页候选、推送到工作区编辑。

- 工作区和 Canvas 补强 SAM2 交互式细化链路:框选后正/反点继续细化同一个候选 mask,反向点请求启用背景过滤,空结果会移除被否定候选;AI 推送到工作区后保留选中态和未保存 draft mask。

- 工作区标注保存闭环补强:未保存 mask 可归档保存,dirty saved mask 可更新,保存后用后端 saved annotation 替换已提交 draft,清空/删除已保存 mask 时同步后端删除。

- Dashboard 任务进度区改为展示 queued、running、success、failed、cancelled 最近任务,处理中统计只计算 queued/running,并保留近期完成记录。

- 时间轴在顶部时间进度条和底部缩略图导航轴之间新增已编辑帧标记带,基于当前项目帧内 masks 标出已有编辑/标注的帧,并支持点击标记跳转。

- 前端测试覆盖 SAM2.1 变体选择、模型状态徽标、AI 页候选隔离、遮罩透明度、候选上追加正/反点、推送工作区保留选择、Canvas 交互式细化、VideoWorkspace 传播/保存、Dashboard 进度和时间轴已编辑帧标记。

- 后端测试覆盖 SAM2.1 变体状态、sam2 alias 兼容、sam3 禁用、semantic 禁用、传播标注保存、Dashboard 最近任务状态和 SAM3 历史测试跳过说明。

- README、AGENTS 和 doc 文档同步当前真实进度,更新 SAM2.1 变体、SAM3 禁用、接口契约、设计冻结、需求冻结、前端元素审计、实施计划、FastAPI docs 说明和测试矩阵。
2026-05-01 23:39:53 +08:00

81 lines
2.5 KiB
Python

import numpy as np
from services.sam2_engine import DEFAULT_SAM2_MODEL_ID, SAM2Engine
class _FakePredictor:
def __init__(self, masks, scores):
self.masks = masks
self.scores = scores
self.calls = []
def set_image(self, _image):
pass
def predict(self, **kwargs):
self.calls.append(kwargs)
return self.masks, self.scores, None
def _mask(offset=0):
mask = np.zeros((32, 32), dtype=np.uint8)
mask[4 + offset:20 + offset, 5 + offset:22 + offset] = 1
return mask
def _ready_engine(monkeypatch, predictor):
monkeypatch.setattr("services.sam2_engine.SAM2_AVAILABLE", True)
engine = SAM2Engine()
engine._model_loaded[DEFAULT_SAM2_MODEL_ID] = True
engine._predictors[DEFAULT_SAM2_MODEL_ID] = predictor
return engine
def test_sam2_point_prediction_requests_single_best_mask(monkeypatch):
predictor = _FakePredictor(
np.array([_mask()], dtype=np.uint8),
np.array([0.92], dtype=np.float32),
)
engine = _ready_engine(monkeypatch, predictor)
polygons, scores = engine.predict_points(
DEFAULT_SAM2_MODEL_ID,
np.zeros((32, 32, 3), dtype=np.uint8),
[[0.5, 0.5]],
[1],
)
assert predictor.calls[0]["multimask_output"] is False
assert len(polygons) == 1
assert scores == [0.9200000166893005]
def test_sam2_auto_prediction_keeps_single_best_mask(monkeypatch):
predictor = _FakePredictor(
np.array([_mask(0), _mask(2), _mask(4)], dtype=np.uint8),
np.array([0.8, 0.7, 0.6], dtype=np.float32),
)
engine = _ready_engine(monkeypatch, predictor)
polygons, scores = engine.predict_auto(DEFAULT_SAM2_MODEL_ID, np.zeros((32, 32, 3), dtype=np.uint8))
assert predictor.calls[0]["multimask_output"] is False
assert len(polygons) == 1
assert scores == [0.800000011920929]
def test_sam2_status_exposes_selectable_variants(monkeypatch, tmp_path):
checkpoint = tmp_path / "sam2.1_hiera_small.pt"
checkpoint.write_bytes(b"model")
monkeypatch.setattr("services.sam2_engine.settings.sam_model_path", str(tmp_path / "sam2.1_hiera_tiny.pt"))
engine = SAM2Engine()
status = engine.status("sam2.1_hiera_small")
assert engine.normalize_model_id("sam2") == DEFAULT_SAM2_MODEL_ID
assert "sam2.1_hiera_small" in engine.variant_ids()
assert status["id"] == "sam2.1_hiera_small"
assert status["label"] == "SAM 2.1 Small"
assert status["checkpoint_exists"] is True
assert status["checkpoint_path"].endswith("sam2.1_hiera_small.pt")