iomgaa
150eb5f6ae
- ceus 迁入 polygons、dark_region、reference_region、tic_fitting、tic、anchor_frames、 perfusion_metrics、trajectory、tic_curves、tic_plot、overlay，旧测试照搬（改 import 与配置读取）。 - 对外两段接口：prepare_segmentation（第 1–2 步）与 quantify_case（第 4–9 步）。 - 第 3 条：删掉 D1 取医生标注时刻的模式，D1 只从拟合曲线算。 - 第 4 条：F1/F2 取视频第 20−X、60−X 秒（X 为视频第 0 秒时已推注的秒数），落在视频外记未定位。 - 绿框融合权重、紫圈颜色阈值照搬；关键帧容差窗口的出处写进 ceus.yaml 注释。 - segment：同一进程调用 MedSAM2，设置写进 ceus.yaml 的 segment 段，全部必填；gpu 测试真跑通过。 - 依赖加 scipy、matplotlib、medsam2（可编辑安装，不编译 CUDA 扩展）；uv 只解析 linux x86_64。 - 导入契约：只有 segment 引用 torch、sam2；segment 不引用其他 cadx 模块。 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
"""segment：分割设置的校验，以及在合成小视频上真跑一次 MedSAM2（`gpu`）。"""
|
||
|
||
from pathlib import Path
|
||
from typing import Any, cast
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import numpy.typing as npt
|
||
import pytest
|
||
import yaml
|
||
from pydantic import ValidationError
|
||
|
||
from cadx.segment import SegmentConfig, medsam2_commit, segment_video
|
||
|
||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||
CEUS_YAML = REPO_ROOT / "config" / "ceus.yaml"
|
||
|
||
|
||
def _segment_section() -> dict[str, Any]:
|
||
document = cast(dict[str, Any], yaml.safe_load(CEUS_YAML.read_text(encoding="utf-8")))
|
||
return cast(dict[str, Any], document["segment"])
|
||
|
||
|
||
def _absolute(section: dict[str, Any]) -> dict[str, Any]:
|
||
"""路径字段换成仓库根目录下的绝对路径（脚本也这样做）。"""
|
||
return {
|
||
**section,
|
||
"medsam2_dir": str(REPO_ROOT / section["medsam2_dir"]),
|
||
"checkpoint": str(REPO_ROOT / section["checkpoint"]),
|
||
}
|
||
|
||
|
||
def test_config_yaml_section_validates() -> None:
|
||
cfg = SegmentConfig.model_validate(_segment_section())
|
||
assert cfg.sam2_config == "configs/sam2.1_hiera_t512.yaml"
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"field",
|
||
[
|
||
"medsam2_dir",
|
||
"checkpoint",
|
||
"sam2_config",
|
||
"device",
|
||
"offload_video_to_cpu",
|
||
"offload_state_to_cpu",
|
||
],
|
||
)
|
||
def test_config_missing_field_raises(field: str) -> None:
|
||
section = _segment_section()
|
||
del section[field]
|
||
with pytest.raises(ValidationError, match=field):
|
||
SegmentConfig.model_validate(section)
|
||
|
||
|
||
def test_config_extra_field_raises() -> None:
|
||
with pytest.raises(ValidationError, match="object_id"):
|
||
SegmentConfig.model_validate({**_segment_section(), "object_id": 1})
|
||
|
||
|
||
def test_config_bool_not_coerced() -> None:
|
||
with pytest.raises(ValidationError, match="offload_video_to_cpu"):
|
||
SegmentConfig.model_validate({**_segment_section(), "offload_video_to_cpu": "false"})
|
||
|
||
|
||
def _moving_disc_video(
|
||
frame_count: int, height: int, width: int, radius: int
|
||
) -> tuple[list[npt.NDArray[np.uint8]], list[npt.NDArray[np.bool_]]]:
|
||
"""暗背景上一个亮圆从左往右匀速移动；交出逐帧 BGR 图像和每帧圆的真值掩码。"""
|
||
frames: list[npt.NDArray[np.uint8]] = []
|
||
truths: list[npt.NDArray[np.bool_]] = []
|
||
yy, xx = np.mgrid[0:height, 0:width]
|
||
for index in range(frame_count):
|
||
cx = radius + 10 + index * 2
|
||
cy = height // 2
|
||
frame = np.full((height, width, 3), 30, dtype=np.uint8)
|
||
cv2.circle(frame, (cx, cy), radius, (220, 220, 220), thickness=-1)
|
||
frames.append(frame)
|
||
truths.append((xx - cx) ** 2 + (yy - cy) ** 2 <= radius**2)
|
||
return frames, truths
|
||
|
||
|
||
@pytest.mark.gpu
|
||
def test_segment_video_on_synthetic_clip() -> None:
|
||
cfg = SegmentConfig.model_validate(_absolute(_segment_section()))
|
||
frame_count, height, width = 40, 120, 200
|
||
frames, truths = _moving_disc_video(frame_count, height, width, radius=18)
|
||
# 两个提示帧，最早的不在第 0 帧，前向和反向传播都要走到。
|
||
prompt_masks = {5: truths[5], 25: truths[25]}
|
||
|
||
masks = segment_video(frames, prompt_masks, cfg)
|
||
|
||
assert masks.shape == (frame_count, height, width)
|
||
assert masks.dtype == np.bool_
|
||
assert len(medsam2_commit(cfg.medsam2_dir)) == 40
|