diff --git a/research-wiki/explanation/结构.md b/research-wiki/explanation/结构.md index 96c8368..5b84e63 100644 --- a/research-wiki/explanation/结构.md +++ b/research-wiki/explanation/结构.md @@ -165,6 +165,7 @@ library/<构建号>/ - `keyframes/<槽位>.png`:下文关键帧表里定位到的那几张(槽位即 B0、D1 等),和 `key.png` 一样去掉附加数据重新编码;每例最多 8 张图(7 张关键帧加 `key.png`)。 - `tic.json`:从这一例的 `summary.json` 里逐字段挑出,此外什么都不写:四条曲线(`roi_all` 病灶整体、`roi_perf` 病灶去暗区、`roi_pan` 紫圈即医生圈的一块正常胰腺、`background_circle` 背景圆)各自的拟合是否成功(`success`)和拟合优度 R²(`r2`);病灶两条曲线的灌注指标(`perfusion_metrics`,拟合没过质量门为 null);紫圈曲线是否可用(`pan_curve_usable`);7 个槽位各自是否定位到(`located`)、定位到的视频时间(`video_time_s`,秒)、没定位的原因(`reason`)。视频号、文件路径、配置、帧号、拟合参数都不写。挑哪些字段只在 `src/cadx/build/_ceus.py` 里列。 +- 核对这次 CEUS 运行是在这份交付包上跑的:运行快照里 `ceus.yaml` 的 `paths.delivery_dir` 与 `build.yaml` 的 `delivery_dir` 规范化后不是同一路径;或逐例比,交付包里这段视频文件的 SHA-256 与 `summary.json` 的 `video_sha256` 不同,或交付包的推注时间与 `summary.json` 的 `injection_offset_s` 不同,整次停下报错。这两个字段不进 `tic.json`。 - CEUS 在这一例上失败(在 `failed.csv` 里)或运行目录里没有这一例的结果,这一例整例跳过,记 C15。失败的病例要不要改成「超声不可评价、仍参与会诊」等 `../待定问题/03-病例排除标准.md` 定。 每一路能看见什么: @@ -371,15 +372,16 @@ Python 环境用 uv;安装与检查命令见 `../guides/开发环境.md`。 - `read_git_state(repo_root)`:读出完整和短提交号、工作区是否干净(含未跟踪文件)、`git diff HEAD` 输出的 SHA-256。 - `GitState`:上面这四项。 - `make_build_id(now, git)`:构建号 `YYYYMMDD-HHMMSS-`,工作区不干净时末尾加 `-dirty`。 -- `write_snapshot(snapshot_dir, build_id, git, started_at, finished_at, uv_lock, configs)`:新建快照目录,复制配置原文并写 `snapshot.json`(构建号、git 状态、Python 版本、`uv.lock` 哈希、起止时间、各配置的哈希)。 +- `write_snapshot(snapshot_dir, build_id, git, started_at, uv_lock, configs)`:开跑时调。新建快照目录,复制配置原文并写 `snapshot.json`(构建号、git 状态、Python 版本、`uv.lock` 哈希、开始时间、各配置的哈希)。运行要跑很久时(CEUS 定量约一小时)在开跑时写,期间改了配置文件,快照里仍是这次用的那份。 +- `finish_snapshot(snapshot_dir, finished_at, record)`:收尾时调。在快照目录里新建 `finished.json`:结束时间加 `record`(运行结束才知道的值,如显存峰值、缓存命中数);已存在就报错,不改开跑时写的文件。 ### ceus(`src/cadx/ceus/`) 对外两段接口,中间由 `scripts/ceus_quant.py` 调 segment 做第 3 步: -- `prepare_segmentation(video_path, cfg)`:第 1–2 步,交出 `SegmentationInput`;视频打不开或读出的不合规(fps、尺寸、时间戳、帧数)抛 `VideoDecodeError`、找不到可信分屏线抛 `SplitDetectionError`、两半配准失败(找不到扇形区、可用模板匹配太少,或两种方法估出的平移相差超过 `registration.max_disagreement_px`)抛 `RegistrationError`(都是第 1 步),两半都有或都没有绿框抛 `ContrastSideError`(第 2 步)。 +- `prepare_segmentation(video_path, cfg)`:第 1–2 步,交出 `SegmentationInput`(逐帧解码几遍,不把整段视频留在内存里,只留分割那一半的画面);视频打不开或读出的不合规(fps、尺寸、时间戳、帧数)抛 `VideoDecodeError`、找不到可信分屏线抛 `SplitDetectionError`、两半配准失败(找不到扇形区、可用模板匹配太少,或两种方法估出的平移相差超过 `registration.max_disagreement_px`)抛 `RegistrationError`(都是第 1 步),两半都有或都没有绿框抛 `ContrastSideError`(第 2 步)。 - `register_halves(frames, left_view, right_view, cfg.registration, cfg.prompt_source, cfg.pan_source)`:第 1 步的两半配准,交出 `HalfRegistration`:左半局部坐标 → 右半局部坐标的平移,`correlation`(两半扇形区边缘图的归一化互相关,采用它)、`template`(左半扇形内多块模板逐帧匹配后取中位数,只用来核对)、各次匹配明细、两者各分量之差的最大值;`contrast_to_target(side)` 换成造影那一半 → 另一半的平移。 -- `SegmentationInput`:视频的 fps、帧数、宽高、逐帧时间戳,分屏线,`registration`(`HalfRegistration`),`offset`(造影那一半局部坐标 + `offset` = `target_view` 局部坐标),`contrast`(`ContrastSide`:哪一半是造影画面、造影视图和分割那一半视图 `target_view`、提示帧),`target_frames`(分割那一半的逐帧 BGR 图像),`prompt_masks`(提示帧帧号 → 绿框按 `offset` 挪到 `target_view` 局部坐标的掩膜),造影那一半的第一帧。 +- `SegmentationInput`:视频的 fps、帧数、宽高、逐帧时间戳,分屏线,`registration`(`HalfRegistration`),`offset`(造影那一半局部坐标 + `offset` = `target_view` 局部坐标),`contrast`(`ContrastSide`:哪一半是造影画面、造影视图和分割那一半视图 `target_view`、提示帧),`target_frames()`(重新解码,逐帧交出分割那一半的 BGR 图像;不存整段),`prompt_masks`(提示帧帧号 → 绿框按 `offset` 挪到 `target_view` 局部坐标的掩膜),造影那一半的第一帧。 - `quantify_case(prepared, target_masks, injection_offset_s, cfg)`:第 4–9 步。`target_masks` 是分割在 `target_view` 局部坐标上的 `(帧数, 高, 宽)` 布尔数组,`injection_offset_s` 是推注时间 X;分割一帧有效轮廓都没有抛 `NoValidContourError`(第 4 步)。拟合没过质量门不抛错。 - `QuantResult`:`trajectory`(`Trajectory`)、`tic`(`TicResult`)、`curves`(`TicCurves`)、推注时间、病灶两条曲线拟合是否都过质量门、`anchors`(7 个槽位的 `AnchorFrameLocation`)、`perfusion`(`roi_all`、`roi_perf` 各自的 `LognormalPerfusionMetrics`,拟合没过为 None)、紫圈检出率。 - `AnchorFrameLocation`:一个关键帧槽位的帧号、视频时间、推注后时刻、状态(`ok` 或未定位的类别)、未定位的原因、锚点时刻。 @@ -392,7 +394,7 @@ Python 环境用 uv;安装与检查命令见 `../guides/开发环境.md`。 ### segment(`src/cadx/segment/`) -- `segment_video(frames, prompt_masks, cfg)`:把提示帧的病灶掩码用 MedSAM2 传播到每一帧;`frames` 是分割那一半的逐帧 BGR 图像,`prompt_masks` 是提示帧帧号 → 同尺寸布尔掩码,交出 `(帧数, 高, 宽)` 布尔数组(即 `quantify_case` 的 `target_masks`)。 +- `segment_video(frames, prompt_masks, cfg)`:把提示帧的病灶掩码用 MedSAM2 传播到每一帧;`frames` 是分割那一半的逐帧 BGR 图像(只遍历一遍,边读边写成 jpg,可以是逐帧解码的生成器),`prompt_masks` 是提示帧帧号 → 同尺寸布尔掩码,交出 `(帧数, 高, 宽)` 布尔数组(即 `quantify_case` 的 `target_masks`)。MedSAM2 的 `init_state` 读帧时换成本模块的读法,帧张量直接建在显卡上,不先在 CPU 上建整段(结果逐位相同)。 - `SegmentConfig`:`config/ceus.yaml` 里 `segment` 段的字段定义(MedSAM2 代码目录、权重、模型配置名、设备、两个 offload 开关;pydantic,字段必填、不许多写、读入后不可改)。 - `medsam2_commit(medsam2_dir)`:读 MedSAM2 代码目录 `.git` 的 HEAD 提交号,给快照和分割缓存键用。 - `cuda_peak_memory_bytes(device)`:本进程在该设备上分配过的显存峰值(字节),给运行快照用;非 CUDA 设备为 None,本进程没初始化过 CUDA(分割全读了缓存)为 0。放在 segment 里,是因为 `cadx` 里只有它引用 `torch`。 diff --git a/scripts/ceus_align_review.py b/scripts/ceus_align_review.py index 4f18592..e8419fc 100644 --- a/scripts/ceus_align_review.py +++ b/scripts/ceus_align_review.py @@ -23,7 +23,7 @@ import numpy.typing as npt from cadx import ceus from cadx.ceus import CeusConfig, HalfRegistration, ViewRect, crop_frame from cadx.delivery import read_delivery -from cadx.runs import make_build_id, read_git_state, write_snapshot +from cadx.runs import finish_snapshot, make_build_id, read_git_state, write_snapshot from ceus_quant import REPO_ROOT, REVIEW_HEADER, file_sha256, load_config, resolve_paths KEY_HEADER = ("id", "video_id", "is_control", "dx", "dy", "video_sha256") @@ -35,13 +35,6 @@ type _Gray = npt.NDArray[np.float32] type _Mask = npt.NDArray[np.bool_] -def _indices(count: int, wanted: int) -> set[int]: - """与 `sample_frames_for_split` 相同的均匀抽帧下标。""" - if count <= wanted: - return set(range(count)) - return {int(i) for i in np.linspace(0, count - 1, num=wanted, dtype=np.int32).tolist()} - - @dataclass(frozen=True) class Material: """一例出图要的东西:两个视图;整帧的时间平均灰度、抽样中间那帧的灰度、绿框紫圈像素;两半扇形。""" @@ -59,8 +52,8 @@ def step1(video_path: Path, cfg: CeusConfig) -> tuple[Material, HalfRegistration """第 1 步:逐帧解码只留分屏和配准要的抽样帧,找分屏线,配准两半。""" video = ceus.VideoFile(video_path) _, frame_count, _, _ = video.probe() - split_at = _indices(frame_count, cfg.views.split_sample_frames) - register_at = _indices(frame_count, cfg.registration.sample_frames) + split_at = set(ceus.sample_indices(frame_count, cfg.views.split_sample_frames)) + register_at = set(ceus.sample_indices(frame_count, cfg.registration.sample_frames)) split_frames, frames = list[_U8](), list[_U8]() decoded = 0 for index, frame in video.iter_frames(): @@ -169,9 +162,9 @@ def render(config_path: Path) -> None: csv.writer(handle).writerows([KEY_HEADER, *sorted(rows)]) with (partial / "failed.csv").open("w", encoding="utf-8", newline="") as handle: csv.writer(handle).writerows([("video_id", "step", "reason"), *failures]) - finished_at, lock = datetime.now().astimezone(), REPO_ROOT / "uv.lock" configs = {"ceus.yaml": config_path} # align_check 段的随机数种子随 ceus.yaml 进快照 - write_snapshot(partial / "snapshot", run_id, git, started_at, finished_at, lock, configs) + write_snapshot(partial / "snapshot", run_id, git, started_at, REPO_ROOT / "uv.lock", configs) + finish_snapshot(partial / "snapshot", datetime.now().astimezone(), {}) if run_dir.exists(): raise FileExistsError(f"运行目录已存在,不覆盖:{run_dir}") partial.rename(run_dir) diff --git a/scripts/ceus_quant.py b/scripts/ceus_quant.py index 3aeef98..2056efe 100644 --- a/scripts/ceus_quant.py +++ b/scripts/ceus_quant.py @@ -7,6 +7,8 @@ 第 3 步(MedSAM2 分割)的结果按内容键缓存在 `/segment_cache/<键>/masks.npz`,键相同就 直接读缓存、不调显卡。某一例在第 1–6 步失败(`CeusCaseError`)记进 `failed.csv`,接着跑下一例; 其余异常整次运行停下,留下 `.partial` 目录。终端只打印计数,不打印视频号。 +快照在开跑时写进 `.partial/snapshot/`(代码、配置和核对表原文、权重与 MedSAM2 版本),配置和核对表 +随后从快照里的副本读;收尾只补 `finished.json`(结束时间、显存峰值、分割缓存命中数)。 """ import argparse @@ -44,7 +46,7 @@ from cadx.ceus import ( tic_plot_markers, ) from cadx.delivery import CaseRecord, read_delivery -from cadx.runs import make_build_id, read_git_state, write_snapshot +from cadx.runs import finish_snapshot, make_build_id, read_git_state, write_snapshot from cadx.segment import SegmentConfig, cuda_peak_memory_bytes, medsam2_commit, segment_video REPO_ROOT = Path(__file__).resolve().parents[1] @@ -241,7 +243,7 @@ def _masks( if final.exists(): with np.load(final / "masks.npz", allow_pickle=False) as data: return data["masks"], True - masks = segment_video(prepared.target_frames, prepared.prompt_masks, segment_cfg) + masks = segment_video(prepared.target_frames(), prepared.prompt_masks, segment_cfg) partial = cache_root / f"{key}.partial" # 上次中途停下留下的 .partial 会让这里报错:work/ 可以删掉重算,删掉它再跑。 partial.mkdir(parents=True, exist_ok=False) @@ -283,6 +285,7 @@ def _run_case( registration = prepared.registration summary = { "video_id": case.video_id, + "video_sha256": video_sha256, "segment_cache_key": key, "registration": { "contrast_side": prepared.contrast.side, @@ -323,10 +326,31 @@ def main(argv: Sequence[str]) -> None: raise FileExistsError(f"运行目录已存在,不覆盖:{path}") partial.mkdir(parents=True, exist_ok=False) - review = read_review(cfg.paths.align_review) - delivery = read_delivery(cfg.paths.delivery_dir, None) checkpoint_sha256 = file_sha256(cfg.segment.checkpoint) commit = medsam2_commit(cfg.segment.medsam2_dir) + # 开跑时就写快照:运行期间改了 ceus.yaml 或 review.csv,快照里仍是这次用的那份。 + snapshot_dir = partial / "snapshot" + write_snapshot( + snapshot_dir=snapshot_dir, + build_id=run_id, + git=git, + started_at=started_at, + uv_lock=REPO_ROOT / "uv.lock", + configs={"ceus.yaml": config_path, "align_review.csv": cfg.paths.align_review}, + ) + segment_record = { + "checkpoint_sha256": checkpoint_sha256, + "medsam2_commit": commit, + "device": cfg.segment.device, + } + (snapshot_dir / "segment.json").write_text( + json.dumps(segment_record, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" + ) + # 配置和核对表读快照里的副本,与快照记的是同一份字节。 + written = load_config(snapshot_dir / "ceus.yaml") + cfg = resolve_paths(written) + review = read_review(snapshot_dir / "align_review.csv") + delivery = read_delivery(cfg.paths.delivery_dir, None) # 缓存键用配置文件里的写法(路径不随仓库位置变)。 def key_of(video_sha256: str) -> str: @@ -352,24 +376,10 @@ def main(argv: Sequence[str]) -> None: writer.writerow(_FAILED_HEADER) writer.writerows(failures) peak = cuda_peak_memory_bytes(cfg.segment.device) - write_snapshot( - snapshot_dir=partial / "snapshot", - build_id=run_id, - git=git, - started_at=started_at, - finished_at=datetime.now().astimezone(), - uv_lock=REPO_ROOT / "uv.lock", - configs={"ceus.yaml": config_path, "align_review.csv": cfg.paths.align_review}, - ) - segment_record = { - "checkpoint_sha256": checkpoint_sha256, - "medsam2_commit": commit, - "device": cfg.segment.device, - "cuda_peak_memory_bytes": peak, - "segment_cache_hits": cache_hits, - } - (partial / "snapshot" / "segment.json").write_text( - json.dumps(segment_record, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" + finish_snapshot( + snapshot_dir, + datetime.now().astimezone(), + {"cuda_peak_memory_bytes": peak, "segment_cache_hits": cache_hits}, ) # POSIX 的 rename 会悄悄覆盖空目录,所以改名前再查一次。 if run_dir.exists(): diff --git a/src/cadx/build/_assemble.py b/src/cadx/build/_assemble.py index 22c3647..73e11f0 100644 --- a/src/cadx/build/_assemble.py +++ b/src/cadx/build/_assemble.py @@ -27,7 +27,7 @@ from cadx.delivery import ( review_sentences, ) from cadx.library import LibraryCase, LibraryCeus, check_key_image, write_library -from cadx.runs import make_build_id, read_git_state, write_snapshot +from cadx.runs import finish_snapshot, make_build_id, read_git_state, write_snapshot ALIAS_ALPHABET = string.ascii_lowercase + string.digits ALIAS_LENGTH = 4 @@ -155,7 +155,11 @@ def run_build(config_path: Path, repo_root: Path, now: datetime) -> BuildSummary delivery = read_delivery(config.delivery_dir, code_maps) passed, key_image_skips = _check_key_images(delivery.cases, config.delivery_dir) skipped = [*delivery.skipped, *key_image_skips] - ceus_run = None if config.ceus_run is None else open_ceus_run(config.ceus_run) + ceus_run = ( + None + if config.ceus_run is None + else open_ceus_run(config.ceus_run, config.delivery_dir, repo_root) + ) ceus_of: dict[str, LibraryCeus | None] = dict.fromkeys((c.case_no for c in passed), None) if ceus_run is not None: with_ceus, ceus_skips = _attach_ceus(passed, ceus_run) @@ -214,10 +218,10 @@ def run_build(config_path: Path, repo_root: Path, now: datetime) -> BuildSummary build_id=build_id, git=git, started_at=now, - finished_at=datetime.now().astimezone(), uv_lock=repo_root / "uv.lock", configs=configs, ) + finish_snapshot(answers_partial / "snapshot", datetime.now().astimezone(), {}) write_answers( answers_dir=answers_partial, build_id=build_id, diff --git a/src/cadx/build/_ceus.py b/src/cadx/build/_ceus.py index e191d92..5baa14b 100644 --- a/src/cadx/build/_ceus.py +++ b/src/cadx/build/_ceus.py @@ -1,16 +1,20 @@ """读一次 CEUS 定量运行(`runs/ceus/<运行号>/`),按白名单挑出超声一路的 tic.json 和关键帧(C15)。 运行目录的布局由 `scripts/ceus_quant.py` 定:`failed.csv`、`snapshot/snapshot.json`、 -`<视频号>/summary.json`、`<视频号>/keyframes/<槽位>.png`。tic.json 的字段只在这里列, -视频号、路径、配置、帧号和拟合参数都不进去。 +`snapshot/ceus.yaml`、`<视频号>/summary.json`、`<视频号>/keyframes/<槽位>.png`。tic.json 的字段 +只在这里列,视频号、路径、配置、帧号和拟合参数都不进去。快照里的交付包路径、summary.json 里记的 +视频 SHA-256 和推注时间要与这份交付包一致,否则运行不是在这份交付包上跑的,整次停下。 """ import csv +import hashlib import json from dataclasses import dataclass from pathlib import Path from typing import Any +import yaml + from cadx.delivery import CaseRecord, SkipRecord from cadx.library import LibraryCeus @@ -51,14 +55,27 @@ class CeusRun: failed: dict[str, tuple[str, str]] -def open_ceus_run(run_dir: Path) -> CeusRun: - """核对运行目录(存在、不是 `.partial`、有 failed.csv 和快照),读 failed.csv。""" +def open_ceus_run(run_dir: Path, delivery_dir: Path, repo_root: Path) -> CeusRun: + """核对运行目录(存在、不是 `.partial`、有 failed.csv 和快照),读 failed.csv。 + + 快照里 `ceus.yaml` 的 `paths.delivery_dir`(相对路径相对 `repo_root`,同 ceus_quant)与这次 + build 的 `delivery_dir` 规范化后不是同一路径,说明运行不是在这份交付包上跑的,报错停下。 + """ if run_dir.name.endswith(".partial"): raise ValueError(f"CEUS 运行目录是没跑完的 .partial:{run_dir}") if not run_dir.is_dir(): raise FileNotFoundError(f"CEUS 运行目录不存在:{run_dir}") if not (run_dir / "snapshot" / "snapshot.json").is_file(): raise FileNotFoundError(f"CEUS 运行目录缺快照 snapshot/snapshot.json:{run_dir}") + ceus_yaml = run_dir / "snapshot" / "ceus.yaml" + ran_on = _typed( + yaml.safe_load(ceus_yaml.read_text(encoding="utf-8"))["paths"]["delivery_dir"], str + ) + if (repo_root / ran_on).resolve() != delivery_dir.resolve(): + raise ValueError( + f"CEUS 运行不是在这份交付包上跑的:{ceus_yaml} 的 paths.delivery_dir 是 {ran_on}," + f"build 的 delivery_dir 是 {delivery_dir}" + ) with (run_dir / "failed.csv").open(encoding="utf-8", newline="") as handle: rows = list(csv.reader(handle)) if not rows or rows[0] != _FAILED_HEADER: @@ -142,6 +159,16 @@ def case_ceus(run: CeusRun, case: CaseRecord) -> LibraryCeus | SkipRecord: summary = json.loads((case_dir / "summary.json").read_text(encoding="utf-8")) if summary["video_id"] != case.video_id: raise ValueError(f"summary.json 的视频号与所在目录不一致:{case_dir}") + if _typed(summary["video_sha256"], str) != hashlib.sha256(case.video.read_bytes()).hexdigest(): + raise ValueError( + f"CEUS 运行不是在这份交付包上跑的:{case.video} 的 SHA-256 与 " + f"{case_dir / 'summary.json'} 的 video_sha256 不同" + ) + if _typed(summary["injection_offset_s"], float) != case.bolus_start_s: + raise ValueError( + f"CEUS 运行不是在这份交付包上跑的:交付包的推注时间 {case.bolus_start_s} 秒与 " + f"{case_dir / 'summary.json'} 的 injection_offset_s 不同" + ) tic = tic_view(summary, case.video_id) located = [slot for slot in _SLOTS if tic["keyframes"][slot]["located"]] images = sorted(p.stem for p in (case_dir / "keyframes").iterdir()) diff --git a/src/cadx/ceus/__init__.py b/src/cadx/ceus/__init__.py index 3317642..8f26311 100644 --- a/src/cadx/ceus/__init__.py +++ b/src/cadx/ceus/__init__.py @@ -67,7 +67,6 @@ from cadx.ceus._geometry import ( from cadx.ceus._green_prompt import ( ColorPrompt, PromptObservation, - collect_prompt_observations, extract_color_prompt, prompt_color_mask, ) @@ -111,6 +110,7 @@ from cadx.ceus._split_screen import ( SplitDetectionError, detect_split_x, sample_frames_for_split, + sample_indices, split_views, view_rect, ) @@ -207,7 +207,6 @@ __all__ = [ "align_polygon_to_previous", "apply_dark_morphology", "apply_prompt_frame_position_constraints", - "collect_prompt_observations", "color_marks", "compute_lognormal_perfusion_metrics", "compute_tic_for_video", @@ -255,6 +254,7 @@ __all__ = [ "render_tic_png", "resample_closed_contour", "sample_frames_for_split", + "sample_indices", "score_dark_candidate", "select_prompt_frames", "select_reference_region", diff --git a/src/cadx/ceus/_contrast_side.py b/src/cadx/ceus/_contrast_side.py index 1a4a1bf..d7ed548 100644 --- a/src/cadx/ceus/_contrast_side.py +++ b/src/cadx/ceus/_contrast_side.py @@ -1,19 +1,19 @@ """第 2 步:判断造影画面在哪一半,并在那一半上选出提示帧。 哪一半是造影画面,各例不一定相同。绿框只画在造影那一半,所以两半各自按同样的办法 -(`collect_prompt_observations` + `select_prompt_frames`,参数相同)选提示帧:恰好一半选出 +(逐帧 `extract_color_prompt` + `select_prompt_frames`,参数相同)选提示帧:恰好一半选出 至少一帧,这一半就是造影画面;两半都选出或都选不出,这一例失败(`ContrastSideError`)。 旧代码按配置写死造影在左半,造影在右半的病例就找不到绿框。 """ -from collections.abc import Sequence +from collections.abc import Iterable from dataclasses import dataclass from typing import Literal from cadx.ceus._config import KeyframesConfig, PromptSourceConfig from cadx.ceus._errors import CeusCaseError from cadx.ceus._geometry import Frame, ViewRect, crop_frame -from cadx.ceus._green_prompt import collect_prompt_observations +from cadx.ceus._green_prompt import PromptObservation, extract_color_prompt from cadx.ceus._prompt_frames import PromptFrame, select_prompt_frames PROMPT_STEP = 2 @@ -42,28 +42,24 @@ class ContrastSide: prompt_frames: list[PromptFrame] -def _prompt_frames_in( - frames: Sequence[Frame], - view: ViewRect, - prompt_cfg: PromptSourceConfig, - keyframes_cfg: KeyframesConfig, -) -> list[PromptFrame]: - observations = collect_prompt_observations( - (crop_frame(frame, view) for frame in frames), prompt_cfg - ) - return select_prompt_frames(observations, keyframes_cfg) - - def locate_contrast_side( - frames: Sequence[Frame], + frames: Iterable[Frame], left_view: ViewRect, right_view: ViewRect, prompt_cfg: PromptSourceConfig, keyframes_cfg: KeyframesConfig, ) -> ContrastSide: - """两半各自选提示帧,恰好一半有就定为造影画面;否则抛 `ContrastSideError`。""" - left = _prompt_frames_in(frames, left_view, prompt_cfg, keyframes_cfg) - right = _prompt_frames_in(frames, right_view, prompt_cfg, keyframes_cfg) + """两半各自选提示帧,恰好一半有就定为造影画面;否则抛 `ContrastSideError`。 + + `frames` 只遍历一遍,逐帧在两半上各找一次绿框,可以直接传逐帧解码的生成器。 + """ + observations: tuple[list[PromptObservation], list[PromptObservation]] = ([], []) + for index, frame in enumerate(frames): + for found, view in zip(observations, (left_view, right_view), strict=True): + prompt = extract_color_prompt(crop_frame(frame, view), prompt_cfg) + found.append(PromptObservation(frame_index=index, prompt=prompt)) + left = select_prompt_frames(observations[0], keyframes_cfg) + right = select_prompt_frames(observations[1], keyframes_cfg) if left and not right: return ContrastSide("left", left_view, right_view, left) if right and not left: diff --git a/src/cadx/ceus/_green_prompt.py b/src/cadx/ceus/_green_prompt.py index 8a01438..a22e220 100644 --- a/src/cadx/ceus/_green_prompt.py +++ b/src/cadx/ceus/_green_prompt.py @@ -3,21 +3,20 @@ - `prompt_color_mask`:TIC 采样时要从亮度统计里剔除医生画上去的线,这里给出那组颜色的掩膜, 受 `exclude_prompt_color` 开关和 `prompt_exclude_dilate` 膨胀核控制; - `extract_color_prompt`:HSV 范围内的像素 → 抹掉顶部文字带 → 形态学闭运算加膨胀把断开的描边 - 连起来 → 外轮廓 → 按外接框面积 w×h 过滤、取最大者 → 实心填充成掩膜;没有候选返回 None; -- `collect_prompt_observations`:逐帧调上一个函数,没检出的帧用 None 占位,帧号与下标对齐。 + 连起来 → 外轮廓 → 按外接框面积 w×h 过滤、取最大者 → 实心填充成掩膜;没有候选返回 None。 颜色来自参数:喂绿框那组配置检绿框,喂紫圈那组配置(`PanSourceConfig`)检紫圈。帧是已经切出的 一半画面,返回的外接框和掩膜都是这一半的局部坐标。 """ -from collections.abc import Iterable, Sequence +from collections.abc import Sequence from dataclasses import dataclass from typing import Protocol import cv2 import numpy as np -from cadx.ceus._config import PromptSourceConfig, TicAnalysisConfig +from cadx.ceus._config import TicAnalysisConfig from cadx.ceus._geometry import AnyArray, Frame, Mask, rect_from_values # 形态学结构元:5×5,闭运算与膨胀各一次。 @@ -174,13 +173,3 @@ def extract_color_prompt( contour_area=best.contour_area, candidate_count=candidate_count, ) - - -def collect_prompt_observations( - source_frames: Iterable[Frame], prompt_cfg: PromptSourceConfig -) -> list[PromptObservation]: - """逐帧找绿框,交出与帧序列等长的观测列表;没检出的帧 `prompt` 为 None。""" - return [ - PromptObservation(frame_index=index, prompt=extract_color_prompt(frame, prompt_cfg)) - for index, frame in enumerate(source_frames) - ] diff --git a/src/cadx/ceus/_pipeline.py b/src/cadx/ceus/_pipeline.py index 5782255..3b69bd2 100644 --- a/src/cadx/ceus/_pipeline.py +++ b/src/cadx/ceus/_pipeline.py @@ -1,8 +1,8 @@ """ceus 对外的两段接口,中间由 `scripts/ceus_quant.py` 调 segment 做第 3 步分割。 - `prepare_segmentation`(第 1–2 步):读视频、找分屏线、配准两半、判断造影画面在哪一半、选提示帧, - 交出要交给分割的那一半(灰阶那一半,`target_view`)的逐帧图像和提示帧掩膜(`target_view` - 局部坐标)。 + 交出提示帧掩膜(`target_view` 局部坐标);要交给分割的那一半(灰阶那一半,`target_view`)的 + 逐帧图像由 `target_frames()` 重新解码交出。 - `quantify_case`(第 4–9 步):拿分割得到的逐帧掩膜(`target_view` 局部坐标)整理成平滑轮廓并与 绿框融合、挪回造影那一半,找暗区、选背景圆、逐帧采 TIC、拟合、定位 7 个关键帧、算灌注指标。 @@ -47,7 +47,7 @@ from cadx.ceus._polygons import ( ) from cadx.ceus._reference_region import select_reference_region from cadx.ceus._registration import HalfRegistration, register_halves -from cadx.ceus._split_screen import split_views +from cadx.ceus._split_screen import sample_indices, split_views from cadx.ceus._tic import CURVE_NAMES, TicResult, compute_tic_for_video from cadx.ceus._tic_curves import TicCurves, from_tic_result from cadx.ceus._tic_fitting import LognormalFitResult, evaluate_lognormal @@ -86,10 +86,10 @@ class SegmentationInput: `offset`:造影那一半局部坐标 → `target_view` 局部坐标的平移 `(dx, dy)`,取互相关的估计; - `contrast`:哪一半是造影画面、两个视图矩形(整帧坐标)、提示帧(造影那一半局部坐标的 绿框掩膜); - - `target_frames`:要交给分割的那一半(`contrast.target_view`,灰阶画面)的逐帧 BGR 图像, - 与帧号一一对应; - `prompt_masks`:提示帧帧号 → 这一帧绿框挪到 `target_view` 局部坐标后的布尔掩膜; - `first_contrast_frame`:造影那一半的第一帧(选背景圆用)。 + + 要交给分割的那一半画面不存在这里,由 `target_frames()` 重新解码逐帧交出。 """ video_path: Path @@ -101,10 +101,14 @@ class SegmentationInput: registration: HalfRegistration offset: Offset contrast: ContrastSide - target_frames: list[Frame] prompt_masks: dict[int, Mask] first_contrast_frame: Frame + def target_frames(self) -> Iterator[Frame]: + """重新解码,逐帧交出分割那一半(`contrast.target_view`,灰阶画面)的 BGR 图像。""" + for _, frame in _checked_frames(VideoFile(self.video_path), self.frame_count): + yield crop_frame(frame, self.contrast.target_view) + @dataclass(frozen=True) class QuantResult: @@ -129,9 +133,25 @@ class QuantResult: pan_detected_fraction: float +def _checked_frames(video: VideoFile, frame_count: int) -> Iterator[tuple[int, Frame]]: + """逐帧解码,遍历完核对解码帧数与容器记录的一致。""" + decoded = 0 + for index, frame in video.iter_frames(): + decoded += 1 + yield index, frame + if decoded != frame_count: + raise VideoDecodeError( + f"解码帧数 {decoded} 与容器记录的帧数 {frame_count} 不一致:{video.path}" + ) + + def prepare_segmentation(video_path: Path, cfg: CeusConfig) -> SegmentationInput: """第 1–2 步:解码、分屏、配准两半、找绿框定造影那一半、选提示帧,交出分割的输入。 + 不把整段视频留在内存里,逐帧解码两遍:第一遍只留找分屏线和配准用的抽样帧;第二遍逐帧在两半 + 上找绿框。抽样帧号与 `split_views`、`register_halves` 自己抽的相同,所以结果与把整段帧交给 + 它们一样。分割那一半的画面由 `SegmentationInput.target_frames()` 到用时再解码。 + 视频读不了或读出的不合规抛 `VideoDecodeError`、找不到可信分屏线抛 `SplitDetectionError`、两半 配准失败(含两种方法不一致)抛 `RegistrationError`(都是第 1 步),两半都有或都没有绿框抛 `ContrastSideError`(第 2 步),都是 `CeusCaseError`。 @@ -139,22 +159,32 @@ def prepare_segmentation(video_path: Path, cfg: CeusConfig) -> SegmentationInput video = VideoFile(video_path) fps, frame_count, width, height = video.probe() frame_times_s = video.frame_times_s() - full_frames = [frame for _, frame in video.iter_frames()] - if len(full_frames) != frame_count: - raise VideoDecodeError( - f"解码帧数 {len(full_frames)} 与容器记录的帧数 {frame_count} 不一致:{video_path}" - ) - split_x, left_view, right_view = split_views(full_frames, cfg.views) + split_at = set(sample_indices(frame_count, cfg.views.split_sample_frames)) + register_at = set(sample_indices(frame_count, cfg.registration.sample_frames)) + split_frames: list[Frame] = [] + register_frames: list[Frame] = [] + for index, frame in _checked_frames(video, frame_count): + if index in split_at: + split_frames.append(frame) + if index in register_at: + register_frames.append(frame) + split_x, left_view, right_view = split_views(split_frames, cfg.views) registration = register_halves( - full_frames, left_view, right_view, cfg.registration, cfg.prompt_source, cfg.pan_source + register_frames, left_view, right_view, cfg.registration, cfg.prompt_source, cfg.pan_source ) + first_frame = split_frames[0] # 抽样帧号从 0 起,第 0 帧总在里面 + del split_frames, register_frames + contrast = locate_contrast_side( - full_frames, left_view, right_view, cfg.prompt_source, cfg.keyframes + (frame for _, frame in _checked_frames(video, frame_count)), + left_view, + right_view, + cfg.prompt_source, + cfg.keyframes, ) offset = registration.contrast_to_target(contrast.side) - target_frames = [crop_frame(frame, contrast.target_view).copy() for frame in full_frames] - first_contrast_frame = crop_frame(full_frames[0], contrast.contrast_view).copy() + first_contrast_frame = crop_frame(first_frame, contrast.contrast_view).copy() prompt_masks = { prompt.frame_index: map_source_mask_to_target(prompt.mask, offset, contrast.target_view) for prompt in contrast.prompt_frames @@ -169,7 +199,6 @@ def prepare_segmentation(video_path: Path, cfg: CeusConfig) -> SegmentationInput registration=registration, offset=offset, contrast=contrast, - target_frames=target_frames, prompt_masks=prompt_masks, first_contrast_frame=first_contrast_frame, ) diff --git a/src/cadx/ceus/_split_screen.py b/src/cadx/ceus/_split_screen.py index fa3c2ef..ae4d2a5 100644 --- a/src/cadx/ceus/_split_screen.py +++ b/src/cadx/ceus/_split_screen.py @@ -162,15 +162,19 @@ def view_rect( return rect +def sample_indices(frame_count: int, sample_count: int) -> list[int]: + """沿时间轴均匀抽 `sample_count` 个帧号(去重、升序);帧数不超过它时全用。""" + if sample_count <= 0 or frame_count <= sample_count: + return list(range(frame_count)) + indices = np.linspace(0, frame_count - 1, num=sample_count, dtype=np.int32) + return sorted({int(index) for index in indices.tolist()}) + + def sample_frames_for_split(frames: Sequence[Frame], sample_count: int) -> list[Frame]: - """沿时间轴均匀抽 `sample_count` 帧(下标去重);帧数不超过它时全用。""" + """按 `sample_indices` 均匀抽帧;已经抽好的帧(不多于 `sample_count`)原样全用。""" if not frames: raise SplitDetectionError("帧序列为空,无法找分屏线。") - if sample_count <= 0 or len(frames) <= sample_count: - return list(frames) - indices = np.linspace(0, len(frames) - 1, num=sample_count, dtype=np.int32) - unique_indices = sorted({int(index) for index in indices.tolist()}) - return [frames[index] for index in unique_indices] + return [frames[index] for index in sample_indices(len(frames), sample_count)] def split_views(frames: Sequence[Frame], views_cfg: ViewsConfig) -> tuple[int, ViewRect, ViewRect]: diff --git a/src/cadx/runs/__init__.py b/src/cadx/runs/__init__.py index ca68e75..b7d9a38 100644 --- a/src/cadx/runs/__init__.py +++ b/src/cadx/runs/__init__.py @@ -1,5 +1,11 @@ """runs:构建号怎么取、运行快照记什么。会诊运行目录等会诊那一步再加。""" -from cadx.runs._snapshot import GitState, make_build_id, read_git_state, write_snapshot +from cadx.runs._snapshot import ( + GitState, + finish_snapshot, + make_build_id, + read_git_state, + write_snapshot, +) -__all__ = ["GitState", "make_build_id", "read_git_state", "write_snapshot"] +__all__ = ["GitState", "finish_snapshot", "make_build_id", "read_git_state", "write_snapshot"] diff --git a/src/cadx/runs/_snapshot.py b/src/cadx/runs/_snapshot.py index 6a37e0c..7d828cc 100644 --- a/src/cadx/runs/_snapshot.py +++ b/src/cadx/runs/_snapshot.py @@ -1,4 +1,9 @@ -"""构建号与运行快照:一次构建用的是哪份代码、哪份配置、哪个环境。""" +"""构建号与运行快照:一次构建用的是哪份代码、哪份配置、哪个环境。 + +快照分两次写:开跑时 `write_snapshot` 记下代码、配置原文和环境(运行期间改了配置文件也不影响 +快照),收尾时 `finish_snapshot` 另写 `finished.json`,只记结束时间和运行结束才知道的值。两次 +都只新建文件,不改已写的。 +""" import hashlib import json @@ -55,7 +60,6 @@ def write_snapshot( build_id: str, git: GitState, started_at: datetime, - finished_at: datetime, uv_lock: Path, configs: Mapping[str, Path], ) -> None: @@ -73,9 +77,22 @@ def write_snapshot( "python_version": sys.version, "uv_lock_sha256": _sha256(uv_lock), "started_at": started_at.isoformat(), - "finished_at": finished_at.isoformat(), "config_sha256": copied, } (snapshot_dir / "snapshot.json").write_text( json.dumps(snapshot, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" ) + + +def finish_snapshot( + snapshot_dir: Path, finished_at: datetime, record: Mapping[str, object] +) -> None: + """收尾时在快照目录里新建 `finished.json`:结束时间加 `record`(运行结束才知道的值)。 + + `finished.json` 已存在(同一快照收两次尾)或 `record` 里有 `finished_at` 键都报错。 + """ + if "finished_at" in record: + raise ValueError("record 里不能有 finished_at,结束时间由参数 finished_at 给") + content = {"finished_at": finished_at.isoformat(), **record} + with (snapshot_dir / "finished.json").open("x", encoding="utf-8") as handle: + handle.write(json.dumps(content, ensure_ascii=False, indent=2) + "\n") diff --git a/src/cadx/segment/_medsam2.py b/src/cadx/segment/_medsam2.py index 1fd68a1..d494f7e 100644 --- a/src/cadx/segment/_medsam2.py +++ b/src/cadx/segment/_medsam2.py @@ -4,13 +4,19 @@ `000000.jpg` 起的 jpg 序列(质量 95)写进一个目录交给 `init_state`;每个提示帧用 `add_new_mask` 加进同一个对象;从最早的提示帧向后传播到末帧,再从它向前传播到第 0 帧;logits > 0 为病灶。 现在在同一进程里调用,jpg 临时目录由这里建、用完即删。 + +`init_state` 读 jpg 时换成 `_load_jpg_frames`:MedSAM2 原来的读法先在 CPU 上建整段 float32 张量 +(N×3×512×512,2288 帧约 7.2GB)再整个搬上显卡,这份 CPU 张量是整个流程的内存峰值。换上的读法 +直接在存放帧的设备上建张量、逐帧写入,算术与原来逐位相同(MedSAM2 源码不改)。 """ +import inspect import os import shutil import subprocess import tempfile -from collections.abc import Mapping, Sequence +from collections.abc import Generator, Iterable, Mapping +from contextlib import contextmanager from pathlib import Path from typing import Any @@ -58,33 +64,25 @@ def cuda_peak_memory_bytes(device: str) -> int | None: return int(torch.cuda.max_memory_allocated(device)) -def _check_inputs( - frames: Sequence[Frame], prompt_masks: Mapping[int, Mask], cfg: SegmentConfig -) -> tuple[int, int]: - """校验输入并返回帧的 (高, 宽)。""" +def _check_files(cfg: SegmentConfig) -> None: for path in (cfg.medsam2_dir / "sam2" / cfg.sam2_config, cfg.checkpoint): if not path.is_file(): raise FileNotFoundError(f"MedSAM2 缺少文件:{path}") - if len(frames) == 0: - raise ValueError("没有要分割的帧") - height, width = frames[0].shape[:2] - for index, frame in enumerate(frames): - if frame.dtype != np.uint8 or frame.shape != (height, width, 3): - raise ValueError( - f"第 {index} 帧应为 ({height}, {width}, 3) 的 uint8 BGR 图像," - f"收到 {frame.shape} {frame.dtype}" - ) + + +def _check_prompts( + prompt_masks: Mapping[int, Mask], frame_count: int, height: int, width: int +) -> None: if len(prompt_masks) == 0: raise ValueError("至少需要一个提示帧") for index, mask in prompt_masks.items(): - if not 0 <= index < len(frames): - raise ValueError(f"提示帧帧号 {index} 不在 [0, {len(frames)}) 内") + if not 0 <= index < frame_count: + raise ValueError(f"提示帧帧号 {index} 不在 [0, {frame_count}) 内") if mask.dtype != np.bool_ or mask.shape != (height, width): raise ValueError( f"提示帧 {index} 的掩码应为 ({height}, {width}) 的布尔数组," f"收到 {mask.shape} {mask.dtype}" ) - return height, width def _check_device(device: str) -> None: @@ -110,70 +108,139 @@ def _check_sam2_source(medsam2_dir: Path) -> None: raise RuntimeError(f"环境里的 sam2 来自 {loaded},与配置的 {expected} 不一致") -def _write_jpgs(frames: Sequence[Frame], frames_dir: Path) -> None: +def _write_jpgs(frames: Iterable[Frame], frames_dir: Path) -> tuple[int, int, int]: + """逐帧写 jpg(帧不留在内存里),返回(帧数, 高, 宽);各帧须是同一尺寸的 uint8 BGR 图像。""" + count, shape = 0, None for index, frame in enumerate(frames): + shape = frame.shape if shape is None else shape + if frame.dtype != np.uint8 or len(shape) != 3 or frame.shape != shape or shape[2] != 3: + raise ValueError( + f"第 {index} 帧应为与第 0 帧同尺寸 {shape} 的 uint8 BGR 图像," + f"收到 {frame.shape} {frame.dtype}" + ) path = frames_dir / f"{index:06d}.jpg" if not cv2.imwrite(str(path), frame, [cv2.IMWRITE_JPEG_QUALITY, _JPEG_QUALITY]): raise RuntimeError(f"写分割用的 jpg 失败:{path}") + count += 1 + if shape is None: + raise ValueError("没有要分割的帧") + return count, shape[0], shape[1] + + +def _load_jpg_frames( + *, + video_path: str, + image_size: int, + offload_video_to_cpu: bool, + async_loading_frames: bool, + compute_device: torch.device, +) -> tuple[torch.Tensor, int, int]: + """代替 `sam2.utils.misc.load_video_frames` 读 `_write_jpgs` 写的目录,参数和返回值与它相同。 + + 照原函数 `load_video_frames_from_jpg_images`:同一个 `_load_img_as_tensor`(float64)逐帧写进 + float32 张量,再减均值、除标准差(均值、标准差取原函数的参数默认值)。不同只在张量直接建在 + 存放设备上(`offload_video_to_cpu` 为否时是显卡),不先在 CPU 上建整段再搬:float64 转 float32 + 是同一种舍入,归一化原来也在这个设备上做,结果逐位相同。 + """ + from sam2.utils.misc import ( # pyright: ignore[reportMissingTypeStubs] + _load_img_as_tensor, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] + load_video_frames_from_jpg_images, # pyright: ignore[reportUnknownVariableType] + ) + + if async_loading_frames: + raise ValueError("_load_jpg_frames 不支持 async_loading_frames") + defaults = inspect.signature(load_video_frames_from_jpg_images).parameters # pyright: ignore[reportUnknownArgumentType] + device = torch.device("cpu") if offload_video_to_cpu else compute_device + img_mean = torch.tensor(defaults["img_mean"].default, dtype=torch.float32)[:, None, None] + img_std = torch.tensor(defaults["img_std"].default, dtype=torch.float32)[:, None, None] + names = sorted(name for name in os.listdir(video_path) if name.endswith(".jpg")) + if not names: + raise RuntimeError(f"分割用的 jpg 目录是空的:{video_path}") + images = torch.zeros(len(names), 3, image_size, image_size, dtype=torch.float32, device=device) + height = width = 0 + for index, name in enumerate(names): + path = os.path.join(video_path, name) + images[index], height, width = _load_img_as_tensor(path, image_size) # pyright: ignore[reportUnknownVariableType] + images -= img_mean.to(device) + images /= img_std.to(device) + return images, height, width + + +@contextmanager +def _jpg_frames_on_device() -> Generator[None]: + """`init_state` 期间把 sam2 视频预测器读帧的函数换成 `_load_jpg_frames`,退出时换回。""" + import sam2.sam2_video_predictor as video_predictor # pyright: ignore[reportMissingTypeStubs] + + original = video_predictor.load_video_frames # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] + video_predictor.load_video_frames = _load_jpg_frames + try: + yield + finally: + video_predictor.load_video_frames = original def _to_mask(mask_logits: torch.Tensor) -> Mask: return np.squeeze((mask_logits > 0.0).detach().cpu().numpy()).astype(np.bool_) -def _propagate(predictor: Any, state: Any, first_prompt: int) -> dict[int, Mask]: - """从最早的提示帧向后传播到末帧,再向前传播到第 0 帧。""" - propagated: dict[int, Mask] = {} +def _propagate(predictor: Any, state: Any, first_prompt: int, masks: Mask) -> Mask: + """从最早的提示帧向后传播到末帧,再向前传播到第 0 帧,逐帧写进 `masks`(后写的覆盖先写的)。 + + 返回哪些帧写过。 + """ + written = np.zeros(masks.shape[0], dtype=np.bool_) + passes = ( + {"start_frame_idx": first_prompt}, + { + "start_frame_idx": first_prompt, + "max_frame_num_to_track": first_prompt + 1, + "reverse": True, + }, + ) with torch.inference_mode(): - for frame_idx, obj_ids, mask_logits in predictor.propagate_in_video( - state, start_frame_idx=first_prompt - ): - propagated[int(frame_idx)] = _to_mask(mask_logits[list(obj_ids).index(_OBJECT_ID)]) - for frame_idx, obj_ids, mask_logits in predictor.propagate_in_video( - state, - start_frame_idx=first_prompt, - max_frame_num_to_track=first_prompt + 1, - reverse=True, - ): - propagated[int(frame_idx)] = _to_mask(mask_logits[list(obj_ids).index(_OBJECT_ID)]) - return propagated + for kwargs in passes: + for frame_idx, obj_ids, mask_logits in predictor.propagate_in_video(state, **kwargs): + mask = _to_mask(mask_logits[list(obj_ids).index(_OBJECT_ID)]) + if mask.shape != masks.shape[1:]: + raise RuntimeError(f"MedSAM2 输出掩码形状 {mask.shape},应为 {masks.shape[1:]}") + masks[int(frame_idx)] = mask + written[int(frame_idx)] = True + return written def segment_video( - frames: Sequence[Frame], prompt_masks: Mapping[int, Mask], cfg: SegmentConfig + frames: Iterable[Frame], prompt_masks: Mapping[int, Mask], cfg: SegmentConfig ) -> Mask: """把提示帧的病灶掩码传播到每一帧。 - `frames` 是逐帧 BGR 图像(同一尺寸),`prompt_masks` 是提示帧帧号 → 与帧同尺寸的布尔掩码。 - 返回 `(帧数, 高, 宽)` 的布尔数组,与输入帧一一对应。 + `frames` 是逐帧 BGR 图像(同一尺寸),只遍历一遍、边读边写成 jpg,可以直接传逐帧解码的生成器; + `prompt_masks` 是提示帧帧号 → 与帧同尺寸的布尔掩码。返回 `(帧数, 高, 宽)` 的布尔数组,与输入帧 + 一一对应。 """ - height, width = _check_inputs(frames, prompt_masks, cfg) + _check_files(cfg) _check_device(cfg.device) _check_sam2_source(cfg.medsam2_dir) # sam2 没有类型标注,参数类型是 Unknown;返回的 predictor 按 Any 用。 from sam2.build_sam import build_sam2_video_predictor # pyright: ignore[reportUnknownVariableType] - predictor: Any = build_sam2_video_predictor( - config_file=cfg.sam2_config, ckpt_path=str(cfg.checkpoint), device=cfg.device - ) with tempfile.TemporaryDirectory(prefix="cadx-medsam2-") as tmp: frames_dir = Path(tmp) - _write_jpgs(frames, frames_dir) - state: Any = predictor.init_state( - str(frames_dir), - offload_video_to_cpu=cfg.offload_video_to_cpu, - offload_state_to_cpu=cfg.offload_state_to_cpu, + frame_count, height, width = _write_jpgs(frames, frames_dir) + _check_prompts(prompt_masks, frame_count, height, width) + predictor: Any = build_sam2_video_predictor( + config_file=cfg.sam2_config, ckpt_path=str(cfg.checkpoint), device=cfg.device ) + with _jpg_frames_on_device(): + state: Any = predictor.init_state( + str(frames_dir), + offload_video_to_cpu=cfg.offload_video_to_cpu, + offload_state_to_cpu=cfg.offload_state_to_cpu, + ) for index in sorted(prompt_masks): predictor.add_new_mask(state, frame_idx=index, obj_id=_OBJECT_ID, mask=prompt_masks[index]) - propagated = _propagate(predictor, state, min(prompt_masks)) - - missing = [index for index in range(len(frames)) if index not in propagated] + masks = np.zeros((frame_count, height, width), dtype=np.bool_) + written = _propagate(predictor, state, min(prompt_masks), masks) + missing = np.flatnonzero(~written).tolist() if missing: raise RuntimeError(f"MedSAM2 传播后缺 {len(missing)} 帧的掩码,前几帧:{missing[:5]}") - stack = np.stack([propagated[index] for index in range(len(frames))], axis=0) - if stack.shape != (len(frames), height, width): - raise RuntimeError( - f"MedSAM2 输出掩码形状 {stack.shape},应为 {(len(frames), height, width)}" - ) - return stack + return masks diff --git a/tests/fake_ceus_run.py b/tests/fake_ceus_run.py index c4d9941..54f343b 100644 --- a/tests/fake_ceus_run.py +++ b/tests/fake_ceus_run.py @@ -1,17 +1,21 @@ """测试用假 CEUS 运行目录:布局照 `scripts/ceus_quant.py` 写出的(failed.csv、snapshot/、 `<视频号>/summary.json`、`<视频号>/keyframes/<槽位>.png`),数值全部编造,不调显卡。 -summary.json 故意带上 tic.json 不该有的字段(视频号、分割缓存键、帧号、拟合参数、失败原因代码、 -推注时间、紫圈检出率),用来测 build 只按白名单挑字段。关键帧 PNG 带文本块和 EXIF。 +summary.json 故意带上 tic.json 不该有的字段(视频号、视频 SHA-256、分割缓存键、帧号、拟合参数、 +失败原因代码、推注时间、紫圈检出率),用来测 build 只按白名单挑字段。视频 SHA-256 和推注时间照 +`fake_delivery` 的假交付包(视频内容 `VIDEO_BYTES`、推注 5 秒)。关键帧 PNG 带文本块和 EXIF。 """ import csv +import hashlib import json from pathlib import Path from typing import Any from PIL import Image, PngImagePlugin +from fake_delivery import VIDEO_BYTES + SLOTS = ("B0", "D1", "D2", "D3", "D4", "F1", "F2") CURVES = ("roi_all", "roi_perf", "roi_pan", "background_circle") UNLOCATED_REASON = "拟合曲线上没有 D4 时刻(编造的原因)" @@ -68,6 +72,7 @@ def summary(video_id: str, located: tuple[str, ...]) -> dict[str, Any]: } return { "video_id": video_id, + "video_sha256": hashlib.sha256(VIDEO_BYTES).hexdigest(), "segment_cache_key": "f" * 64, "injection_offset_s": 5.0, "reference_valid": False, @@ -99,12 +104,16 @@ def keyframe_png(path: Path, shade: int) -> None: def write_run( run_dir: Path, + delivery_dir: Path, results: dict[str, tuple[str, ...]], failed: list[tuple[str, int, str]], ) -> Path: - """写一个已完成的运行目录:`results` 是视频号 → 定位到的槽位,`failed` 是 failed.csv 的行。""" + """写一个已完成的运行目录:`delivery_dir` 记进快照里的 ceus.yaml,`results` 是视频号 → 定位到的 + 槽位,`failed` 是 failed.csv 的行。""" (run_dir / "snapshot").mkdir(parents=True) (run_dir / "snapshot" / "snapshot.json").write_text("{}", encoding="utf-8") + ceus_yaml = f"paths:\n delivery_dir: {json.dumps(str(delivery_dir), ensure_ascii=False)}\n" + (run_dir / "snapshot" / "ceus.yaml").write_text(ceus_yaml, encoding="utf-8") with (run_dir / "failed.csv").open("w", encoding="utf-8", newline="") as handle: writer = csv.writer(handle) writer.writerow(("video_id", "step", "reason")) diff --git a/tests/fake_delivery.py b/tests/fake_delivery.py index 6afbee0..204d117 100644 --- a/tests/fake_delivery.py +++ b/tests/fake_delivery.py @@ -62,6 +62,8 @@ MR_PATH = "MR_agent/MR报告文本_40例.xlsx" # 编造的报告文字:不含亚型名、结论段标题、性别字、「数字岁」和复核措辞。 ULTRASOUND_TEXT = "胰体部见一无回声区,边界清晰,壁薄。造影后囊壁呈等增强,内部未见增强。" +# 假交付包里每段 CEUS 视频文件的内容(不是真视频)。 +VIDEO_BYTES = b"not a real video" MR_TEXT = "胰体部见一类圆形长T2信号灶,边界清楚,内见细分隔;增强后分隔轻度强化。" # 测试用的样例映射表:含义是编的,只用来验证解码路径。 @@ -204,7 +206,7 @@ class FakeDelivery: rows = case.case_report_rows or [list(ULTRASOUND_HEADER), row] _write_workbook(folder / name, case.case_report_sheet, rows) else: - (folder / name).write_bytes(b"not a real video") + (folder / name).write_bytes(VIDEO_BYTES) (root / "交付说明_40例_20260908.md").write_text("说明不读", encoding="utf-8") return root diff --git a/tests/test_build_ceus.py b/tests/test_build_ceus.py index d016f99..e60dea7 100644 --- a/tests/test_build_ceus.py +++ b/tests/test_build_ceus.py @@ -76,6 +76,7 @@ def _key_tree(node: Any) -> Any: def test_c15_skips_failed_and_missing_cases(tmp_path: Path) -> None: run = write_run( tmp_path / "runs" / RUN_ID, + tmp_path / "delivery", {"U001": LOCATED, "U004": LOCATED}, [("U002", 2, "两半都没有绿框")], ) @@ -95,7 +96,11 @@ def test_c15_skips_failed_and_missing_cases(tmp_path: Path) -> None: def test_tic_json_has_exactly_the_listed_fields(tmp_path: Path) -> None: - result = _build(tmp_path, 1, write_run(tmp_path / "runs" / RUN_ID, {"U001": LOCATED}, [])) + result = _build( + tmp_path, + 1, + write_run(tmp_path / "runs" / RUN_ID, tmp_path / "delivery", {"U001": LOCATED}, []), + ) (alias,) = list_aliases(result.library_dir) tic = read_ultrasound(result.library_dir, alias).tic perfusion = { @@ -126,7 +131,7 @@ def test_tic_json_has_exactly_the_listed_fields(tmp_path: Path) -> None: def test_keyframes_metadata_dropped_and_pixels_kept(tmp_path: Path) -> None: - run = write_run(tmp_path / "runs" / RUN_ID, {"U001": LOCATED}, []) + run = write_run(tmp_path / "runs" / RUN_ID, tmp_path / "delivery", {"U001": LOCATED}, []) result = _build(tmp_path, 1, run) (alias,) = list_aliases(result.library_dir) route = read_ultrasound(result.library_dir, alias) @@ -159,7 +164,9 @@ def test_null_ceus_run_has_no_ceus_files(tmp_path: Path) -> None: def test_scan_finds_nothing_then_hits_video_id_in_tic_json(tmp_path: Path) -> None: - run = write_run(tmp_path / "runs" / RUN_ID, {"U001": LOCATED, "U002": LOCATED}, []) + run = write_run( + tmp_path / "runs" / RUN_ID, tmp_path / "delivery", {"U001": LOCATED, "U002": LOCATED}, [] + ) result = _build(tmp_path, 2, run) forbidden = ["P001", "P002", "U001", "U002", "U001_roi2.mp4", "U002_roi2.mp4", "U999"] assert scan_for_leaks(result.library_dir, forbidden, ("name",)) == [] @@ -173,7 +180,7 @@ def test_scan_finds_nothing_then_hits_video_id_in_tic_json(tmp_path: Path) -> No def test_reason_with_video_id_stops_build(tmp_path: Path) -> None: - run = write_run(tmp_path / "runs" / RUN_ID, {"U001": LOCATED}, []) + run = write_run(tmp_path / "runs" / RUN_ID, tmp_path / "delivery", {"U001": LOCATED}, []) content = summary("U001", LOCATED) content["anchor_frames"]["D4"]["reason"] = "U001 的 D4 不在视频里" (run / "U001" / "summary.json").write_text(json.dumps(content), encoding="utf-8") @@ -183,7 +190,9 @@ def test_reason_with_video_id_stops_build(tmp_path: Path) -> None: def test_keyframe_that_cannot_be_reencoded_skips_case_as_c14(tmp_path: Path) -> None: - run = write_run(tmp_path / "runs" / RUN_ID, {"U001": LOCATED, "U002": LOCATED}, []) + run = write_run( + tmp_path / "runs" / RUN_ID, tmp_path / "delivery", {"U001": LOCATED, "U002": LOCATED}, [] + ) Image.new("I;16", (4, 4)).save(run / "U002" / "keyframes" / "D2.png", format="PNG") result = _build(tmp_path, 2, run) assert [r.case_no for r in read_aliases(result.answers_dir)] == ["001"] @@ -199,22 +208,72 @@ def _assert_stops(tmp_path: Path, run: Path, error: type[Exception]) -> None: assert not (tmp_path / "library").exists() and not (tmp_path / "answers").exists() +def test_video_changed_since_run_stops(tmp_path: Path) -> None: + """交付包里的视频内容和 CEUS 运行时记的 SHA-256 不同:运行不是在这份交付包上跑的,整次停下。""" + run = write_run(tmp_path / "runs" / RUN_ID, tmp_path / "delivery", {"U001": LOCATED}, []) + root = make_delivery(1).write(tmp_path / "delivery") + (root / "CEUS_agent" / "001" / "U001_roi2.mp4").write_bytes(b"another video") + with pytest.raises(ValueError, match=r"U001_roi2\.mp4 的 SHA-256"): + run_build(write_config(tmp_path, root, 1, None, run), REPO_ROOT, NOW) + assert not (tmp_path / "library").exists() and not (tmp_path / "answers").exists() + + +@pytest.mark.parametrize( + ("delivery_bolus", "summary_offset", "message"), + [ + (7, 5.0, "推注时间 7 秒"), # 交付包的推注时间改了 + (5, 5, "应为 float"), # summary 里的推注时间不是浮点数 + ], +) +def test_injection_offset_mismatch_stops( + tmp_path: Path, delivery_bolus: int, summary_offset: object, message: str +) -> None: + run = write_run(tmp_path / "runs" / RUN_ID, tmp_path / "delivery", {"U001": LOCATED}, []) + content = summary("U001", LOCATED) + content["injection_offset_s"] = summary_offset + (run / "U001" / "summary.json").write_text(json.dumps(content), encoding="utf-8") + delivery = make_delivery(1) + delivery.case(1).ultrasound["推注开始时间s"] = delivery_bolus + root = delivery.write(tmp_path / "delivery") + with pytest.raises(ValueError, match=message): + run_build(write_config(tmp_path, root, 1, None, run), REPO_ROOT, NOW) + assert not (tmp_path / "library").exists() and not (tmp_path / "answers").exists() + + +def test_run_on_other_delivery_dir_stops(tmp_path: Path) -> None: + """快照里记的交付包路径与 build 的不同(规范化后比):整次停下;写法不同但同一路径照常装。""" + root = make_delivery(1).write(tmp_path / "delivery") + same = write_run( + tmp_path / "runs" / RUN_ID, tmp_path / "x" / ".." / "delivery", {"U001": LOCATED}, [] + ) + result = run_build(write_config(tmp_path, root, 1, None, same), REPO_ROOT, NOW) + assert len(list_aliases(result.library_dir)) == 1 + other = write_run(tmp_path / "runs" / "other", tmp_path / "delivery_old", {"U001": LOCATED}, []) + build_dir = tmp_path / "second" + build_dir.mkdir() + with pytest.raises(ValueError, match=r"paths\.delivery_dir"): + run_build(write_config(build_dir, root, 1, None, other), REPO_ROOT, NOW) + assert not (build_dir / "library").exists() and not (build_dir / "answers").exists() + + def test_run_dir_missing_stops(tmp_path: Path) -> None: _assert_stops(tmp_path, tmp_path / "runs" / RUN_ID, FileNotFoundError) def test_partial_run_dir_stops(tmp_path: Path) -> None: - run = write_run(tmp_path / "runs" / f"{RUN_ID}.partial", {"U001": LOCATED}, []) + run = write_run( + tmp_path / "runs" / f"{RUN_ID}.partial", tmp_path / "delivery", {"U001": LOCATED}, [] + ) _assert_stops(tmp_path, run, ValueError) def test_run_dir_without_failed_csv_stops(tmp_path: Path) -> None: - run = write_run(tmp_path / "runs" / RUN_ID, {"U001": LOCATED}, []) + run = write_run(tmp_path / "runs" / RUN_ID, tmp_path / "delivery", {"U001": LOCATED}, []) (run / "failed.csv").unlink() _assert_stops(tmp_path, run, FileNotFoundError) def test_run_dir_without_snapshot_stops(tmp_path: Path) -> None: - run = write_run(tmp_path / "runs" / RUN_ID, {"U001": LOCATED}, []) + run = write_run(tmp_path / "runs" / RUN_ID, tmp_path / "delivery", {"U001": LOCATED}, []) (run / "snapshot" / "snapshot.json").unlink() _assert_stops(tmp_path, run, FileNotFoundError) diff --git a/tests/test_build_outputs.py b/tests/test_build_outputs.py index 52c18cc..ec7a790 100644 --- a/tests/test_build_outputs.py +++ b/tests/test_build_outputs.py @@ -365,6 +365,7 @@ def test_answers_contents(tmp_path: Path) -> None: "review.csv", "skipped.csv", "snapshot/build.yaml", + "snapshot/finished.json", "snapshot/snapshot.json", ] diff --git a/tests/test_ceus_contrast_side.py b/tests/test_ceus_contrast_side.py index f568b65..8702c7b 100644 --- a/tests/test_ceus_contrast_side.py +++ b/tests/test_ceus_contrast_side.py @@ -72,6 +72,21 @@ def test_green_box_in_right_half(ceus_config: CeusConfig) -> None: assert abs(int(xs.max()) - 169) <= 3 +def test_frames_are_iterated_once(ceus_config: CeusConfig) -> None: + """帧只遍历一遍:传逐帧交出的生成器,结果与传列表相同。""" + frames = _frames([250]) + _, left, right = split_views(frames, ceus_config.views) + from_list = locate_contrast_side( + frames, left, right, ceus_config.prompt_source, ceus_config.keyframes + ) + from_generator = locate_contrast_side( + (frame for frame in frames), left, right, ceus_config.prompt_source, ceus_config.keyframes + ) + assert from_generator.side == from_list.side == "right" + assert [p.frame_index for p in from_generator.prompt_frames] == [0] + assert np.array_equal(from_generator.prompt_frames[0].mask, from_list.prompt_frames[0].mask) + + def test_green_box_in_both_halves_fails_case(ceus_config: CeusConfig) -> None: frames = _frames([40, 250]) _, left, right = split_views(frames, ceus_config.views) diff --git a/tests/test_ceus_green_prompt.py b/tests/test_ceus_green_prompt.py index 019fc95..81fe4ac 100644 --- a/tests/test_ceus_green_prompt.py +++ b/tests/test_ceus_green_prompt.py @@ -1,6 +1,6 @@ """B3 green_prompt 绿框提示检测单元测试。 -钉死四件事的关键性质(参考 ``references/ROI/roi_tracking_pipeline.py`` P:241-332/989-1004): +钉死这几件事的关键性质(参考 ``references/ROI/roi_tracking_pipeline.py`` P:241-332/989-1004): - ``extract_color_prompt``:合成在已知位置画绿色多边形的 BGR 左视图帧 → 返回该绿框 bbox(含端点 ``x2=x1+w-1`` / ``y2=y1+h-1``)+ 充满轮廓的布尔 mask; @@ -9,8 +9,6 @@ 绿框在带外则正常返回。 - ``min_area``:按 **bbox 面积 w×h**(findings §3.2,非轮廓面积)过滤—— bbox 面积 < min_area 的候选被丢弃。 -- ``collect_prompt_observations``:对多帧序列返回逐帧观测列表, - 无绿帧以 ``None`` 占位(``observation.prompt is None``)。 - ``prompt_color_mask``:TIC 端"排除绿框色"掩膜——``exclude_prompt_color=False`` 全 False; ``prompt_exclude_dilate=1`` 不膨胀(mask == 原始 inRange);增大膨胀核则 mask 更大。 @@ -24,10 +22,8 @@ from cadx.ceus import ( CeusConfig, ColorPrompt, PanSourceConfig, - PromptObservation, PromptSourceConfig, TicAnalysisConfig, - collect_prompt_observations, extract_color_prompt, prompt_color_mask, ) @@ -198,24 +194,6 @@ def test_min_area_picks_largest_bbox_candidate() -> None: assert abs(prompt.bbox_xyxy[1] - 60) <= 3 -# Phase 4: collect_prompt_observations —— 多帧观测列表(None 占位) -def test_collect_observations_none_placeholder_for_no_green_frames() -> None: - frames = [ - _frame_with_green_rect(120, 200, (40, 40, 120, 90)), # 有绿 - _gray_frame(120, 200), # 无绿 → None 占位 - _frame_with_green_rect(120, 200, (50, 40, 130, 95)), # 有绿 - ] - observations = collect_prompt_observations(frames, _cfg()) - assert len(observations) == 3 - assert all(isinstance(o, PromptObservation) for o in observations) - assert [o.frame_index for o in observations] == [0, 1, 2] - assert observations[0].prompt is not None - assert observations[1].prompt is None # 无绿帧占位 - assert observations[2].prompt is not None - assert observations[1].detected is False - assert observations[0].detected is True - - # Phase 5: prompt_color_mask —— TIC 端"排除绿框色"掩膜(受 tic_cfg 门控) def test_prompt_color_mask_marks_green_pixels() -> None: frame = _frame_with_green_rect(120, 200, (40, 30, 120, 90)) diff --git a/tests/test_ceus_pipeline.py b/tests/test_ceus_pipeline.py index c21353e..c354fde 100644 --- a/tests/test_ceus_pipeline.py +++ b/tests/test_ceus_pipeline.py @@ -24,14 +24,19 @@ from cadx.ceus import ( CeusConfig, NoValidContourError, SegmentationInput, + VideoFile, + crop_frame, keyframe_images, + locate_contrast_side, pack_tic_curves, pack_trajectory, prepare_segmentation, quant_summary, quantify_case, + register_halves, render_case_overlay, render_tic_png, + split_views, tic_plot_markers, unpack_tic_curves, unpack_trajectory, @@ -77,8 +82,9 @@ def test_prepare_segmentation_finds_contrast_half_and_maps_prompts( target = prepared.contrast.target_view assert (target.x1, target.x2) == (0, prepared.split_x) assert abs(prepared.split_x - SPLIT) <= 1 - assert len(prepared.target_frames) == FRAMES - assert all(frame.shape == (target.height, target.width, 3) for frame in prepared.target_frames) + target_frames = list(prepared.target_frames()) + assert len(target_frames) == FRAMES + assert all(frame.shape == (target.height, target.width, 3) for frame in target_frames) assert sorted(prepared.prompt_masks) == [p.frame_index for p in prepared.contrast.prompt_frames] assert 0 in prepared.prompt_masks xs = np.where(prepared.prompt_masks[0])[1] @@ -87,6 +93,37 @@ def test_prepare_segmentation_finds_contrast_half_and_maps_prompts( assert abs(int(xs.max()) - 169) <= 5 +def test_streaming_prepare_matches_whole_video_in_memory( + prepared: SegmentationInput, config: CeusConfig +) -> None: + """逐帧解码的第 1–2 步与「整段帧读进内存再交给各步」逐像素相同(改成流式之前就是后一种)。""" + frames = [frame for _, frame in VideoFile(prepared.video_path).iter_frames()] + split_x, left, right = split_views(frames, config.views) + registration = register_halves( + frames, left, right, config.registration, config.prompt_source, config.pan_source + ) + contrast = locate_contrast_side(frames, left, right, config.prompt_source, config.keyframes) + assert prepared.split_x == split_x + assert prepared.registration == registration + assert prepared.contrast.side == contrast.side + assert (prepared.contrast.contrast_view, prepared.contrast.target_view) == ( + contrast.contrast_view, + contrast.target_view, + ) + assert [p.frame_index for p in prepared.contrast.prompt_frames] == [ + p.frame_index for p in contrast.prompt_frames + ] + for ours, whole in zip(prepared.contrast.prompt_frames, contrast.prompt_frames, strict=True): + assert np.array_equal(ours.mask, whole.mask) + target_frames = list(prepared.target_frames()) + assert len(target_frames) == len(frames) + for ours, frame in zip(target_frames, frames, strict=True): + assert np.array_equal(ours, crop_frame(frame, contrast.target_view)) + assert np.array_equal( + prepared.first_contrast_frame, crop_frame(frames[0], contrast.contrast_view) + ) + + def test_quantify_case_end_to_end(prepared: SegmentationInput, config: CeusConfig) -> None: result = quantify_case(prepared, _lesion_masks(prepared), _INJECTION_OFFSET_S, config) diff --git a/tests/test_ceus_quant.py b/tests/test_ceus_quant.py index 181bf19..20fe47d 100644 --- a/tests/test_ceus_quant.py +++ b/tests/test_ceus_quant.py @@ -13,7 +13,7 @@ import csv import hashlib import importlib.util import json -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping from datetime import datetime from pathlib import Path from types import ModuleType @@ -88,11 +88,12 @@ class _CountingSegment: def __call__( self, - frames: Sequence[npt.NDArray[np.uint8]], + 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 @@ -170,7 +171,11 @@ def test_run_writes_outputs_and_failed_case( 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["cuda_peak_memory_bytes"] is None # device 是 cpu + 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 @@ -178,6 +183,37 @@ def test_run_writes_outputs_and_failed_case( 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]], diff --git a/tests/test_runs_snapshot.py b/tests/test_runs_snapshot.py new file mode 100644 index 0000000..abbe508 --- /dev/null +++ b/tests/test_runs_snapshot.py @@ -0,0 +1,45 @@ +"""runs 的快照收尾:`finish_snapshot` 只新建 finished.json,不改开跑时写的快照,也不收两次尾。""" + +import json +from datetime import datetime +from pathlib import Path + +import pytest + +from cadx.runs import GitState, finish_snapshot, write_snapshot + +_GIT = GitState(commit="a" * 40, short_commit="aaaaaaa", clean=True, diff_sha256="b" * 64) + + +def _snapshot(tmp_path: Path) -> Path: + lock = tmp_path / "uv.lock" + lock.write_text("lock", encoding="utf-8") + config = tmp_path / "x.yaml" + config.write_text("a: 1\n", encoding="utf-8") + snapshot_dir = tmp_path / "snapshot" + started = datetime(2026, 10, 16, 9, 0, 0).astimezone() + write_snapshot(snapshot_dir, "run-1", _GIT, started, lock, {"x.yaml": config}) + return snapshot_dir + + +def test_finish_writes_separate_file_and_keeps_snapshot(tmp_path: Path) -> None: + snapshot_dir = _snapshot(tmp_path) + before = (snapshot_dir / "snapshot.json").read_bytes() + finished = datetime(2026, 10, 16, 10, 0, 0).astimezone() + finish_snapshot(snapshot_dir, finished, {"segment_cache_hits": 3}) + assert (snapshot_dir / "snapshot.json").read_bytes() == before + record = json.loads((snapshot_dir / "finished.json").read_text(encoding="utf-8")) + assert record == {"finished_at": finished.isoformat(), "segment_cache_hits": 3} + + +def test_finish_twice_or_with_finished_at_in_record_raises(tmp_path: Path) -> None: + snapshot_dir = _snapshot(tmp_path) + now = datetime(2026, 10, 16, 10, 0, 0).astimezone() + with pytest.raises(ValueError, match="finished_at"): + finish_snapshot(snapshot_dir, now, {"finished_at": "x"}) + assert not (snapshot_dir / "finished.json").exists() + finish_snapshot(snapshot_dir, now, {}) + first = (snapshot_dir / "finished.json").read_bytes() + with pytest.raises(FileExistsError): + finish_snapshot(snapshot_dir, now, {"segment_cache_hits": 1}) + assert (snapshot_dir / "finished.json").read_bytes() == first diff --git a/tests/test_segment.py b/tests/test_segment.py index f03667a..bc41bde 100644 --- a/tests/test_segment.py +++ b/tests/test_segment.py @@ -93,3 +93,36 @@ def test_segment_video_on_synthetic_clip() -> None: assert masks.shape == (frame_count, height, width) assert masks.dtype == np.bool_ assert len(medsam2_commit(cfg.medsam2_dir)) == 40 + + +@pytest.mark.gpu +@pytest.mark.parametrize("offload_video_to_cpu", [False, True]) +def test_jpg_loader_matches_medsam2_loader_bitwise( + tmp_path: Path, offload_video_to_cpu: bool +) -> None: + """换上的读帧函数与 MedSAM2 原来的 `load_video_frames` 交出逐位相同的张量和原始尺寸。""" + import torch + from sam2.utils.misc import load_video_frames # pyright: ignore[reportUnknownVariableType] + + from cadx.segment._medsam2 import ( + _load_jpg_frames, # pyright: ignore[reportPrivateUsage] + _write_jpgs, # pyright: ignore[reportPrivateUsage] + ) + + frames, _ = _moving_disc_video(12, 120, 200, radius=18) + noise = np.random.default_rng(0).integers(0, 40, size=(12, 120, 200, 3), dtype=np.uint8) + _write_jpgs([frame + n for frame, n in zip(frames, noise, strict=True)], tmp_path) + kwargs: dict[str, Any] = { + "video_path": str(tmp_path), + "image_size": 512, + "offload_video_to_cpu": offload_video_to_cpu, + "async_loading_frames": False, + "compute_device": torch.device("cuda"), + } + expected, expected_h, expected_w = load_video_frames(**kwargs) # pyright: ignore[reportUnknownVariableType] + assert isinstance(expected, torch.Tensor) # 不异步读时原函数交出整段张量 + images, height, width = _load_jpg_frames(**kwargs) + assert (height, width) == (expected_h, expected_w) == (120, 200) + assert images.device == expected.device + assert images.dtype == torch.float32 + assert torch.equal(images, expected)