admin/Pre_Seg_Server
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

feat: 完善 AI 分割与工作区标注闭环

Browse Source 
功能增加:

- 将视频导入和生成帧拆成两个明确动作，项目库生成帧时选择 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 需求/设计/接口/测试矩阵，按当前实现冻结功能状态。
This commit is contained in:
admin
2026-05-01 21:50:17 +08:00
parent 5ab4602535
commit 8a9247075e
31 changed files with 920 additions and 216 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
backend/routers/ai.py
View File

@@ -340,7 +340,21 @@ def predict(payload: PredictRequest, db: Session = Depends(get_db)) -> dict:
elif prompt_type == "semantic":
text = payload.prompt_data if isinstance(payload.prompt_data, str) else ""
polygons, scores = sam_registry.predict_semantic(payload.model, image, text)
min_score = options.get("min_score")
confidence_threshold = None
if min_score is not None:
try:
parsed_min_score = float(min_score)
if parsed_min_score > 0:
confidence_threshold = parsed_min_score
except (TypeError, ValueError):
confidence_threshold = None
polygons, scores = sam_registry.predict_semantic(
payload.model,
image,
text,
confidence_threshold=confidence_threshold,
)
else:
raise HTTPException(status_code=400, detail=f"Unsupported prompt_type: {prompt_type}")
@@ -352,6 +366,13 @@ def predict(payload: PredictRequest, db: Session = Depends(get_db)) -> dict:
raise HTTPException(status_code=400, detail=str(exc)) from exc
polygons, scores = _filter_predictions(polygons, scores, options, negative_points)
logger.info(
"AI predict completed model=%s prompt_type=%s frame_id=%s polygons=%d",
payload.model or "default",
prompt_type,
payload.image_id,
len(polygons),
)
return {"polygons": polygons, "scores": scores}

8
backend/services/sam2_engine.py
View File

@@ -207,7 +207,7 @@ class SAM2Engine:
masks, scores, _ = self._predictor.predict(
point_coords=pts,
point_labels=lbls,
multimask_output=True,
multimask_output=False,
)
polygons = []
@@ -335,16 +335,16 @@ class SAM2Engine:
masks, scores, _ = self._predictor.predict(
point_coords=pts,
point_labels=lbls,
multimask_output=True,
multimask_output=False,
)
polygons = []
for m in masks[:3]: # Limit to top 3 masks
for m in masks[:1]:
poly = self._mask_to_polygon(m)
if poly:
polygons.append(poly)
return polygons, scores[:3].tolist()
return polygons, scores[:1].tolist()
except Exception as exc: # noqa: BLE001
logger.error("SAM2 auto prediction failed: %s", exc)
return self._dummy_polygons(image.shape[1], image.shape[0]), [0.5]

30
backend/services/sam3_engine.py
View File

@@ -260,6 +260,7 @@ class SAM3Engine:
*,
text: str = "",
box: list[float] | None = None,
confidence_threshold: float | None = None,
) -> tuple[list[list[list[float]]], list[float]]:
status = self._external_status(force=True)
if not status.get("available"):
@@ -279,7 +280,11 @@ class SAM3Engine:
"box": box,
"model_version": settings.sam3_model_version,
"checkpoint_path": self._checkpoint_path(),
"confidence_threshold": settings.sam3_confidence_threshold,
"confidence_threshold": (
confidence_threshold
if confidence_threshold is not None
else settings.sam3_confidence_threshold
),
},
ensure_ascii=False,
),
@@ -312,8 +317,18 @@ class SAM3Engine:
raise RuntimeError(str(payload["error"]))
return payload.get("polygons", []), payload.get("scores", [])
def _predict_semantic_external(self, image: np.ndarray, text: str) -> tuple[list[list[list[float]]], list[float]]:
return self._predict_external(image, "semantic", text=text)
def _predict_semantic_external(
self,
image: np.ndarray,
text: str,
confidence_threshold: float | None = None,
) -> tuple[list[list[list[float]]], list[float]]:
return self._predict_external(
image,
"semantic",
text=text,
confidence_threshold=confidence_threshold,
)
def _predict_box_external(self, image: np.ndarray, box: list[float]) -> tuple[list[list[list[float]]], list[float]]:
return self._predict_external(image, "box", box=box)
@@ -378,11 +393,16 @@ class SAM3Engine:
raise RuntimeError(str(payload["error"]))
return payload.get("frames", [])
def predict_semantic(self, image: np.ndarray, text: str) -> tuple[list[list[list[float]]], list[float]]:
def predict_semantic(
self,
image: np.ndarray,
text: str,
confidence_threshold: float | None = None,
) -> tuple[list[list[list[float]]], list[float]]:
if not text.strip():
raise ValueError("SAM 3 semantic prompt requires non-empty text.")
if not self._can_load() and self._external_status().get("available"):
return self._predict_semantic_external(image, text)
return self._predict_semantic_external(image, text, confidence_threshold=confidence_threshold)
if not self._ensure_ready():
raise RuntimeError(self.status()["message"])

