Files
CADx-4/tests/test_ceus_pipeline.py
T
iomgaa 150eb5f6ae feat(ceus,segment): 任务书 02 第二阶段——分割和后半段算法
- ceus 迁入 polygons、dark_region、reference_region、tic_fitting、tic、anchor_frames、
  perfusion_metrics、trajectory、tic_curves、tic_plot、overlay,旧测试照搬(改 import 与配置读取)。
- 对外两段接口:prepare_segmentation(第 1–2 步)与 quantify_case(第 4–9 步)。
- 第 3 条:删掉 D1 取医生标注时刻的模式,D1 只从拟合曲线算。
- 第 4 条:F1/F2 取视频第 20−X、60−X 秒(X 为视频第 0 秒时已推注的秒数),落在视频外记未定位。
- 绿框融合权重、紫圈颜色阈值照搬;关键帧容差窗口的出处写进 ceus.yaml 注释。
- segment:同一进程调用 MedSAM2,设置写进 ceus.yaml 的 segment 段,全部必填;gpu 测试真跑通过。
- 依赖加 scipy、matplotlib、medsam2(可编辑安装,不编译 CUDA 扩展);uv 只解析 linux x86_64。
- 导入契约:只有 segment 引用 torch、sam2;segment 不引用其他 cadx 模块。

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-16 05:07:02 -04:00

