iomgaa
150eb5f6ae
- 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>
261 lines
11 KiB
Python
261 lines
11 KiB
Python
"""trajectory npz 打包/解包单元测试（照搬旧项目 test_ceus_trajectory.py）。
|
||
|
||
旧标题：B11 trajectory npz 打包/解包单元测试。
|
||
|
||
钉死轨迹工件的单一 schema：完整往返还原全部键（沿用参考 ``write_roi_trajectory_npz``
|
||
键集 + 新增 ``keyframe_indices`` / ``dark_polygons`` / ``calibration_json``）、
|
||
各键 dtype 精确（int32 / float32 / uint8 / str）、变长 ``dark_polygons``（逐帧
|
||
不同数量多边形）与 ``keyframe_indices`` 往返、``reference_region`` / ``calibration``
|
||
dict 往返，以及防御：坏字节解包 raise、形状不一致 pack raise、禁 pickle
|
||
（含 object pickle 的恶意 npz 解包被拒）。测试用真实数据的二次构造（合成多边形 /
|
||
掩膜），边缘场景靠实现里的 assertion/raise 保证。
|
||
"""
|
||
|
||
import io
|
||
import pickle
|
||
|
||
import numpy as np
|
||
import pytest
|
||
|
||
from cadx.ceus import (
|
||
Trajectory,
|
||
ViewRect,
|
||
pack_trajectory,
|
||
unpack_trajectory,
|
||
)
|
||
|
||
# Phase 1: 构造一个真实形态的 Trajectory（二次构造的合成几何）
|
||
|
||
|
||
def _make_polygon(offset: float) -> np.ndarray:
|
||
"""构造 (64, 2) float32 闭合多边形（单位圆平移），与重采样输出形态一致。"""
|
||
angles = np.linspace(0.0, 2.0 * np.pi, 64, endpoint=False)
|
||
points = np.stack([np.cos(angles), np.sin(angles)], axis=1) * 10.0 + offset
|
||
return points.astype(np.float32)
|
||
|
||
|
||
def _make_trajectory() -> Trajectory:
|
||
num_frames = 3
|
||
target_polygons = [_make_polygon(float(i)) for i in range(num_frames)]
|
||
source_polygons = [_make_polygon(float(i) + 100.0) for i in range(num_frames)]
|
||
# 逐帧暗区多边形数量不同：0 个 / 1 个 / 2 个（变长嵌套）
|
||
dark_polygons: list[list[np.ndarray]] = [
|
||
[],
|
||
[_make_polygon(5.0)],
|
||
[_make_polygon(6.0), _make_polygon(7.0)],
|
||
]
|
||
reference_mask = np.zeros((20, 30), dtype=np.uint8)
|
||
reference_mask[5:10, 5:15] = 1
|
||
return Trajectory(
|
||
frame_indices=np.arange(num_frames, dtype=np.int32),
|
||
fps=29.97,
|
||
target_polygons=np.asarray(target_polygons, dtype=np.float32),
|
||
source_polygons=np.asarray(source_polygons, dtype=np.float32),
|
||
source_view_rect=ViewRect(0, 0, 100, 200),
|
||
target_view_rect=ViewRect(100, 0, 200, 200),
|
||
source_effective_rect=ViewRect(5, 5, 95, 195),
|
||
target_effective_rect=ViewRect(105, 5, 195, 195),
|
||
source_polygon_coordinate_space="source_view_local",
|
||
target_polygon_coordinate_space="target_effective_view_local",
|
||
reference_mask=reference_mask,
|
||
reference_region={"frame_index": 1, "area": 50.0, "source": "auto"},
|
||
keyframe_indices=np.asarray([1, 2], dtype=np.int32),
|
||
dark_polygons=dark_polygons,
|
||
calibration={
|
||
"source": "config/ceus.yaml",
|
||
"config_hash": "abc123",
|
||
"cli_override": False,
|
||
"detected_split_x": 100,
|
||
},
|
||
)
|
||
|
||
|
||
# Phase 2: 完整往返 —— 全部键还原
|
||
|
||
|
||
def test_roundtrip_preserves_all_keys() -> None:
|
||
trajectory = _make_trajectory()
|
||
restored = unpack_trajectory(pack_trajectory(trajectory))
|
||
|
||
np.testing.assert_array_equal(restored.frame_indices, trajectory.frame_indices)
|
||
assert restored.fps == pytest.approx(trajectory.fps)
|
||
np.testing.assert_allclose(restored.target_polygons, trajectory.target_polygons)
|
||
np.testing.assert_allclose(restored.source_polygons, trajectory.source_polygons)
|
||
assert restored.source_view_rect == trajectory.source_view_rect
|
||
assert restored.target_view_rect == trajectory.target_view_rect
|
||
assert restored.source_effective_rect == trajectory.source_effective_rect
|
||
assert restored.target_effective_rect == trajectory.target_effective_rect
|
||
assert restored.source_polygon_coordinate_space == trajectory.source_polygon_coordinate_space
|
||
assert restored.target_polygon_coordinate_space == trajectory.target_polygon_coordinate_space
|
||
np.testing.assert_array_equal(restored.reference_mask, trajectory.reference_mask)
|
||
assert restored.reference_region == trajectory.reference_region
|
||
np.testing.assert_array_equal(restored.keyframe_indices, trajectory.keyframe_indices)
|
||
assert restored.calibration == trajectory.calibration
|
||
|
||
|
||
# Phase 3: dtype 精确
|
||
|
||
|
||
def test_roundtrip_dtypes_are_exact() -> None:
|
||
restored = unpack_trajectory(pack_trajectory(_make_trajectory()))
|
||
assert restored.frame_indices.dtype == np.int32
|
||
assert restored.target_polygons.dtype == np.float32
|
||
assert restored.source_polygons.dtype == np.float32
|
||
assert restored.reference_mask.dtype == np.uint8
|
||
assert restored.keyframe_indices.dtype == np.int32
|
||
assert isinstance(restored.fps, float)
|
||
assert isinstance(restored.source_polygon_coordinate_space, str)
|
||
assert isinstance(restored.target_polygon_coordinate_space, str)
|
||
|
||
|
||
def test_polygon_shapes_are_n_64_2() -> None:
|
||
restored = unpack_trajectory(pack_trajectory(_make_trajectory()))
|
||
assert restored.target_polygons.shape == (3, 64, 2)
|
||
assert restored.source_polygons.shape == (3, 64, 2)
|
||
|
||
|
||
# Phase 3b: 存档层 schema —— 直接在 npz 原始存档上断言键名 / dtype / shape
|
||
# （对象层往返不足以守住 pack 的存储 schema；pack 若 schema 漂移须在此被发现）
|
||
|
||
|
||
def test_archive_schema_keys_dtypes_shapes() -> None:
|
||
"""打开 pack 产出的原始 npz 存档，逐一断言键集 / 存储 dtype / shape。
|
||
|
||
与对象层往返互补：往返只验 unpack(pack(x))==x，无法发现 pack 写入存档时的
|
||
dtype / shape 漂移（如 fps 被写成 float64、polygons 退化）。此处用
|
||
``allow_pickle=False`` 直接读取存档，钉死存储 schema。
|
||
"""
|
||
archive_bytes = pack_trajectory(_make_trajectory())
|
||
with np.load(io.BytesIO(archive_bytes), allow_pickle=False) as archive:
|
||
# 键集合 == 全部必要键（无多无少）
|
||
expected_keys = {
|
||
"frame_indices",
|
||
"fps",
|
||
"target_polygons",
|
||
"source_polygons",
|
||
"source_view_rect",
|
||
"target_view_rect",
|
||
"source_effective_rect",
|
||
"target_effective_rect",
|
||
"source_polygon_coordinate_space",
|
||
"target_polygon_coordinate_space",
|
||
"reference_mask",
|
||
"reference_region_json",
|
||
"keyframe_indices",
|
||
"dark_polygons",
|
||
"calibration_json",
|
||
}
|
||
assert set(archive.files) == expected_keys
|
||
|
||
# 存档内 dtype 精确（存储 dtype 是 schema 的一部分）
|
||
assert archive["frame_indices"].dtype == np.int32
|
||
assert archive["keyframe_indices"].dtype == np.int32
|
||
assert archive["target_polygons"].dtype == np.float32
|
||
assert archive["source_polygons"].dtype == np.float32
|
||
# fps 存档 dtype 必须 float32（即使 in-memory unpack 还原为 Python float）
|
||
assert archive["fps"].dtype == np.float32
|
||
for rect_key in (
|
||
"source_view_rect",
|
||
"target_view_rect",
|
||
"source_effective_rect",
|
||
"target_effective_rect",
|
||
):
|
||
assert archive[rect_key].dtype == np.int32
|
||
assert archive["reference_mask"].dtype == np.uint8
|
||
|
||
# shape 精确
|
||
assert archive["target_polygons"].shape == (3, 64, 2)
|
||
assert archive["source_polygons"].shape == (3, 64, 2)
|
||
assert archive["frame_indices"].shape == (3,)
|
||
for rect_key in (
|
||
"source_view_rect",
|
||
"target_view_rect",
|
||
"source_effective_rect",
|
||
"target_effective_rect",
|
||
):
|
||
assert archive[rect_key].shape == (4,)
|
||
assert archive["reference_mask"].ndim == 2
|
||
|
||
|
||
# Phase 4: 变长 dark_polygons 往返
|
||
|
||
|
||
def test_roundtrip_variable_length_dark_polygons() -> None:
|
||
trajectory = _make_trajectory()
|
||
restored = unpack_trajectory(pack_trajectory(trajectory))
|
||
assert len(restored.dark_polygons) == 3
|
||
assert len(restored.dark_polygons[0]) == 0
|
||
assert len(restored.dark_polygons[1]) == 1
|
||
assert len(restored.dark_polygons[2]) == 2
|
||
for original_frame, restored_frame in zip(
|
||
trajectory.dark_polygons, restored.dark_polygons, strict=True
|
||
):
|
||
for original_poly, restored_poly in zip(original_frame, restored_frame, strict=True):
|
||
assert isinstance(restored_poly, np.ndarray)
|
||
assert restored_poly.dtype == np.float32
|
||
np.testing.assert_allclose(restored_poly, original_poly)
|
||
|
||
|
||
def test_roundtrip_empty_dark_polygons() -> None:
|
||
trajectory = _make_trajectory()
|
||
trajectory.dark_polygons = [[], [], []]
|
||
restored = unpack_trajectory(pack_trajectory(trajectory))
|
||
assert restored.dark_polygons == [[], [], []]
|
||
|
||
|
||
# Phase 5: 防御 —— 坏字节 / 形状不一致 / 禁 pickle
|
||
|
||
|
||
def test_unpack_bad_bytes_raises() -> None:
|
||
with pytest.raises((ValueError, OSError)):
|
||
unpack_trajectory(b"this-is-not-a-valid-npz-archive")
|
||
|
||
|
||
def test_pack_rejects_mismatched_frame_count() -> None:
|
||
trajectory = _make_trajectory()
|
||
# frame_indices 长度与多边形帧数不一致
|
||
trajectory.frame_indices = np.arange(5, dtype=np.int32)
|
||
with pytest.raises(ValueError):
|
||
pack_trajectory(trajectory)
|
||
|
||
|
||
def test_pack_rejects_wrong_polygon_shape() -> None:
|
||
trajectory = _make_trajectory()
|
||
# 把 target_polygons 改成 (N, 32, 2)，违反 (N, 64, 2)
|
||
trajectory.target_polygons = np.zeros((3, 32, 2), dtype=np.float32)
|
||
with pytest.raises(ValueError):
|
||
pack_trajectory(trajectory)
|
||
|
||
|
||
def test_pack_rejects_non_2d_reference_mask() -> None:
|
||
trajectory = _make_trajectory()
|
||
trajectory.reference_mask = np.zeros((20, 30, 3), dtype=np.uint8)
|
||
with pytest.raises(ValueError):
|
||
pack_trajectory(trajectory)
|
||
|
||
|
||
def test_unpack_rejects_pickled_object_payload() -> None:
|
||
"""构造真实含 object pickle 的恶意 .npz（须 allow_pickle 才能 save/load），解包须被拒。
|
||
|
||
关键：用 ``np.savez_compressed`` 生成真正的 .npz 存档（而非 .npy），其中含一个
|
||
object dtype 数组（背后是 pickle payload）。这覆盖了「合法 .npz 容器、但成员需
|
||
pickle 反序列化」的注入面——``unpack_trajectory`` 内部 ``allow_pickle=False`` 必须
|
||
在访问该成员时 raise，绝不静默反序列化执行。
|
||
"""
|
||
buffer = io.BytesIO()
|
||
# 把 object dtype 数组放进 unpack 首先访问的必要键 frame_indices，确保命中
|
||
# allow_pickle=False 对 object 成员的拒绝路径（而非仅缺键 KeyError）。
|
||
# object dtype 数组在 savez 时即以 pickle 存储，构造真实 .npz 字节。
|
||
np.savez_compressed(buffer, frame_indices=np.array([{"evil": 1}], dtype=object))
|
||
raw = buffer.getvalue()
|
||
# 访问 frame_indices 这一 object 成员时，allow_pickle=False 必须 raise（ValueError）
|
||
with pytest.raises((ValueError, OSError, KeyError)):
|
||
unpack_trajectory(raw)
|
||
|
||
|
||
def test_pickle_payload_cannot_smuggle_through_unpack() -> None:
|
||
"""即便 pickle 字节伪装，unpack 内部 allow_pickle=False 也不得执行反序列化。"""
|
||
payload = pickle.dumps({"evil": "code"})
|
||
with pytest.raises((ValueError, OSError, KeyError)):
|
||
unpack_trajectory(payload)
|