添加Docker自包含部署分支
- 新增 Seg_Server_Docker 自包含部署内容,包含前后端、FastAPI、Celery、PostgreSQL、Redis、MinIO、演示视频和 DICOM 数据。 - 保留 demo 数据以支持恢复演示出厂设置,排除 SAM 2.1 .pt 权重并在 README 中补充下载命令。 - 补充 GPU 部署、backend/worker 镜像复用、frpc/frps + NPM 公网域名反代部署说明。 - 在 .env/.env.example 中用 # XXXX 标注局域网和公网域名部署需要修改的配置项。 - 添加部署分支 .gitignore,忽略本地模型权重、构建产物、缓存和日志。
This commit is contained in:
130
backend/services/sam_registry.py
Normal file
130
backend/services/sam_registry.py
Normal file
@@ -0,0 +1,130 @@
|
||||
"""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 DEFAULT_SAM2_MODEL_ID, TORCH_AVAILABLE, sam_engine as sam2_engine
|
||||
|
||||
# SAM 3 integration is intentionally disabled for the current product flow.
|
||||
# The source files are kept in the repository so the integration can be
|
||||
# restored later, but the active registry only exposes SAM 2.
|
||||
# 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 DEFAULT_SAM2_MODEL_ID).lower()
|
||||
if self._engines["sam2"].is_sam2_model(selected):
|
||||
return self._engines["sam2"].normalize_model_id(selected)
|
||||
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]:
|
||||
selected = self.normalize_model_id(selected_model)
|
||||
return {
|
||||
"selected_model": selected,
|
||||
"gpu": self.gpu_status(),
|
||||
"models": [sam2_engine.status(model_id) for model_id in sam2_engine.variant_ids()],
|
||||
}
|
||||
|
||||
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:
|
||||
normalized = self.normalize_model_id(model_id)
|
||||
if self._engines["sam2"].is_sam2_model(normalized):
|
||||
return self._engines["sam2"]
|
||||
return self._engines[normalized]
|
||||
|
||||
def _ensure_available(self, model_id: str | None) -> Any:
|
||||
normalized = self.normalize_model_id(model_id)
|
||||
engine = self._engine(model_id)
|
||||
status = engine.status(normalized) if engine is sam2_engine else 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]):
|
||||
model = self.normalize_model_id(model_id)
|
||||
return self._ensure_available(model).predict_points(model, image, points, labels)
|
||||
|
||||
def predict_box(self, model_id: str | None, image: Any, box: list[float]):
|
||||
model = self.normalize_model_id(model_id)
|
||||
return self._ensure_available(model).predict_box(model, 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 not sam2_engine.is_sam2_model(model):
|
||||
raise NotImplementedError("Interactive box + point refinement is currently supported by SAM 2.")
|
||||
return self._ensure_available(model).predict_interactive(model, image, box, points, labels)
|
||||
|
||||
def predict_auto(self, model_id: str | None, image: Any):
|
||||
model = self.normalize_model_id(model_id)
|
||||
return self._ensure_available(model).predict_auto(model, image)
|
||||
|
||||
def predict_semantic(
|
||||
self,
|
||||
model_id: str | None,
|
||||
image: Any,
|
||||
text: str,
|
||||
confidence_threshold: float | None = None,
|
||||
):
|
||||
self.normalize_model_id(model_id)
|
||||
raise NotImplementedError("Semantic text prompting is disabled; use SAM 2 point or box prompts.")
|
||||
|
||||
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,
|
||||
):
|
||||
model = self.normalize_model_id(model_id)
|
||||
return self._ensure_available(model).propagate_video(
|
||||
model,
|
||||
frame_paths,
|
||||
source_frame_index,
|
||||
seed,
|
||||
direction=direction,
|
||||
max_frames=max_frames,
|
||||
)
|
||||
|
||||
|
||||
sam_registry = SAMRegistry()
|
||||
Reference in New Issue
Block a user