Files
Pre_Seg_Server/backend/tests/test_sam3_engine.py
admin 8a9247075e feat: 完善 AI 分割与工作区标注闭环
功能增加:

- 将视频导入和生成帧拆成两个明确动作,项目库生成帧时选择 FPS,工作区不再自动触发拆帧。

- 为工作区新增调整多边形工具,支持选中 mask、拖动顶点、边中点插点、双击边界按位置插点,并保留多 polygon 子区域编辑。

- 打通 AI 页 SAM2/SAM3 结果到工作区的联动,生成 mask 后自动选中,可在右侧分类树换标签,并推送到工作区继续编辑。

- 增强 Dashboard WebSocket 连接状态与心跳,使用真实 onopen/onclose/onerror 状态驱动前端显示。

- 完善 SAM3 external worker 适配,支持 box prompt、semantic 请求级阈值和 video tracker 路径。

bugfix:

- 修复 SAM2 文本语义误走自动分割的问题,改为提示使用点提示或切换 SAM3。

- 修复 SAM2 多候选重叠显示的问题,点提示和 auto fallback 默认只采用最高分候选。

- 修复 SAM2 反向点看起来无效的问题,带负点时启用背景过滤,过滤为空时移除旧候选。

- 修复 SAM3 单个 2D mask 结果无法转 polygon、低阈值 semantic 返回被默认阈值吞掉的问题。

- 修复 AI 页 mask 未选中导致分类树无法修改 SAM2 结果标签的问题。

测试和文档:

- 补充 CanvasArea、AISegmentation、ProjectLibrary、VideoWorkspace、Dashboard、websocket 和 SAM engine/API 测试。

- 新增 backend/tests/test_sam2_engine.py,覆盖 SAM2 单候选请求和 auto fallback 行为。

- 更新 README、AGENTS 和 doc 需求/设计/接口/测试矩阵,按当前实现冻结功能状态。
2026-05-01 21:50:17 +08:00

294 lines
10 KiB
Python

