iomgaa
1de88a5d0a
- 内存：prepare_segmentation 改为流式解码，只留分屏和配准的抽样帧，绿框逐帧在两半上找； 分割那一半由 SegmentationInput.target_frames() 到用时再解码，segment_video 边读边写 jpg； MedSAM2 init_state 的读帧换成直接在存放设备上建张量（不改 MedSAM2 源码，结果逐位相同）； 传播结果写进预分配数组。帧数最多的一段视频常驻内存峰值 11.2 GiB → 2.3 GiB，分辨率最大的 一段原代码被 OOM 杀掉、现为 4.2 GiB；两段的掩码和全部产物与改前逐字节相同，33 段视频的 第 1–2 步结果与改前一致，29 例配准与 review.csv 一致。 - 快照：ceus_quant 开跑时写 snapshot/（git、ceus.yaml、review.csv、uv.lock、segment.json）， 配置和核对表从快照副本读；收尾由新函数 runs.finish_snapshot 只新建 finished.json。 build 与 ceus_align_review 同样改用 finish_snapshot 记结束时间。 - build：summary.json 加 video_sha256；装 CEUS 结果时核对快照里的 delivery_dir、逐例核对 视频 SHA-256 和推注时间，不同就整次停下；新字段做类型检查，不进 tic.json。 来源：主会话任务（问题 1–3 及补充的 delivery_dir 比对）。 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
408 lines
18 KiB
Python
408 lines
18 KiB
Python
"""`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 hashlib
|
||
import importlib.util
|
||
import json
|
||
from collections.abc import Callable, Iterable, Mapping
|
||
from datetime import datetime
|
||
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.answers import read_aliases
|
||
from cadx.build import run_build
|
||
from cadx.ceus import ANCHOR_SLOTS
|
||
from cadx.library import list_aliases, read_ultrasound
|
||
from conftest import CEUS_YAML
|
||
from fake_ceus_video import HEIGHT, LESION, SPLIT, WIDTH, write_video
|
||
from fake_delivery import REPO_ROOT, make_delivery, write_config
|
||
|
||
_SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "ceus_quant.py"
|
||
# review.csv 的列，按任务书 02「配准的目视核对」第 3 步独立写，不取脚本里的常量。
|
||
_REVIEW_HEADER = ("video_sha256", "dx", "dy", "passed", "verdict_a", "verdict_b")
|
||
|
||
|
||
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 _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"),
|
||
"align_runs_root": str(tmp_path / "runs_align"),
|
||
"align_review": str(tmp_path / "review.csv"),
|
||
}
|
||
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
|
||
|
||
|
||
def _write_review(document: dict[str, Any], entries: dict[str, tuple[int, int, bool]]) -> None:
|
||
"""按视频号写配准目视核对表：(dx, dy, 是否通过)，视频 SHA-256 现算。"""
|
||
videos = Path(document["paths"]["delivery_dir"]) / "CEUS_agent"
|
||
lines = [",".join(_REVIEW_HEADER)]
|
||
for video_id, (dx, dy, passed) in entries.items():
|
||
video = videos / f"{int(video_id[1:]):03d}" / f"{video_id}_roi2.mp4"
|
||
sha = hashlib.sha256(video.read_bytes()).hexdigest()
|
||
lines.append(f"{sha},{dx},{dy},{str(passed).lower()},对齐,对齐")
|
||
Path(document["paths"]["align_review"]).write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||
|
||
|
||
class _CountingSegment:
|
||
"""假分割：每帧都是病灶方块挪到左半局部坐标后的位置（略往里收 3 像素）。"""
|
||
|
||
def __init__(self) -> None:
|
||
self.calls = 0
|
||
|
||
def __call__(
|
||
self,
|
||
frames: Iterable[npt.NDArray[np.uint8]],
|
||
prompt_masks: Mapping[int, npt.NDArray[np.bool_]],
|
||
cfg: object,
|
||
) -> npt.NDArray[np.bool_]:
|
||
self.calls += 1
|
||
frames = list(frames)
|
||
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)
|
||
document = _config_dict(tmp_path, delivery_dir)
|
||
# 两段合成视频两半扇形对齐，配准估出 (0, 0)（test_ceus_pipeline 断言过），目视核对都通过。
|
||
_write_review(document, {"U001": (0, 0, True), "U002": (0, 0, True)})
|
||
return module, segment, document, 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["registration"]["contrast_side"] == "right"
|
||
assert summary["registration"]["contrast_to_target_offset"] == [0, 0]
|
||
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()
|
||
assert (snapshot / "align_review.csv").read_bytes() == (tmp_path / "review.csv").read_bytes()
|
||
segment_record = json.loads((snapshot / "segment.json").read_text(encoding="utf-8"))
|
||
assert segment_record["medsam2_commit"] == "c" * 40
|
||
assert segment_record["checkpoint_sha256"] == hashlib.sha256(b"fake weights").hexdigest()
|
||
finished = json.loads((snapshot / "finished.json").read_text(encoding="utf-8"))
|
||
assert finished["cuda_peak_memory_bytes"] is None # device 是 cpu
|
||
assert finished["segment_cache_hits"] == 0
|
||
assert "finished_at" in finished
|
||
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_snapshot_written_before_cases_run(
|
||
tmp_path: Path,
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
|
||
) -> None:
|
||
"""跑病例时快照已在 .partial 里；运行中途改了 ceus.yaml 和 review.csv，快照仍是开跑时那份。"""
|
||
module, segment, document, next_run_id = setup
|
||
config = _write_config(tmp_path, document, "ceus.yaml")
|
||
review = Path(document["paths"]["align_review"])
|
||
original_config, original_review = config.read_bytes(), review.read_bytes()
|
||
snapshot = tmp_path / "runs" / "run-a.partial" / "snapshot"
|
||
seen: list[bool] = []
|
||
|
||
def editing_segment(*args: Any) -> npt.NDArray[np.bool_]:
|
||
seen.append(all((snapshot / name).is_file() for name in ("snapshot.json", "segment.json")))
|
||
config.write_text("# 运行中途改了\n", encoding="utf-8")
|
||
review.write_text("改了\n", encoding="utf-8")
|
||
return segment(*args)
|
||
|
||
monkeypatch.setattr(module, "segment_video", editing_segment)
|
||
next_run_id("run-a")
|
||
module.main(["--config", str(config)])
|
||
|
||
assert seen == [True]
|
||
done = tmp_path / "runs" / "run-a" / "snapshot"
|
||
assert (done / "ceus.yaml").read_bytes() == original_config
|
||
assert (done / "align_review.csv").read_bytes() == original_review
|
||
snapshot_json = json.loads((done / "snapshot.json").read_text(encoding="utf-8"))
|
||
assert "finished_at" not in snapshot_json and "started_at" in snapshot_json
|
||
|
||
|
||
def test_build_reads_real_run_layout(
|
||
tmp_path: Path,
|
||
setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
|
||
) -> None:
|
||
"""build 读这个脚本真实写出的运行目录：第 1 例带 tic.json 和关键帧进库，第 2 例记 C15。"""
|
||
module, _, 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"
|
||
delivery_dir = Path(document["paths"]["delivery_dir"])
|
||
build_dir = tmp_path / "build"
|
||
build_dir.mkdir()
|
||
result = run_build(
|
||
write_config(build_dir, delivery_dir, 1, None, run_dir),
|
||
REPO_ROOT,
|
||
datetime(2026, 10, 16, 12, 0, 0).astimezone(),
|
||
)
|
||
assert [r.case_no for r in read_aliases(result.answers_dir)] == ["001"]
|
||
assert result.skipped_by_check == {"C15": 1}
|
||
(alias,) = list_aliases(result.library_dir)
|
||
route = read_ultrasound(result.library_dir, alias)
|
||
summary = json.loads((run_dir / "U001" / "summary.json").read_text(encoding="utf-8"))
|
||
located = {s for s, a in summary["anchor_frames"].items() if a["status"] == "ok"}
|
||
assert set(route.keyframes) == located
|
||
assert route.tic is not None
|
||
assert set(route.tic) == {"fits", "perfusion_metrics", "pan_curve_usable", "keyframes"}
|
||
assert route.tic["keyframes"]["F1"]["video_time_s"] == pytest.approx(15.0, abs=0.1)
|
||
assert route.tic["keyframes"]["F2"]["located"] is False
|
||
|
||
|
||
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)
|
||
_write_review(document, {"U002": (0, 0, True)}) # 第 1 步失败的视频不进核对表
|
||
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_review_rejected_case_fails_at_step_1(
|
||
tmp_path: Path,
|
||
setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
|
||
) -> None:
|
||
"""目视核对没通过的病例第 1 步失败、不调分割；另一例照常（这里是第 2 步失败）。"""
|
||
module, segment, document, next_run_id = setup
|
||
_write_review(document, {"U001": (0, 0, False), "U002": (0, 0, True)})
|
||
next_run_id("run-a")
|
||
module.main(["--config", str(_write_config(tmp_path, document, "ceus.yaml"))])
|
||
|
||
with (tmp_path / "runs" / "run-a" / "failed.csv").open(encoding="utf-8", newline="") as handle:
|
||
rows = list(csv.reader(handle))
|
||
assert rows[1] == ["U001", "1", "配准目视核对未通过"]
|
||
assert rows[2][:2] == ["U002", "2"]
|
||
assert segment.calls == 0
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("entries", "message"),
|
||
[
|
||
({"U002": (0, 0, True)}, "没有这段视频"), # U001 不在表里
|
||
({"U001": (3, 0, True), "U002": (0, 0, True)}, "不同"), # 表里的平移与这次算出的不同
|
||
],
|
||
)
|
||
def test_review_table_mismatch_stops_run(
|
||
tmp_path: Path,
|
||
setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
|
||
entries: dict[str, tuple[int, int, bool]],
|
||
message: str,
|
||
) -> None:
|
||
module, segment, document, next_run_id = setup
|
||
_write_review(document, entries)
|
||
next_run_id("run-a")
|
||
with pytest.raises(RuntimeError, match=message):
|
||
module.main(["--config", str(_write_config(tmp_path, document, "ceus.yaml"))])
|
||
assert not (tmp_path / "runs" / "run-a").exists()
|
||
assert segment.calls == 0
|
||
|
||
|
||
def test_missing_review_file_stops_run(
|
||
tmp_path: Path,
|
||
setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
|
||
) -> None:
|
||
module, _, document, next_run_id = setup
|
||
Path(document["paths"]["align_review"]).unlink()
|
||
next_run_id("run-a")
|
||
with pytest.raises(FileNotFoundError):
|
||
module.main(["--config", str(_write_config(tmp_path, document, "ceus.yaml"))])
|
||
|
||
|
||
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),
|
||
("ceus", ("registration", "fan_std_threshold"), 5.0, True), # 第 1 步配准
|
||
("ceus", ("registration", "max_disagreement_px"), 9, True),
|
||
("ceus", ("pan_source", "hsv_lower"), [120, 60, 60], 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)
|