fix(ceus): 降低长视频内存峰值,快照开跑即写,build 核对 CEUS 运行与交付包

- 内存: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>
This commit is contained in:
2026-10-16 09:01:44 -04:00
parent e165b32b9d
commit 1de88a5d0a
23 changed files with 561 additions and 202 deletions
+6 -4
View File
@@ -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-<git 短哈希>`,工作区不干净时末尾加 `-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`。
+5 -12
View File
@@ -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)
+32 -22
View File
@@ -7,6 +7,8 @@
第 3 步(MedSAM2 分割)的结果按内容键缓存在 `<work_root>/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():
+7 -3
View File
@@ -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,
+31 -4
View File
@@ -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())
+2 -2
View File
@@ -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",
+15 -19
View File
@@ -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:
+3 -14
View File
@@ -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)
]
+46 -17
View File
@@ -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,
)
+10 -6
View File
@@ -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]:
+8 -2
View File
@@ -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"]
+20 -3
View File
@@ -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")
+116 -49
View File
@@ -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,55 +108,129 @@ 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]
with tempfile.TemporaryDirectory(prefix="cadx-medsam2-") as tmp:
frames_dir = Path(tmp)
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 tempfile.TemporaryDirectory(prefix="cadx-medsam2-") as tmp:
frames_dir = Path(tmp)
_write_jpgs(frames, frames_dir)
with _jpg_frames_on_device():
state: Any = predictor.init_state(
str(frames_dir),
offload_video_to_cpu=cfg.offload_video_to_cpu,
@@ -166,14 +238,9 @@ def segment_video(
)
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
+12 -3
View File
@@ -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"))
+3 -1
View File
@@ -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
+67 -8
View File
@@ -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)
+1
View File
@@ -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",
]
+15
View File
@@ -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)
+1 -23
View File
@@ -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))
+39 -2
View File
@@ -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)
+39 -3
View File
@@ -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]],
+45
View File
@@ -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
+33
View File
@@ -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)