feat(ceus): 任务书 02 第三阶段——scripts/ceus_quant.py

- delivery 的 CaseRecord 加视频路径和推注时间（C04、C08 已检查过）。 - scripts/ceus_quant.py：读 ceus.yaml，跑检查全部通过的病例；runs/ceus/<运行号>/ 先写 .partial 再改名；第 1–6 步的单例失败记 failed.csv 后接着跑；终端只打印计数。 - 分割缓存 work/segment_cache/<键>/：键含视频 SHA-256、第 1–3 步参数、权重 SHA-256、MedSAM2 提交号。 - 视频解码错误改为第 1 步的单例失败（VideoDecodeError），不再停掉整次运行（审阅意见）。 - 暗区 kmeans 用 cv2 进程级随机数，结果受病例顺序影响；ceus 配置加必填 random_seed，每例第 4 步前设种子。 - segment 导出 cuda_peak_memory_bytes，供运行快照记录显存峰值。 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-16 05:30:29 -04:00
parent 150eb5f6ae
commit 3524b31e3e
12 changed files with 667 additions and 26 deletions
@@ -0,0 +1,286 @@
+"""`scripts/ceus_quant.py`：逐例循环、failed.csv、输出布局、分割缓存。
+
+假交付包两例：第 1 例是能过质量门的合成分屏造影视频（造影在右半，绿框是病灶的空心描边），第 2 例
+同样的画面但没有绿框。分割换成计数的假函数（monkeypatch 脚本里引用的 `segment_video`），交出病灶
+方块挪到左半局部坐标的掩码，不调显卡。
+
+手算的预期（视频 400×200、10 fps、40 秒，推注时间 X = 5 秒，同 test_ceus_pipeline）：
+- 第 2 例两半都没有绿框 → 第 2 步失败，进 failed.csv，没有输出目录，也不调分割；
+- 第 1 例 F1 在视频第 20 − 5 = 15 秒，有关键帧图；F2 在第 55 秒，超出视频，未定位，没有图。
+"""
+
+import csv
+import importlib.util
+import json
+from collections.abc import Callable, Mapping, Sequence
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+import pytest
+import yaml
+
+from cadx.ceus import ANCHOR_SLOTS, LognormalParameters, evaluate_lognormal
+from conftest import CEUS_YAML
+from fake_delivery import make_delivery
+
+_SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "ceus_quant.py"
+_WIDTH, _HEIGHT, _SPLIT = 400, 200, 200
+_FPS, _FRAMES = 10.0, 400
+_LESION = (250, 60, 369, 159)  # 整帧坐标 x1, y1, x2, y2（含端点）
+_CYST = (290, 95, 320, 125)
+_BOLUS = LognormalParameters(area=1300.0, mu=2.4, sigma=0.45, arrival_offset_s=4.0, baseline=40.0)
+
+
+def _load_script() -> ModuleType:
+    spec = importlib.util.spec_from_file_location("ceus_quant", _SCRIPT)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def _write_video(path: Path, green_box: bool) -> None:
+    rng = np.random.default_rng(3)
+    intensity = evaluate_lognormal(np.arange(_FRAMES) / _FPS, _BOLUS)
+    writer = cv2.VideoWriter(str(path), cv2.VideoWriter.fourcc(*"mp4v"), _FPS, (_WIDTH, _HEIGHT))
+    x1, y1, x2, y2 = _LESION
+    for index in range(_FRAMES):
+        frame = np.zeros((_HEIGHT, _WIDTH, 3), dtype=np.uint8)
+        frame[:, :_SPLIT] = 120
+        frame[:, _SPLIT + 1 :] = 30
+        lesion = np.clip(
+            intensity[index] + rng.normal(0.0, 6.0, (y2 - y1 + 1, x2 - x1 + 1)), 0, 255
+        )
+        frame[y1 : y2 + 1, x1 : x2 + 1] = lesion.astype(np.uint8)[:, :, None]
+        cx1, cy1, cx2, cy2 = _CYST
+        frame[cy1 : cy2 + 1, cx1 : cx2 + 1] = 10
+        if green_box:
+            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), thickness=3)
+        writer.write(frame)
+    writer.release()
+
+
+def _config_dict(tmp_path: Path, delivery_dir: Path) -> dict[str, Any]:
+    document: dict[str, Any] = yaml.safe_load(CEUS_YAML.read_text(encoding="utf-8"))
+    document["paths"] = {
+        "delivery_dir": str(delivery_dir),
+        "runs_root": str(tmp_path / "runs"),
+        "work_root": str(tmp_path / "work"),
+    }
+    checkpoint = tmp_path / "fake_checkpoint.pt"
+    checkpoint.write_bytes(b"fake weights")
+    document["segment"].update(
+        {"medsam2_dir": str(tmp_path / "MedSAM2"), "checkpoint": str(checkpoint), "device": "cpu"}
+    )
+    return document
+
+
+def _write_config(tmp_path: Path, document: dict[str, Any], name: str) -> Path:
+    path = tmp_path / name
+    path.write_text(yaml.safe_dump(document, allow_unicode=True), encoding="utf-8")
+    return path
+
+
+class _CountingSegment:
+    """假分割：每帧都是病灶方块挪到左半局部坐标后的位置（略往里收 3 像素）。"""
+
+    def __init__(self) -> None:
+        self.calls = 0
+
+    def __call__(
+        self,
+        frames: Sequence[npt.NDArray[np.uint8]],
+        prompt_masks: Mapping[int, npt.NDArray[np.bool_]],
+        cfg: object,
+    ) -> npt.NDArray[np.bool_]:
+        self.calls += 1
+        height, width = frames[0].shape[:2]
+        masks = np.zeros((len(frames), height, width), dtype=bool)
+        x1, y1, x2, y2 = _LESION
+        masks[:, y1 + 3 : y2 - 2, x1 - _SPLIT + 3 : x2 - _SPLIT - 2] = True
+        return masks
+
+
+@pytest.fixture
+def setup(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]]:
+    """写假交付包（两例真视频），载入脚本并换掉分割、MedSAM2 提交号和运行号。"""
+    delivery_dir = make_delivery(2).write(tmp_path / "delivery")
+    _write_video(delivery_dir / "CEUS_agent" / "001" / "U001_roi2.mp4", green_box=True)
+    _write_video(delivery_dir / "CEUS_agent" / "002" / "U002_roi2.mp4", green_box=False)
+    module = _load_script()
+    segment = _CountingSegment()
+    monkeypatch.setattr(module, "segment_video", segment)
+    run_ids: list[str] = []
+
+    def fake_commit(medsam2_dir: Path) -> str:
+        return "c" * 40
+
+    def fake_run_id(now: object, git: object) -> str:
+        return run_ids[-1]
+
+    monkeypatch.setattr(module, "medsam2_commit", fake_commit)
+    monkeypatch.setattr(module, "make_build_id", fake_run_id)
+    return module, segment, _config_dict(tmp_path, delivery_dir), run_ids.append
+
+
+def test_run_writes_outputs_and_failed_case(
+    tmp_path: Path,
+    setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
+    capsys: pytest.CaptureFixture[str],
+) -> None:
+    module, segment, document, next_run_id = setup
+    next_run_id("run-a")
+    module.main(["--config", str(_write_config(tmp_path, document, "ceus.yaml"))])
+
+    run_dir = tmp_path / "runs" / "run-a"
+    assert sorted(p.name for p in (tmp_path / "runs").iterdir()) == ["run-a"]
+    with (run_dir / "failed.csv").open(encoding="utf-8", newline="") as handle:
+        rows = list(csv.reader(handle))
+    assert rows[0] == ["video_id", "step", "reason"]
+    assert [row[:2] for row in rows[1:]] == [["U002", "2"]]
+    assert "都没有" in rows[1][2]
+    assert not (run_dir / "U002").exists()
+    assert segment.calls == 1
+
+    case_dir = run_dir / "U001"
+    for name in ("trajectory.npz", "tic_curves.npz", "summary.json", "tic.png", "overlay.mp4"):
+        assert (case_dir / name).stat().st_size > 0, name
+    summary = json.loads((case_dir / "summary.json").read_text(encoding="utf-8"))
+    assert summary["video_id"] == "U001"
+    assert summary["injection_offset_s"] == 5.0
+    assert list(summary["anchor_frames"]) == list(ANCHOR_SLOTS)
+    assert summary["perfusion_metrics"]["roi_perf"] is not None
+    assert "pan_detected_fraction" in summary
+    assert summary["anchor_frames"]["F2"]["status"] == "not_covered"
+    located = {s for s, a in summary["anchor_frames"].items() if a["frame_index"] is not None}
+    assert "F1" in located and "F2" not in located
+    keyframes = {p.stem for p in (case_dir / "keyframes").iterdir()}
+    assert keyframes == located
+    image = cv2.imread(str(case_dir / "keyframes" / "F1.png"))
+    assert image is not None and image.shape == (_HEIGHT, _WIDTH, 3)  # 整幅原始画面
+
+    snapshot = run_dir / "snapshot"
+    assert (snapshot / "ceus.yaml").read_bytes() == (tmp_path / "ceus.yaml").read_bytes()
+    segment_record = json.loads((snapshot / "segment.json").read_text(encoding="utf-8"))
+    assert segment_record["medsam2_commit"] == "c" * 40
+    assert segment_record["cuda_peak_memory_bytes"] is None  # device 是 cpu
+    assert (snapshot / "snapshot.json").is_file()
+
+    out = capsys.readouterr().out
+    assert "成功：1" in out and "第 2 步 1" in out
+    assert "U001" not in out and "U002" not in out
+
+
+def test_unreadable_video_fails_at_step_1_and_next_case_runs(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
+) -> None:
+    module, segment, document, next_run_id = setup
+    # 第 1 例换成损坏的视频（cv2 打不开），第 2 例换成正常视频。
+    videos = Path(document["paths"]["delivery_dir"]) / "CEUS_agent"
+    (videos / "001" / "U001_roi2.mp4").write_bytes(b"not a real video")
+    _write_video(videos / "002" / "U002_roi2.mp4", green_box=True)
+    next_run_id("run-a")
+    module.main(["--config", str(_write_config(tmp_path, document, "ceus.yaml"))])
+
+    run_dir = tmp_path / "runs" / "run-a"
+    with (run_dir / "failed.csv").open(encoding="utf-8", newline="") as handle:
+        rows = list(csv.reader(handle))
+    assert [row[:2] for row in rows[1:]] == [["U001", "1"]]
+    assert "打不开" in rows[1][2]
+    assert not (run_dir / "U001").exists()
+    assert (run_dir / "U002" / "summary.json").is_file()
+    assert segment.calls == 1
+
+
+def test_segment_cache_reused_and_keyed_on_step_1_3_params(
+    tmp_path: Path,
+    setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
+) -> None:
+    module, segment, document, next_run_id = setup
+    config = _write_config(tmp_path, document, "ceus.yaml")
+    next_run_id("run-a")
+    module.main(["--config", str(config)])
+    assert segment.calls == 1
+    cache = tmp_path / "work" / "segment_cache"
+    assert [p.name.endswith(".partial") for p in cache.iterdir()] == [False]
+
+    next_run_id("run-b")
+    module.main(["--config", str(config)])
+    assert segment.calls == 1  # 同样的输入，读缓存
+    # 同一进程里第二次跑（cv2 的随机数已被第一次推进过）结果也逐位相同：quantify_case 每例重设种子。
+    first = json.loads((tmp_path / "runs/run-a/U001/summary.json").read_text(encoding="utf-8"))
+    second = json.loads((tmp_path / "runs/run-b/U001/summary.json").read_text(encoding="utf-8"))
+    assert first == second
+
+    # 改设备不改结果，不重算；改第 1 步的参数，键变了，重新分割。
+    document["segment"]["device"] = "cuda:0"
+    document["segment"]["offload_state_to_cpu"] = True
+    next_run_id("run-c")
+    module.main(["--config", str(_write_config(tmp_path, document, "device.yaml"))])
+    assert segment.calls == 1
+    document["ceus"]["views"]["split_sample_frames"] = 20
+    next_run_id("run-d")
+    module.main(["--config", str(_write_config(tmp_path, document, "views.yaml"))])
+    assert segment.calls == 2
+    assert len(list(cache.iterdir())) == 2
+
+
+def test_existing_run_directory_is_not_overwritten(
+    tmp_path: Path,
+    setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
+) -> None:
+    module, _, document, next_run_id = setup
+    config = _write_config(tmp_path, document, "ceus.yaml")
+    next_run_id("run-a")
+    module.main(["--config", str(config)])
+    before = sorted(str(p) for p in (tmp_path / "runs").rglob("*"))
+    with pytest.raises(FileExistsError, match="运行目录已存在"):
+        module.main(["--config", str(config)])
+    assert sorted(str(p) for p in (tmp_path / "runs").rglob("*")) == before
+
+
+def _key(module: ModuleType, document: dict[str, Any]) -> str:
+    config = module.QuantConfig.model_validate(document)
+    return module.segment_cache_key("v" * 64, config, "w" * 64, "c" * 40)
+
+
+@pytest.mark.parametrize(
+    ("section", "path", "value", "changes"),
+    [
+        ("ceus", ("views", "split_search_band"), [0.3, 0.7], True),
+        ("ceus", ("prompt_source", "min_area"), 9000, True),
+        ("ceus", ("keyframes", "iou_threshold"), 0.9, True),
+        ("segment", ("sam2_config",), "configs/other.yaml", True),
+        ("segment", ("checkpoint",), "models/other.pt", True),
+        ("segment", ("device",), "cuda:1", False),
+        ("segment", ("offload_video_to_cpu",), True, False),
+        ("ceus", ("smoothing", "temporal_window"), 9, False),  # 第 4 步
+        ("ceus", ("tic_analysis", "fitting", "minimum_r2"), 0.5, False),  # 第 7 步
+    ],
+)
+def test_cache_key_covers_exactly_step_1_3_params(
+    tmp_path: Path,
+    section: str,
+    path: tuple[str, ...],
+    value: object,
+    changes: bool,
+) -> None:
+    module = _load_script()
+    document = _config_dict(tmp_path, tmp_path)
+    base = _key(module, document)
+    node = document[section]
+    for name in path[:-1]:
+        node = node[name]
+    node[path[-1]] = value
+    assert (_key(module, document) != base) is changes
+    assert _key(module, document) == _key(module, document)  # 同输入同键
+    other = module.QuantConfig.model_validate(document)
+    assert module.segment_cache_key("x" * 64, other, "w" * 64, "c" * 40) != _key(module, document)