169 lines
7.0 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""ceus 对外两段接口的端到端测试:合成一段分屏造影视频,第二段喂合成掩膜(不调分割模型)。
视频 400×200、10 fps、40 秒(400 帧),分屏线在 x=200(这一列全黑)。造影画面在右半:病灶是整帧
[250, 369] × [60, 159] 的方块,医生的绿框是它的空心描边;病灶内部亮度按 Lognormal 团注曲线从
基线 40 升到约 170 再回落,内部另有一块一直发黑的囊(暗区);每帧加一点随机噪声,免得相邻帧
被当成重复帧。左半是不变的灰阶画面。
手算的预期:
- 绿框在右半 → 造影画面是右半,分割那一半是左半 [0, 200);
- 绿框挪到左半局部坐标:x = 250 − 200 = 50 到 369 − 200 = 169;
- 推注时间 X = 5 秒 → F1 在视频第 20 − 5 = 15 秒(第 150 帧),F2 在第 60 − 5 = 55 秒,
超过 40 秒的视频末尾,未定位。
"""
import json
import math
from pathlib import Path
import cv2
import numpy as np
import pytest
from cadx.ceus import (
ANCHOR_SLOTS,
CeusConfig,
LognormalParameters,
NoValidContourError,
SegmentationInput,
evaluate_lognormal,
keyframe_images,
pack_tic_curves,
pack_trajectory,
prepare_segmentation,
quant_summary,
quantify_case,
render_case_overlay,
render_tic_png,
tic_plot_markers,
unpack_tic_curves,
unpack_trajectory,
)
_WIDTH, _HEIGHT, _SPLIT = 400, 200, 200
_FPS, _FRAMES = 10.0, 400
_LESION = (250, 60, 369, 159) # 整帧坐标 x1, y1, x2, y2(含端点)
_CYST = (290, 95, 320, 125)
_BOLUS = LognormalParameters(area=1300.0, mu=2.4, sigma=0.45, arrival_offset_s=4.0, baseline=40.0)
_INJECTION_OFFSET_S = 5.0
def _write_video(path: Path) -> None:
rng = np.random.default_rng(3)
times = np.arange(_FRAMES) / _FPS
intensity = evaluate_lognormal(times, _BOLUS)
writer = cv2.VideoWriter(str(path), cv2.VideoWriter.fourcc(*"mp4v"), _FPS, (_WIDTH, _HEIGHT))
x1, y1, x2, y2 = _LESION
for index in range(_FRAMES):
frame = np.zeros((_HEIGHT, _WIDTH, 3), dtype=np.uint8)
frame[:, :_SPLIT] = 120
frame[:, _SPLIT + 1 :] = 30
lesion = np.clip(
intensity[index] + rng.normal(0.0, 6.0, (y2 - y1 + 1, x2 - x1 + 1)), 0, 255
)
frame[y1 : y2 + 1, x1 : x2 + 1] = lesion.astype(np.uint8)[:, :, None]
cx1, cy1, cx2, cy2 = _CYST
frame[cy1 : cy2 + 1, cx1 : cx2 + 1] = 10
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), thickness=3)
writer.write(frame)
writer.release()
@pytest.fixture(scope="module")
def config() -> CeusConfig:
from conftest import load_ceus_section
return CeusConfig.model_validate(load_ceus_section())
@pytest.fixture(scope="module")
def prepared(tmp_path_factory: pytest.TempPathFactory, config: CeusConfig) -> SegmentationInput:
path = tmp_path_factory.mktemp("ceus") / "synthetic.mp4"
_write_video(path)
return prepare_segmentation(path, config)
def _lesion_masks(prepared: SegmentationInput) -> np.ndarray:
"""合成「分割结果」:每帧都是病灶方块挪到左半局部坐标后的位置(略往里收 3 像素)。"""
view = prepared.contrast.target_view
masks = np.zeros((prepared.frame_count, view.height, view.width), dtype=bool)
x1, y1, x2, y2 = _LESION
masks[:, y1 + 3 : y2 - 2, x1 - _SPLIT + 3 : x2 - _SPLIT - 2] = True
return masks
def test_prepare_segmentation_finds_contrast_half_and_maps_prompts(
prepared: SegmentationInput,
) -> None:
assert prepared.frame_count == _FRAMES
assert prepared.frame_size == (_WIDTH, _HEIGHT)
assert prepared.contrast.side == "right"
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)
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]
# 描边 3 像素宽、找绿框时再膨胀 2 像素、mp4 压缩会糊开边缘,容差放到 5 像素。
assert abs(int(xs.min()) - 50) <= 5
assert abs(int(xs.max()) - 169) <= 5
def test_quantify_case_end_to_end(prepared: SegmentationInput, config: CeusConfig) -> None:
result = quantify_case(prepared, _lesion_masks(prepared), _INJECTION_OFFSET_S, config)
assert result.trajectory.source_polygons.shape == (_FRAMES, 64, 2)
assert result.lesion_fit_ready is True
assert list(result.anchors) == list(ANCHOR_SLOTS)
for slot in ("D1", "D2", "D3"):
assert result.anchors[slot].status == "ok", result.anchors[slot].reason
assert result.anchors["F1"].frame_index == 150
assert result.anchors["F1"].frame_time_video_s == pytest.approx(15.0, abs=0.06)
assert result.anchors["F2"].status == "not_covered"
assert result.anchors["F2"].reason is not None
assert result.perfusion["roi_perf"] is not None
peak_time = _BOLUS.arrival_offset_s + math.exp(_BOLUS.mu - _BOLUS.sigma**2)
assert result.perfusion["roi_perf"].ttp_s == pytest.approx(peak_time, abs=0.5)
# 画面上没有紫圈。
assert result.pan_detected_fraction == 0.0
assert result.tic.reference_valid is False
# 暗区(囊)被找出来,去暗区曲线比病灶整体亮。
peak_row = result.tic.rows[round(peak_time * _FPS)]
assert peak_row["dark_area_px"] > 0
assert peak_row["roi_perf_raw"] > peak_row["roi_all_raw"]
restored = unpack_trajectory(pack_trajectory(result.trajectory))
assert restored.calibration["contrast_side"] == "right"
curves = unpack_tic_curves(pack_tic_curves(result.curves))
assert curves.injection_offset_s == pytest.approx(_INJECTION_OFFSET_S)
summary = quant_summary(result)
json.dumps(summary, allow_nan=False)
assert summary["anchor_frames"]["F1"]["frame_index"] == 150
assert summary["perfusion_metrics"]["roi_perf"] is not None
markers = tic_plot_markers(result.anchors)
assert [marker.label.split()[0] for marker in markers][:3] == ["D1", "D2", "D3"]
assert render_tic_png(result.tic, markers=markers).startswith(b"\x89PNG")
assert len(render_case_overlay(prepared, result)) > 0
images = keyframe_images(prepared, result.anchors)
assert "F1" in images and "F2" not in images
assert all(image.shape == (_HEIGHT, _WIDTH, 3) for image in images.values())
def test_empty_segmentation_fails_case_at_step_4(
prepared: SegmentationInput, config: CeusConfig
) -> None:
empty = np.zeros_like(_lesion_masks(prepared))
with pytest.raises(NoValidContourError) as raised:
quantify_case(prepared, empty, _INJECTION_OFFSET_S, config)
assert raised.value.step == 4
def test_mask_shape_mismatch_is_rejected(prepared: SegmentationInput, config: CeusConfig) -> None:
masks = _lesion_masks(prepared)[:-1]
with pytest.raises(ValueError, match="分割掩膜"):
quantify_case(prepared, masks, _INJECTION_OFFSET_S, config)