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>
This commit is contained in:
@@ -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)
|
||||
@@ -418,3 +418,19 @@ def test_skip_record_names_file_and_column(tmp_path: Path) -> None:
|
||||
assert record.files == (ULTRASOUND_PATH,)
|
||||
assert record.columns == ("推注开始时间s",)
|
||||
assert record.value_json == '{"推注开始时间s": 2.5}'
|
||||
|
||||
|
||||
def test_case_record_carries_video_path_and_bolus_time(tmp_path: Path) -> None:
|
||||
delivery = make_delivery(3)
|
||||
delivery.case(2).dir_files = ["U002_key.png", "U002.mp4", "U002_report.xlsx"]
|
||||
delivery.case(3).ultrasound["推注开始时间s"] = 0 # 0 是合法取值:视频第 0 秒正好推注
|
||||
root = delivery.write(tmp_path / "delivery")
|
||||
result = read_delivery(root, None)
|
||||
folder = root / "CEUS_agent"
|
||||
assert [c.video for c in result.cases] == [
|
||||
folder / "001" / "U001_roi2.mp4",
|
||||
folder / "002" / "U002.mp4",
|
||||
folder / "003" / "U003_roi2.mp4",
|
||||
]
|
||||
# 假交付包的推注时间是 5,第 3 例改成 0。
|
||||
assert [c.bolus_start_s for c in result.cases] == [5, 5, 0]
|
||||
|
||||
Reference in New Issue
Block a user