CADx-4/tests/test_ceus_registration.py

"""两半配准：每段视频单独估左半 → 右半的平移（任务书 02 第一阶段第 1 条）。

合成分屏画面（`fake_split_scene`）：每半 240×200，左半扇形是一个梯形，右半的扇形和纹理是左半平移
(dx, dy) 后的样子，再换一种亮度、加独立噪声。预期值就是造图时设定的平移，不从实现里推。
"""

import numpy as np
import pytest

from cadx.ceus import (
    CeusCaseError,
    CeusConfig,
    RegistrationError,
    ViewRect,
    register_halves,
)
from fake_split_scene import SplitScene, make_scene, scene_gray, to_bgr

_HEIGHT, _HALF = 200, 240
_LEFT, _RIGHT = ViewRect(0, 0, _HALF, _HEIGHT), ViewRect(_HALF, 0, 2 * _HALF, _HEIGHT)
# 左半局部坐标的梯形扇形：上窄下宽，顶边在顶部信息栏（200 × 0.0834 ≈ 17 行）之下。
_FAN = [(80, 30), (170, 30), (215, 185), (35, 185)]


def _frames(scene: SplitScene, *, green_box: bool) -> list[np.ndarray]:
    rng = np.random.default_rng(11)
    frames: list[np.ndarray] = []
    for index in range(12):
        frame = to_bgr(scene_gray(scene, index, rng))
        if green_box:
            # 绿框只画在左半：配准前要修补掉，否则它的边缘只在一半上有。
            frame[90:140, 100] = frame[90:140, 160] = (0, 255, 0)
            frame[90, 100:161] = frame[140, 100:161] = (0, 255, 0)
        frames.append(frame)
    return frames


@pytest.mark.parametrize(
    ("offset", "green_box"),
    [
        ((-37, 0), False),  # 负的水平偏移（真实视频里最常见：造影那一半扇形左边有黑边）
        ((25, 4), False),  # 正的水平偏移、非零竖直偏移
        ((-12, -6), True),  # 两个分量都为负，左半画着绿框
        ((0, 0), False),
    ],
)
def test_known_offset_is_recovered_by_both_methods(
    ceus_config: CeusConfig, offset: tuple[int, int], green_box: bool
) -> None:
    scene = make_scene(
        height=_HEIGHT,
        half_width=_HALF,
        fan_polygon=_FAN,
        offset=offset,
        interior_shift=None,
        texture_std=25.0,
        right_gain=0.8,
        right_bias=10.0,
        seed=5,
    )
    registration = register_halves(
        _frames(scene, green_box=green_box),
        _LEFT,
        _RIGHT,
        ceus_config.registration,
        ceus_config.prompt_source,
        ceus_config.pan_source,
    )
    assert registration.correlation == offset
    assert registration.template == offset
    assert registration.disagreement_px == 0
    # 造影在左半时平移照用；造影在右半时（把右半的东西挪到左半）取反。
    assert registration.contrast_to_target("left") == offset
    assert registration.contrast_to_target("right") == (-offset[0], -offset[1])


# 冲突画面用更宽的扇形：模板块落在扇形中部，纹理平移后真正的匹配位置仍在右半扇形内。
_WIDE_FAN = [(40, 25), (200, 25), (235, 195), (5, 195)]


def _conflicting_scene() -> SplitScene:
    """扇形轮廓按 (15, 0) 平移，里面的纹理按 (−15, 0) 平移：两种方法各信一样。

    纹理起伏（标准差 5）远小于扇形内外的亮度差（80），边缘图上占主导的是扇形轮廓，互相关跟着
    轮廓走；模板块整块落在扇形里，只看得到纹理，跟着纹理走。
    """
    return make_scene(
        height=_HEIGHT,
        half_width=_HALF,
        fan_polygon=_WIDE_FAN,
        offset=(15, 0),
        interior_shift=(-15, 0),
        texture_std=5.0,
        right_gain=0.8,
        right_bias=10.0,
        seed=5,
    )


def test_methods_disagree_on_conflicting_scene(ceus_config: CeusConfig) -> None:
    """门槛放到极大时两种方法各自的估计：互相关跟轮廓 (15, 0)，模板匹配跟纹理 (−15, 0)。"""
    loose = ceus_config.registration.model_copy(update={"max_disagreement_px": 10_000})
    registration = register_halves(
        _frames(_conflicting_scene(), green_box=False),
        _LEFT,
        _RIGHT,
        loose,
        ceus_config.prompt_source,
        ceus_config.pan_source,
    )
    assert registration.correlation == (15, 0)
    assert registration.template == (-15, 0)
    assert registration.disagreement_px == 30


def test_disagreement_beyond_threshold_fails_case_at_step_1(ceus_config: CeusConfig) -> None:
    with pytest.raises(RegistrationError, match="两半配准不一致") as raised:
        register_halves(
            _frames(_conflicting_scene(), green_box=False),
            _LEFT,
            _RIGHT,
            ceus_config.registration,
            ceus_config.prompt_source,
            ceus_config.pan_source,
        )
    assert isinstance(raised.value, CeusCaseError)
    assert raised.value.step == 1
    # 原因里附两个估计值。
    assert "(15, 0)" in raised.value.reason and "(-15, 0)" in raised.value.reason


def test_static_half_has_no_fan_and_fails(ceus_config: CeusConfig) -> None:
    """一半画面完全不动（没有扇形区）：第 1 步失败，不猜一个平移。"""
    frames = [np.full((_HEIGHT, 2 * _HALF, 3), 60, dtype=np.uint8) for _ in range(12)]
    with pytest.raises(RegistrationError, match="找不到扇形区") as raised:
        register_halves(
            frames,
            _LEFT,
            _RIGHT,
            ceus_config.registration,
            ceus_config.prompt_source,
            ceus_config.pan_source,
        )
    assert raised.value.step == 1