iomgaa
00ebbffb28
- 第 1 步分屏后在两半扇形区内用边缘图归一化互相关估 (dx, dy)，再用逐帧多块模板匹配的中位数独立核对； 两者各分量之差超过门槛（15 像素）这一例第 1 步失败。删掉「两个视图左上角相减」的实现与测试。 - 门槛依据：34 例中走完第 1 步的 29 例，差值中位 2、95 分位 6.2、最大 13；参数与门槛在同一批视频上调出。 - scripts/ceus_align_review.py：只做第 1 步，出扇形区棋盘图与边缘叠加图，混入错开 15–30 像素的 阴性对照、随机编号；计分时任一对照未被判「没对齐」即作废。 - ceus_quant 读 review.csv：不通过的病例第 1 步失败；表里没有这段视频或平移对不上即停。 - 配准与紫圈颜色参数进分割缓存键。 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
394 lines
16 KiB
Python
394 lines
16 KiB
Python
"""CEUS 定量：对交付包里检查全部通过的病例逐例跑第 1–9 步，写进 `runs/ceus/<运行号>/`。
|
||
|
||
用法：uv run python scripts/ceus_quant.py --config config/ceus.yaml
|
||
第 1 步算出两半配准后，按视频 SHA-256 查配准目视核对的结果表（`paths.align_review`，由
|
||
`scripts/ceus_align_review.py score` 写出）：核对没通过的病例第 1 步失败；表里没有这段视频，或配准
|
||
结果与表里记的不同，说明核对的不是这次跑的东西，整次运行停下。
|
||
第 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, Mapping, Sequence
|
||
from dataclasses import dataclass
|
||
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, field_validator
|
||
|
||
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（分屏）、registration（两半配准）、
|
||
# prompt_source（绿框颜色，配准和选提示帧都用）、pan_source（紫圈颜色，配准时修补掉）、keyframes
|
||
# （提示帧选择）。它们决定交给分割的帧和提示掩码，进分割缓存键。
|
||
_STEP_1_2_CEUS_FIELDS = frozenset(
|
||
{"views", "registration", "prompt_source", "pan_source", "keyframes"}
|
||
)
|
||
# segment 段里不改变分割结果的开关，不进缓存键；其余字段都进。
|
||
_RESULT_NEUTRAL_SEGMENT_FIELDS = frozenset(
|
||
{"device", "offload_video_to_cpu", "offload_state_to_cpu"}
|
||
)
|
||
_FAILED_HEADER = ("video_id", "step", "reason")
|
||
# 配准目视核对结果表 review.csv 的列：(dx, dy) 是左半 → 右半的平移
|
||
# （`HalfRegistration.correlation`），passed 取 true / false，后两列是两位判读者的判定。
|
||
REVIEW_HEADER = ("video_sha256", "dx", "dy", "passed", "verdict_a", "verdict_b")
|
||
|
||
|
||
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` 段：交付包、运行目录的上级、中间文件目录、配准目视核对的运行目录上级和结果表。
|
||
|
||
相对路径相对仓库根目录。`align_review` 在 ceus_quant 运行时必须存在，出图时用不到。
|
||
"""
|
||
|
||
model_config = ConfigDict(strict=True, extra="forbid", frozen=True)
|
||
|
||
delivery_dir: _PathField
|
||
runs_root: _PathField
|
||
work_root: _PathField
|
||
align_runs_root: _PathField
|
||
align_review: _PathField
|
||
|
||
|
||
class AlignCheckConfig(BaseModel):
|
||
"""`align_check` 段：配准目视核对出图用的随机数种子、阴性对照例数、对照错开距离区间（像素）。"""
|
||
|
||
model_config = ConfigDict(strict=True, extra="forbid", frozen=True)
|
||
|
||
seed: int
|
||
control_count: int = Field(ge=1)
|
||
control_shift_px: list[int]
|
||
|
||
@field_validator("control_shift_px")
|
||
@classmethod
|
||
def _shift_range(cls, value: list[int]) -> list[int]:
|
||
if len(value) != 2 or not 0 < value[0] <= value[1]:
|
||
raise ValueError(
|
||
f"align_check.control_shift_px 必须是 [lo, hi] 且 0 < lo ≤ hi，收到 {value}"
|
||
)
|
||
return value
|
||
|
||
|
||
class QuantConfig(BaseModel):
|
||
"""`ceus.yaml` 整份：`segment` 段交 segment 校验，`ceus` 段交 ceus 校验。"""
|
||
|
||
model_config = ConfigDict(strict=True, extra="forbid", frozen=True)
|
||
|
||
paths: PathsConfig
|
||
align_check: AlignCheckConfig
|
||
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 resolve_paths(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,
|
||
"align_runs_root": REPO_ROOT / paths.align_runs_root,
|
||
"align_review": REPO_ROOT / paths.align_review,
|
||
}
|
||
),
|
||
"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()
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ReviewEntry:
|
||
"""review.csv 的一行：核对时的配准结果（左半 → 右半）和是否通过。"""
|
||
|
||
offset: tuple[int, int]
|
||
passed: bool
|
||
|
||
|
||
def read_review(path: Path) -> dict[str, ReviewEntry]:
|
||
"""读配准目视核对结果表，视频 SHA-256 → 这一行。表头、取值不合规或同一视频出现两次就报错。"""
|
||
with path.open(encoding="utf-8", newline="") as handle:
|
||
rows = list(csv.reader(handle))
|
||
if not rows or tuple(rows[0]) != REVIEW_HEADER:
|
||
raise ValueError(f"{path} 的表头应为 {REVIEW_HEADER}")
|
||
table: dict[str, ReviewEntry] = {}
|
||
for number, row in enumerate(rows[1:], start=2):
|
||
if len(row) != len(REVIEW_HEADER) or row[3] not in ("true", "false"):
|
||
raise ValueError(f"{path} 第 {number} 行不合规：{row}")
|
||
if row[0] in table:
|
||
raise ValueError(f"{path} 第 {number} 行：同一段视频出现两次")
|
||
table[row[0]] = ReviewEntry(offset=(int(row[1]), int(row[2])), passed=row[3] == "true")
|
||
return table
|
||
|
||
|
||
def check_review(
|
||
prepared: SegmentationInput, video_sha256: str, review: Mapping[str, ReviewEntry]
|
||
) -> None:
|
||
"""第 1 步配准结果对照目视核对表：没通过抛 `CeusCaseError`（第 1 步）；对不上表就报错停下。"""
|
||
if video_sha256 not in review:
|
||
raise RuntimeError(
|
||
f"配准目视核对表里没有这段视频（SHA-256 {video_sha256}），核对的不是这批视频"
|
||
)
|
||
entry = review[video_sha256]
|
||
offset = prepared.registration.correlation
|
||
if offset != entry.offset:
|
||
raise RuntimeError(
|
||
f"这段视频（SHA-256 {video_sha256}）的配准结果 {offset} 与目视核对表里的 "
|
||
f"{entry.offset} 不同，核对的不是这次跑的东西"
|
||
)
|
||
if not entry.passed:
|
||
raise CeusCaseError(1, "配准目视核对未通过")
|
||
|
||
|
||
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],
|
||
review: Mapping[str, ReviewEntry],
|
||
run_dir: Path,
|
||
) -> tuple[tuple[str, ...], bool]:
|
||
"""跑一例并写出产物，返回（定位到的关键帧槽位，分割是否命中缓存）。
|
||
|
||
`key_of` 把视频的 SHA-256 换成分割缓存键（其余几样跨病例不变）；`review` 是配准目视核对表。
|
||
第 1–2 步都在 `prepare_segmentation` 里，所以第 2 步失败的病例不查核对表。
|
||
"""
|
||
prepared = prepare_segmentation(case.video, cfg.ceus)
|
||
video_sha256 = file_sha256(case.video)
|
||
check_review(prepared, video_sha256, review)
|
||
key = key_of(video_sha256)
|
||
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))
|
||
registration = prepared.registration
|
||
summary = {
|
||
"video_id": case.video_id,
|
||
"segment_cache_key": key,
|
||
"registration": {
|
||
"contrast_side": prepared.contrast.side,
|
||
"contrast_to_target_offset": list(prepared.offset),
|
||
"left_to_right_correlation": list(registration.correlation),
|
||
"left_to_right_template": list(registration.template),
|
||
"disagreement_px": registration.disagreement_px,
|
||
},
|
||
**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 = resolve_paths(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)
|
||
|
||
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)
|
||
|
||
# 缓存键用配置文件里的写法（路径不随仓库位置变）。
|
||
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, review, 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, "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"
|
||
)
|
||
# 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:])
|