4
backend/services/sam3_external_worker.py
View File

@@ -190,7 +190,9 @@ def _video_outputs_to_response(outputs: dict[str, Any]) -> dict[str, Any]:
def _prediction_to_response(output: dict[str, Any]) -> dict[str, Any]:
masks = _to_numpy(output.get("masks", []))
scores = _to_numpy(output.get("scores", []))
if masks.ndim == 4:
if masks.ndim == 2:
masks = masks[None, :, :]
elif masks.ndim == 4:
masks = masks[:, 0]
elif masks.ndim == 3 and masks.shape[0] == 1:
masks = masks[None, 0]

14
backend/services/sam_registry.py
View File

@@ -83,10 +83,20 @@ class SAMRegistry:
def predict_auto(self, model_id: str | None, image: Any):
return self._ensure_available(model_id).predict_auto(image)
def predict_semantic(self, model_id: str | None, image: Any, text: str):
def predict_semantic(
self,
model_id: str | None,
image: Any,
text: str,
confidence_threshold: float | None = None,
):
model = self.normalize_model_id(model_id)
if model == "sam3":
return self._ensure_available(model).predict_semantic(image, text)
return self._ensure_available(model).predict_semantic(
image,
text,
confidence_threshold=confidence_threshold,
)
return self._ensure_available(model).predict_auto(image)
def propagate_video(

25
backend/tests/test_ai.py
View File

@@ -89,15 +89,25 @@ def test_predict_applies_crop_and_background_filter_options(client, monkeypatch)
def test_predict_box_and_semantic_fallback(client, monkeypatch):
_, frame, _ = _create_project_and_frame(client)
calls = {}
monkeypatch.setattr("routers.ai._load_frame_image", lambda frame: np.zeros((10, 10, 3), dtype=np.uint8))
monkeypatch.setattr("routers.ai.sam_registry.predict_box", lambda model, image, box: (
[[[0.2, 0.2], [0.8, 0.2], [0.8, 0.8]]],
[0.8],
))
monkeypatch.setattr("routers.ai.sam_registry.predict_semantic", lambda model, image, text: (
[[[0.0, 0.0], [1.0, 0.0], [1.0, 1.0]]],
[0.5],
))
def fake_predict_semantic(model, image, text, confidence_threshold=None):
calls["semantic"] = {
"model": model,
"text": text,
"confidence_threshold": confidence_threshold,
}
return (
[[[0.0, 0.0], [1.0, 0.0], [1.0, 1.0]]],
[0.5],
)
monkeypatch.setattr("routers.ai.sam_registry.predict_semantic", fake_predict_semantic)
box_response = client.post("/api/ai/predict", json={
"image_id": frame["id"],
@@ -108,12 +118,19 @@ def test_predict_box_and_semantic_fallback(client, monkeypatch):
"image_id": frame["id"],
"prompt_type": "semantic",
"prompt_data": "胆囊",
"model": "sam3",
"options": {"min_score": 0.05},
})
assert box_response.status_code == 200
assert box_response.json()["scores"] == [0.8]
assert semantic_response.status_code == 200
assert semantic_response.json()["scores"] == [0.5]
assert calls["semantic"] == {
"model": "sam3",
"text": "胆囊",
"confidence_threshold": 0.05,
}
def test_predict_interactive_combines_box_and_points(client, monkeypatch):

63
backend/tests/test_sam2_engine.py Normal file
View File

@@ -0,0 +1,63 @@
import numpy as np
from services.sam2_engine import 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 = True
engine._predictor = 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(
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(np.zeros((32, 32, 3), dtype=np.uint8))
assert predictor.calls[0]["multimask_output"] is False
assert len(polygons) == 1
assert scores == [0.800000011920929]

50
backend/tests/test_sam3_engine.py
View File

@@ -4,7 +4,7 @@ from pathlib import Path
import numpy as np
from services.sam3_engine import SAM3Engine
from services.sam3_external_worker import _to_numpy
from services.sam3_external_worker import _prediction_to_response, _to_numpy
class _Completed:
@@ -98,6 +98,41 @@ def test_sam3_predict_semantic_uses_external_worker(tmp_path, monkeypatch):
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")
@@ -243,3 +278,16 @@ def test_sam3_worker_casts_floating_tensors_before_numpy():
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]