feat(ceus): 任务书 02 第三阶段——scripts/ceus_quant.py
- delivery 的 CaseRecord 加视频路径和推注时间（C04、C08 已检查过）。 - scripts/ceus_quant.py：读 ceus.yaml，跑检查全部通过的病例；runs/ceus/<运行号>/ 先写 .partial 再改名；第 1–6 步的单例失败记 failed.csv 后接着跑；终端只打印计数。 - 分割缓存 work/segment_cache/<键>/：键含视频 SHA-256、第 1–3 步参数、权重 SHA-256、MedSAM2 提交号。 - 视频解码错误改为第 1 步的单例失败（VideoDecodeError），不再停掉整次运行（审阅意见）。 - 暗区 kmeans 用 cv2 进程级随机数，结果受病例顺序影响；ceus 配置加必填 random_seed，每例第 4 步前设种子。 - segment 导出 cuda_peak_memory_bytes，供运行快照记录显存峰值。 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,293 @@
|
||||
"""CEUS 定量：对交付包里检查全部通过的病例逐例跑第 1–9 步，写进 `runs/ceus/<运行号>/`。
|
||||
|
||||
用法：uv run python scripts/ceus_quant.py --config config/ceus.yaml
|
||||
第 3 步（MedSAM2 分割）的结果按内容键缓存在 `<work_root>/segment_cache/<键>/masks.npz`，键相同就
|
||||
直接读缓存、不调显卡。某一例在第 1–6 步失败（`CeusCaseError`）记进 `failed.csv`，接着跑下一例；
|
||||
其余异常整次运行停下，留下 `.partial` 目录。终端只打印计数，不打印视频号。
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import hashlib
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from collections import Counter
|
||||
from collections.abc import Callable, Hashable, Sequence
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Any
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import yaml
|
||||
from pydantic import BaseModel, BeforeValidator, ConfigDict, Field
|
||||
|
||||
from cadx.ceus import (
|
||||
ANCHOR_SLOTS,
|
||||
CeusCaseError,
|
||||
CeusConfig,
|
||||
SegmentationInput,
|
||||
keyframe_images,
|
||||
pack_tic_curves,
|
||||
pack_trajectory,
|
||||
prepare_segmentation,
|
||||
quant_summary,
|
||||
quantify_case,
|
||||
render_case_overlay,
|
||||
render_tic_png,
|
||||
tic_plot_markers,
|
||||
)
|
||||
from cadx.delivery import CaseRecord, read_delivery
|
||||
from cadx.runs import 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]
|
||||
|
||||
# 第 1–2 步（`prepare_segmentation`）读的 ceus 段子段：views（分屏）、prompt_source（绿框颜色）、
|
||||
# keyframes（提示帧选择）。它们决定交给分割的帧和提示掩码，进分割缓存键。
|
||||
_STEP_1_2_CEUS_FIELDS = frozenset({"views", "prompt_source", "keyframes"})
|
||||
# segment 段里不改变分割结果的开关，不进缓存键；其余字段都进。
|
||||
_RESULT_NEUTRAL_SEGMENT_FIELDS = frozenset(
|
||||
{"device", "offload_video_to_cpu", "offload_state_to_cpu"}
|
||||
)
|
||||
_FAILED_HEADER = ("video_id", "step", "reason")
|
||||
|
||||
|
||||
def _reject_empty_path(value: object) -> object:
|
||||
# Path("") 会变成 "."（当前目录），空串在转换前就拦下。
|
||||
if value == "":
|
||||
raise ValueError("路径不能是空串")
|
||||
return value
|
||||
|
||||
|
||||
# YAML 里的路径是字符串；路径字段是唯一允许从字符串转换类型的地方。
|
||||
type _PathField = Annotated[Path, Field(strict=False), BeforeValidator(_reject_empty_path)]
|
||||
|
||||
|
||||
class PathsConfig(BaseModel):
|
||||
"""`paths` 段：交付包、运行目录的上级、中间文件目录。相对路径相对仓库根目录。"""
|
||||
|
||||
model_config = ConfigDict(strict=True, extra="forbid", frozen=True)
|
||||
|
||||
delivery_dir: _PathField
|
||||
runs_root: _PathField
|
||||
work_root: _PathField
|
||||
|
||||
|
||||
class QuantConfig(BaseModel):
|
||||
"""`ceus.yaml` 整份：`segment` 段交 segment 校验，`ceus` 段交 ceus 校验。"""
|
||||
|
||||
model_config = ConfigDict(strict=True, extra="forbid", frozen=True)
|
||||
|
||||
paths: PathsConfig
|
||||
segment: SegmentConfig
|
||||
ceus: CeusConfig
|
||||
|
||||
|
||||
class _UniqueKeyLoader(yaml.SafeLoader):
|
||||
"""同一层出现重复键就报错：PyYAML 默认让后一个悄悄覆盖前一个。"""
|
||||
|
||||
# 签名（含 deep 的默认值）照抄 PyYAML 的基类方法，覆盖时必须一致。
|
||||
def construct_mapping(self, node: yaml.MappingNode, deep: bool = False) -> dict[Hashable, Any]:
|
||||
seen: set[Hashable] = set()
|
||||
for key_node, _ in node.value:
|
||||
key = self.construct_object(key_node, deep=deep) # pyright: ignore[reportUnknownMemberType]
|
||||
if key in seen:
|
||||
raise ValueError(f"YAML 第 {key_node.start_mark.line + 1} 行重复的键：{key!r}")
|
||||
seen.add(key)
|
||||
return super().construct_mapping(node, deep=deep)
|
||||
|
||||
|
||||
def load_config(path: Path) -> QuantConfig:
|
||||
"""读并校验 `ceus.yaml`，路径保持文件里的写法（缓存键用这一份，不随仓库位置变）。"""
|
||||
text = path.read_text(encoding="utf-8")
|
||||
return QuantConfig.model_validate(yaml.load(text, Loader=_UniqueKeyLoader)) # noqa: S506 继承 SafeLoader
|
||||
|
||||
|
||||
def _resolved(config: QuantConfig) -> QuantConfig:
|
||||
"""相对路径换成仓库根目录下的绝对路径（绝对路径原样保留）。"""
|
||||
paths, segment = config.paths, config.segment
|
||||
return config.model_copy(
|
||||
update={
|
||||
"paths": paths.model_copy(
|
||||
update={
|
||||
"delivery_dir": REPO_ROOT / paths.delivery_dir,
|
||||
"runs_root": REPO_ROOT / paths.runs_root,
|
||||
"work_root": REPO_ROOT / paths.work_root,
|
||||
}
|
||||
),
|
||||
"segment": segment.model_copy(
|
||||
update={
|
||||
"medsam2_dir": REPO_ROOT / segment.medsam2_dir,
|
||||
"checkpoint": REPO_ROOT / segment.checkpoint,
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def file_sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
while chunk := handle.read(1 << 20):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def segment_cache_key(
|
||||
video_sha256: str, config: QuantConfig, checkpoint_sha256: str, commit: str
|
||||
) -> str:
|
||||
"""分割缓存键：视频、第 1–3 步参数、权重、MedSAM2 提交号拼成规范 JSON 后的 SHA-256。"""
|
||||
payload = {
|
||||
"video_sha256": video_sha256,
|
||||
"ceus": config.ceus.model_dump(mode="json", include=set(_STEP_1_2_CEUS_FIELDS)),
|
||||
"segment": config.segment.model_dump(
|
||||
mode="json", exclude=set(_RESULT_NEUTRAL_SEGMENT_FIELDS)
|
||||
),
|
||||
"checkpoint_sha256": checkpoint_sha256,
|
||||
"medsam2_commit": commit,
|
||||
}
|
||||
canonical = json.dumps(payload, sort_keys=True, ensure_ascii=False, separators=(",", ":"))
|
||||
return hashlib.sha256(canonical.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _masks(
|
||||
prepared: SegmentationInput, key: str, cache_root: Path, segment_cfg: SegmentConfig
|
||||
) -> tuple[npt.NDArray[np.bool_], bool]:
|
||||
"""读分割缓存；没有就调 MedSAM2，先写 `<键>.partial` 再改名。返回（掩码，是否命中缓存）。"""
|
||||
final = cache_root / key
|
||||
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)
|
||||
partial = cache_root / f"{key}.partial"
|
||||
# 上次中途停下留下的 .partial 会让这里报错：work/ 可以删掉重算，删掉它再跑。
|
||||
partial.mkdir(parents=True, exist_ok=False)
|
||||
np.savez_compressed(partial / "masks.npz", masks=masks)
|
||||
partial.rename(final)
|
||||
return masks, False
|
||||
|
||||
|
||||
def _write_png(path: Path, image: npt.NDArray[np.uint8]) -> None:
|
||||
ok, encoded = cv2.imencode(".png", image)
|
||||
if not ok:
|
||||
raise RuntimeError(f"关键帧编码 PNG 失败：{path}")
|
||||
path.write_bytes(encoded.tobytes())
|
||||
|
||||
|
||||
def _run_case(
|
||||
case: CaseRecord, cfg: QuantConfig, key_of: Callable[[str], str], run_dir: Path
|
||||
) -> tuple[tuple[str, ...], bool]:
|
||||
"""跑一例并写出产物，返回（定位到的关键帧槽位，分割是否命中缓存）。
|
||||
|
||||
`key_of` 把视频的 SHA-256 换成分割缓存键（其余几样跨病例不变）。
|
||||
"""
|
||||
prepared = prepare_segmentation(case.video, cfg.ceus)
|
||||
key = key_of(file_sha256(case.video))
|
||||
masks, hit = _masks(prepared, key, cfg.paths.work_root / "segment_cache", cfg.segment)
|
||||
result = quantify_case(prepared, masks, float(case.bolus_start_s), cfg.ceus)
|
||||
|
||||
case_dir = run_dir / case.video_id
|
||||
(case_dir / "keyframes").mkdir(parents=True, exist_ok=False)
|
||||
(case_dir / "trajectory.npz").write_bytes(pack_trajectory(result.trajectory))
|
||||
(case_dir / "tic_curves.npz").write_bytes(pack_tic_curves(result.curves))
|
||||
summary = {"video_id": case.video_id, "segment_cache_key": key, **quant_summary(result)}
|
||||
(case_dir / "summary.json").write_text(
|
||||
json.dumps(summary, ensure_ascii=False, indent=2, allow_nan=False) + "\n", encoding="utf-8"
|
||||
)
|
||||
markers = tic_plot_markers(result.anchors)
|
||||
(case_dir / "tic.png").write_bytes(render_tic_png(result.tic, markers=markers))
|
||||
(case_dir / "overlay.mp4").write_bytes(render_case_overlay(prepared, result))
|
||||
images = keyframe_images(prepared, result.anchors)
|
||||
for slot, image in images.items():
|
||||
_write_png(case_dir / "keyframes" / f"{slot}.png", image)
|
||||
return tuple(images), hit
|
||||
|
||||
|
||||
def main(argv: Sequence[str]) -> None:
|
||||
parser = argparse.ArgumentParser(description="CEUS 定量：逐例跑第 1–9 步")
|
||||
parser.add_argument("--config", type=Path, required=True, help="ceus.yaml 的路径")
|
||||
config_path: Path = parser.parse_args(argv).config
|
||||
started_at = datetime.now().astimezone()
|
||||
clock = time.monotonic()
|
||||
|
||||
written = load_config(config_path)
|
||||
cfg = _resolved(written)
|
||||
git = read_git_state(REPO_ROOT)
|
||||
run_id = make_build_id(started_at, git)
|
||||
run_dir = cfg.paths.runs_root / run_id
|
||||
partial = cfg.paths.runs_root / f"{run_id}.partial"
|
||||
for path in (run_dir, partial):
|
||||
if path.exists():
|
||||
raise FileExistsError(f"运行目录已存在，不覆盖：{path}")
|
||||
partial.mkdir(parents=True, exist_ok=False)
|
||||
|
||||
delivery = read_delivery(cfg.paths.delivery_dir, None)
|
||||
checkpoint_sha256 = file_sha256(cfg.segment.checkpoint)
|
||||
commit = medsam2_commit(cfg.segment.medsam2_dir)
|
||||
|
||||
# 缓存键用配置文件里的写法（路径不随仓库位置变）。
|
||||
def key_of(video_sha256: str) -> str:
|
||||
return segment_cache_key(video_sha256, written, checkpoint_sha256, commit)
|
||||
|
||||
failures: list[tuple[str, int, str]] = []
|
||||
located: Counter[str] = Counter()
|
||||
cache_hits = 0
|
||||
total = len(delivery.cases)
|
||||
for number, case in enumerate(delivery.cases, start=1):
|
||||
try:
|
||||
slots, hit = _run_case(case, cfg, key_of, partial)
|
||||
except CeusCaseError as error:
|
||||
failures.append((case.video_id, error.step, error.reason))
|
||||
print(f"[{number}/{total}] 失败在第 {error.step} 步", flush=True)
|
||||
continue
|
||||
located.update(slots)
|
||||
cache_hits += hit
|
||||
print(f"[{number}/{total}] 完成", flush=True)
|
||||
|
||||
with (partial / "failed.csv").open("w", encoding="utf-8", newline="") as handle:
|
||||
writer = csv.writer(handle)
|
||||
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},
|
||||
)
|
||||
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"
|
||||
)
|
||||
# POSIX 的 rename 会悄悄覆盖空目录，所以改名前再查一次。
|
||||
if run_dir.exists():
|
||||
raise FileExistsError(f"运行目录已存在，不覆盖：{run_dir}")
|
||||
partial.rename(run_dir)
|
||||
|
||||
by_step = Counter(step for _, step, _ in failures)
|
||||
print(f"运行目录：{run_dir}")
|
||||
print(f"病例总数（检查全部通过）：{total}")
|
||||
print(f"成功：{total - len(failures)}")
|
||||
failed_text = "，".join(f"第 {step} 步 {n}" for step, n in sorted(by_step.items()))
|
||||
print(f"按步分的失败数：{failed_text or '无'}")
|
||||
print("关键帧按槽位分的定位数：" + "，".join(f"{s} {located[s]}" for s in ANCHOR_SLOTS))
|
||||
print(f"分割缓存命中：{cache_hits}")
|
||||
print(f"耗时：{time.monotonic() - clock:.1f} 秒")
|
||||
peak_text = "不适用（非 CUDA 设备）" if peak is None else f"{peak / 2**20:.0f} MiB"
|
||||
print(f"显存峰值（torch 分配）：{peak_text}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(sys.argv[1:])
|
||||
Reference in New Issue
Block a user