iomgaa
1de88a5d0a
- 内存：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>
129 lines
4.7 KiB
Python
129 lines
4.7 KiB
Python
"""segment：分割设置的校验，以及在合成小视频上真跑一次 MedSAM2（`gpu`）。"""
|
||
|
||
from pathlib import Path
|
||
from typing import Any, cast
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import numpy.typing as npt
|
||
import pytest
|
||
import yaml
|
||
from pydantic import ValidationError
|
||
|
||
from cadx.segment import SegmentConfig, medsam2_commit, segment_video
|
||
|
||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||
CEUS_YAML = REPO_ROOT / "config" / "ceus.yaml"
|
||
|
||
|
||
def _segment_section() -> dict[str, Any]:
|
||
document = cast(dict[str, Any], yaml.safe_load(CEUS_YAML.read_text(encoding="utf-8")))
|
||
return cast(dict[str, Any], document["segment"])
|
||
|
||
|
||
def _absolute(section: dict[str, Any]) -> dict[str, Any]:
|
||
"""路径字段换成仓库根目录下的绝对路径（脚本也这样做）。"""
|
||
return {
|
||
**section,
|
||
"medsam2_dir": str(REPO_ROOT / section["medsam2_dir"]),
|
||
"checkpoint": str(REPO_ROOT / section["checkpoint"]),
|
||
}
|
||
|
||
|
||
def test_config_yaml_section_validates() -> None:
|
||
cfg = SegmentConfig.model_validate(_segment_section())
|
||
assert cfg.sam2_config == "configs/sam2.1_hiera_t512.yaml"
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"field",
|
||
[
|
||
"medsam2_dir",
|
||
"checkpoint",
|
||
"sam2_config",
|
||
"device",
|
||
"offload_video_to_cpu",
|
||
"offload_state_to_cpu",
|
||
],
|
||
)
|
||
def test_config_missing_field_raises(field: str) -> None:
|
||
section = _segment_section()
|
||
del section[field]
|
||
with pytest.raises(ValidationError, match=field):
|
||
SegmentConfig.model_validate(section)
|
||
|
||
|
||
def test_config_extra_field_raises() -> None:
|
||
with pytest.raises(ValidationError, match="object_id"):
|
||
SegmentConfig.model_validate({**_segment_section(), "object_id": 1})
|
||
|
||
|
||
def test_config_bool_not_coerced() -> None:
|
||
with pytest.raises(ValidationError, match="offload_video_to_cpu"):
|
||
SegmentConfig.model_validate({**_segment_section(), "offload_video_to_cpu": "false"})
|
||
|
||
|
||
def _moving_disc_video(
|
||
frame_count: int, height: int, width: int, radius: int
|
||
) -> tuple[list[npt.NDArray[np.uint8]], list[npt.NDArray[np.bool_]]]:
|
||
"""暗背景上一个亮圆从左往右匀速移动；交出逐帧 BGR 图像和每帧圆的真值掩码。"""
|
||
frames: list[npt.NDArray[np.uint8]] = []
|
||
truths: list[npt.NDArray[np.bool_]] = []
|
||
yy, xx = np.mgrid[0:height, 0:width]
|
||
for index in range(frame_count):
|
||
cx = radius + 10 + index * 2
|
||
cy = height // 2
|
||
frame = np.full((height, width, 3), 30, dtype=np.uint8)
|
||
cv2.circle(frame, (cx, cy), radius, (220, 220, 220), thickness=-1)
|
||
frames.append(frame)
|
||
truths.append((xx - cx) ** 2 + (yy - cy) ** 2 <= radius**2)
|
||
return frames, truths
|
||
|
||
|
||
@pytest.mark.gpu
|
||
def test_segment_video_on_synthetic_clip() -> None:
|
||
cfg = SegmentConfig.model_validate(_absolute(_segment_section()))
|
||
frame_count, height, width = 40, 120, 200
|
||
frames, truths = _moving_disc_video(frame_count, height, width, radius=18)
|
||
# 两个提示帧，最早的不在第 0 帧，前向和反向传播都要走到。
|
||
prompt_masks = {5: truths[5], 25: truths[25]}
|
||
|
||
masks = segment_video(frames, prompt_masks, cfg)
|
||
|
||
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)
|