import json
from pathlib import Path
import numpy as np
from services.sam3_engine import SAM3Engine
from services.sam3_external_worker import _prediction_to_response, _to_numpy
class _Completed:
def __init__(self, returncode=0, stdout="", stderr=""):
self.returncode = returncode
self.stdout = stdout
self.stderr = stderr
def _external_settings(monkeypatch, python_path: Path):
checkpoint_path = python_path.with_name("sam3.pt")
checkpoint_path.write_bytes(b"checkpoint")
python_path.write_text("#!/usr/bin/env python\n", encoding="utf-8")
python_path.chmod(0o755)
monkeypatch.setattr("services.sam3_engine.SAM3_PACKAGE_AVAILABLE", False)
monkeypatch.setattr("services.sam3_engine.TORCH_AVAILABLE", False)
monkeypatch.setattr("services.sam3_engine.settings.sam3_external_enabled", True)
monkeypatch.setattr("services.sam3_engine.settings.sam3_external_python", str(python_path))
monkeypatch.setattr("services.sam3_engine.settings.sam3_timeout_seconds", 10)
monkeypatch.setattr("services.sam3_engine.settings.sam3_status_cache_seconds", 30)
monkeypatch.setattr("services.sam3_engine.settings.sam3_confidence_threshold", 0.4)
monkeypatch.setattr("services.sam3_engine.settings.sam3_checkpoint_path", str(checkpoint_path))
def test_sam3_status_reports_external_runtime_ready(tmp_path, monkeypatch):
_external_settings(monkeypatch, tmp_path / "python")
def fake_run(args, **_kwargs):
assert "--status" in args
assert _kwargs["env"]["SAM3_CHECKPOINT_PATH"].endswith("sam3.pt")
return _Completed(stdout=json.dumps({
"available": True,
"package_available": True,
"checkpoint_access": True,
"checkpoint_path": _kwargs["env"]["SAM3_CHECKPOINT_PATH"],
"python_ok": True,
"torch_ok": True,
"cuda_available": True,
"device": "cuda",
"message": "ready",
}))
monkeypatch.setattr("services.sam3_engine.subprocess.run", fake_run)
status = SAM3Engine().status()
assert status["available"] is True
assert status["external_available"] is True
assert status["package_available"] is True
assert status["python_ok"] is True
assert status["checkpoint_exists"] is True
assert status["checkpoint_path"].endswith("sam3.pt")
assert status["supports"] == ["semantic", "box", "video_track"]
assert status["message"] == "SAM 3 external runtime is ready; local checkpoint will load in the helper process on inference."
def test_sam3_predict_semantic_uses_external_worker(tmp_path, monkeypatch):
_external_settings(monkeypatch, tmp_path / "python")
calls = []
def fake_run(args, **_kwargs):
calls.append(args)
if "--status" in args:
return _Completed(stdout=json.dumps({
"available": True,
"package_available": True,
"checkpoint_access": True,
"python_ok": True,
"torch_ok": True,
"cuda_available": True,
"device": "cuda",
"message": "ready",
}))
request_path = Path(args[-1])
request = json.loads(request_path.read_text(encoding="utf-8"))
assert request["text"] == "vessel"
assert request["confidence_threshold"] == 0.4
assert request["checkpoint_path"].endswith("sam3.pt")
assert Path(request["image_path"]).exists()
return _Completed(stdout=json.dumps({
"polygons": [[[0.1, 0.1], [0.9, 0.1], [0.9, 0.9]]],
"scores": [0.91],
}))
monkeypatch.setattr("services.sam3_engine.subprocess.run", fake_run)
polygons, scores = SAM3Engine().predict_semantic(np.zeros((8, 8, 3), dtype=np.uint8), " vessel ")
assert polygons == [[[0.1, 0.1], [0.9, 0.1], [0.9, 0.9]]]
assert scores == [0.91]
assert any("--request" in args for args in calls)
def test_sam3_predict_semantic_allows_request_threshold_override(tmp_path, monkeypatch):
_external_settings(monkeypatch, tmp_path / "python")
def fake_run(args, **_kwargs):
if "--status" in args:
return _Completed(stdout=json.dumps({
"available": True,
"package_available": True,
"checkpoint_access": True,
"python_ok": True,
"torch_ok": True,
"cuda_available": True,
"device": "cuda",
"message": "ready",
}))
request_path = Path(args[-1])
request = json.loads(request_path.read_text(encoding="utf-8"))
assert request["confidence_threshold"] == 0.05
return _Completed(stdout=json.dumps({
"polygons": [[[0.2, 0.2], [0.6, 0.2], [0.6, 0.6]]],
"scores": [0.07],
}))
monkeypatch.setattr("services.sam3_engine.subprocess.run", fake_run)
polygons, scores = SAM3Engine().predict_semantic(
np.zeros((8, 8, 3), dtype=np.uint8),
"surgical scene",
confidence_threshold=0.05,
)
assert len(polygons) == 1
assert scores == [0.07]
def test_sam3_predict_box_uses_external_worker(tmp_path, monkeypatch):
_external_settings(monkeypatch, tmp_path / "python")
def fake_run(args, **_kwargs):
if "--status" in args:
return _Completed(stdout=json.dumps({
"available": True,
"package_available": True,
"checkpoint_access": True,
"python_ok": True,
"torch_ok": True,
"cuda_available": True,
"device": "cuda",
"message": "ready",
}))
request_path = Path(args[-1])
request = json.loads(request_path.read_text(encoding="utf-8"))
assert request["prompt_type"] == "box"
assert request["box"] == [0.1, 0.2, 0.7, 0.8]
assert request["text"] == ""
return _Completed(stdout=json.dumps({
"polygons": [[[0.1, 0.2], [0.7, 0.2], [0.7, 0.8]]],
"scores": [0.88],
}))
monkeypatch.setattr("services.sam3_engine.subprocess.run", fake_run)
polygons, scores = SAM3Engine().predict_box(
np.zeros((8, 8, 3), dtype=np.uint8),
[0.1, 0.2, 0.7, 0.8],
)
assert polygons == [[[0.1, 0.2], [0.7, 0.2], [0.7, 0.8]]]
assert scores == [0.88]
def test_sam3_propagate_video_uses_external_worker(tmp_path, monkeypatch):
_external_settings(monkeypatch, tmp_path / "python")
frame_dir = tmp_path / "frames"
frame_dir.mkdir()
frame_paths = []
for index in range(2):
frame_path = frame_dir / f"frame_{index:06d}.jpg"
frame_path.write_bytes(b"jpeg")
frame_paths.append(str(frame_path))
def fake_run(args, **_kwargs):
if "--status" in args:
return _Completed(stdout=json.dumps({
"available": True,
"package_available": True,
"checkpoint_access": True,
"python_ok": True,
"torch_ok": True,
"cuda_available": True,
"device": "cuda",
"message": "ready",
}))
request_path = Path(args[-1])
request = json.loads(request_path.read_text(encoding="utf-8"))
assert request["prompt_type"] == "video_track"
assert request["frame_dir"] == str(frame_dir)
assert request["source_frame_index"] == 0
assert request["direction"] == "forward"
assert request["max_frames"] == 2
assert request["seed"]["bbox"] == [0.1, 0.1, 0.2, 0.2]
return _Completed(stdout=json.dumps({
"frames": [
{
"frame_index": 1,
"polygons": [[[0.2, 0.2], [0.4, 0.2], [0.4, 0.4]]],
"scores": [0.7],
"object_ids": [1],
}
]
}))
monkeypatch.setattr("services.sam3_engine.subprocess.run", fake_run)
frames = SAM3Engine().propagate_video(
frame_paths,
0,
{"bbox": [0.1, 0.1, 0.2, 0.2]},
direction="forward",
max_frames=2,
)
assert frames[0]["frame_index"] == 1
assert frames[0]["scores"] == [0.7]
def test_sam3_predict_semantic_reports_external_errors(tmp_path, monkeypatch):
_external_settings(monkeypatch, tmp_path / "python")
def fake_run(args, **_kwargs):
if "--status" in args:
return _Completed(stdout=json.dumps({
"available": True,
"package_available": True,
"checkpoint_access": True,
"python_ok": True,
"torch_ok": True,
"cuda_available": True,
"device": "cuda",
"message": "ready",
}))
return _Completed(returncode=1, stderr=json.dumps({"error": "HF access denied"}))
monkeypatch.setattr("services.sam3_engine.subprocess.run", fake_run)
try:
SAM3Engine().predict_semantic(np.zeros((8, 8, 3), dtype=np.uint8), "vessel")
except RuntimeError as exc:
assert "HF access denied" in str(exc)
else:
raise AssertionError("Expected SAM 3 external inference failure.")
def test_sam3_worker_casts_floating_tensors_before_numpy():
class FakeTensor:
def __init__(self):
self.float_called = False
def detach(self):
return self
def is_floating_point(self):
return True
def float(self):
self.float_called = True
return self
def cpu(self):
return self
def numpy(self):
return np.array([1.0], dtype=np.float32)
tensor = FakeTensor()
result = _to_numpy(tensor)
assert tensor.float_called is True
assert result.tolist() == [1.0]
def test_sam3_worker_converts_single_2d_mask_to_polygon():
mask = np.zeros((12, 12), dtype=np.uint8)
mask[2:10, 3:9] = 1
result = _prediction_to_response({
"masks": mask,
"scores": np.array([0.82], dtype=np.float32),
})
assert len(result["polygons"]) == 1
assert result["scores"] == [0.8199999928474426]