admin
8a9247075e
功能增加: - 将视频导入和生成帧拆成两个明确动作，项目库生成帧时选择 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 需求/设计/接口/测试矩阵，按当前实现冻结功能状态。
121 lines
4.1 KiB
Python
121 lines
4.1 KiB
Python
"""Model registry for SAM runtimes and GPU status."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from config import settings
|
|
from services.sam2_engine import TORCH_AVAILABLE, sam_engine as sam2_engine
|
|
from services.sam3_engine import sam3_engine
|
|
|
|
try:
|
|
import torch
|
|
except Exception: # noqa: BLE001
|
|
torch = None # type: ignore[assignment]
|
|
|
|
|
|
class ModelUnavailableError(RuntimeError):
|
|
"""Raised when a selected model cannot run in this environment."""
|
|
|
|
|
|
class SAMRegistry:
|
|
"""Dispatch predictions to the selected SAM backend."""
|
|
|
|
def __init__(self) -> None:
|
|
self._engines = {
|
|
"sam2": sam2_engine,
|
|
"sam3": sam3_engine,
|
|
}
|
|
|
|
def normalize_model_id(self, model_id: str | None) -> str:
|
|
selected = (model_id or settings.sam_default_model or "sam2").lower()
|
|
if selected not in self._engines:
|
|
raise ValueError(f"Unsupported model: {model_id}")
|
|
return selected
|
|
|
|
def runtime_status(self, selected_model: str | None = None) -> dict[str, Any]:
|
|
return {
|
|
"selected_model": self.normalize_model_id(selected_model),
|
|
"gpu": self.gpu_status(),
|
|
"models": [engine.status() for engine in self._engines.values()],
|
|
}
|
|
|
|
def gpu_status(self) -> dict[str, Any]:
|
|
cuda_available = bool(TORCH_AVAILABLE and torch is not None and torch.cuda.is_available())
|
|
return {
|
|
"available": cuda_available,
|
|
"device": "cuda" if cuda_available else "cpu",
|
|
"name": torch.cuda.get_device_name(0) if cuda_available else None,
|
|
"torch_available": bool(TORCH_AVAILABLE),
|
|
"torch_version": getattr(torch, "__version__", None) if torch is not None else None,
|
|
"cuda_version": getattr(torch.version, "cuda", None) if torch is not None else None,
|
|
}
|
|
|
|
def _engine(self, model_id: str | None) -> Any:
|
|
return self._engines[self.normalize_model_id(model_id)]
|
|
|
|
def _ensure_available(self, model_id: str | None) -> Any:
|
|
engine = self._engine(model_id)
|
|
status = engine.status()
|
|
if not status["available"]:
|
|
raise ModelUnavailableError(status["message"])
|
|
return engine
|
|
|
|
def predict_points(self, model_id: str | None, image: Any, points: list[list[float]], labels: list[int]):
|
|
return self._ensure_available(model_id).predict_points(image, points, labels)
|
|
|
|
def predict_box(self, model_id: str | None, image: Any, box: list[float]):
|
|
return self._ensure_available(model_id).predict_box(image, box)
|
|
|
|
def predict_interactive(
|
|
self,
|
|
model_id: str | None,
|
|
image: Any,
|
|
box: list[float] | None,
|
|
points: list[list[float]],
|
|
labels: list[int],
|
|
):
|
|
model = self.normalize_model_id(model_id)
|
|
if model != "sam2":
|
|
raise NotImplementedError("Interactive box + point refinement is currently supported by SAM 2.")
|
|
return self._ensure_available(model).predict_interactive(image, box, points, labels)
|
|
|
|
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,
|
|
confidence_threshold: float | None = None,
|
|
):
|
|
model = self.normalize_model_id(model_id)
|
|
if model == "sam3":
|
|
return self._ensure_available(model).predict_semantic(
|
|
image,
|
|
text,
|
|
confidence_threshold=confidence_threshold,
|
|
)
|
|
return self._ensure_available(model).predict_auto(image)
|
|
|
|
def propagate_video(
|
|
self,
|
|
model_id: str | None,
|
|
frame_paths: list[str],
|
|
source_frame_index: int,
|
|
seed: dict[str, Any],
|
|
direction: str,
|
|
max_frames: int | None,
|
|
):
|
|
return self._ensure_available(model_id).propagate_video(
|
|
frame_paths,
|
|
source_frame_index,
|
|
seed,
|
|
direction=direction,
|
|
max_frames=max_frames,
|
|
)
|
|
|
|
|
|
sam_registry = SAMRegistry()
|