hookvision-rk/tests/test_calibration.py

"""标定功能测试：偏移计算、1 秒平均采集、持久化与标定后 READY。"""

import time

import pytest

from hookvision.models import HookState, Config, BucketModel
from hookvision.config_loader import load_config
from hookvision import main_window_v2 as mwmod
from hookvision.main_window_v2 import MainWindow


def _obj(label, x, y, w=20, h=20):
    return {"x": x, "y": y, "w": w, "h": h, "label": label}


def _fresh(mw):
    mw.hook_state = HookState.IDLE
    mw.prev_state = HookState.IDLE
    mw.state_counter = 0
    mw._pair_traj = {"left": [], "right": []}
    mw._voice_suppress_until = 0.0
    mw.prev_left_hook = None
    mw.prev_left_lug = None
    mw.prev_right_hook = None
    mw.prev_right_lug = None


def _write_minimal_config(path):
    path.write_text(
        """[App]
name = 测试系统
copyright = 测试
screen = 0

[Video]
rtsp_timeout = 5

[Hardware]
module_type = RFID

[RFID]
ip = 192.168.1.1
port = 60000
stability_time = 5
stability_count = 3

[AI]
enabled = true

[Bucket:0]
id = 1
description = 测试工位
code_x = TEST001
left_rtsp =
right_rtsp =
left_topic =
right_topic =
""", encoding="utf-8"
    )


class TestComputeOffsetRatios:
    def test_normalized_ratios(self, qapp):
        mw = MainWindow(Config())
        # hook 中心 (40,40)，lug 中心 (50,50)，lug 20x20 → xr=yr=-0.5
        mw.left_objects = [_obj("hook", 30, 30), _obj("lug", 40, 40)]
        assert mw._compute_offset_ratios("left") == (-0.5, -0.5)

    def test_no_intersecting_pair_none(self, qapp):
        mw = MainWindow(Config())
        mw.left_objects = [_obj("hook", 0, 0), _obj("lug", 200, 200)]
        assert mw._compute_offset_ratios("left") is None

    def test_missing_objects_none(self, qapp):
        mw = MainWindow(Config())
        assert mw._compute_offset_ratios("left") is None
        mw.right_objects = [_obj("hook", 0, 0)]
        assert mw._compute_offset_ratios("right") is None

    def test_best_intersecting_pair_ignores_far_pair(self, qapp):
        mw = MainWindow(Config())
        # 远处的配对不相交，应被忽略；只有近处相交配对参与
        mw.left_objects = [_obj("hook", 0, 0), _obj("lug", 200, 200),
                           _obj("hook", 30, 30), _obj("lug", 40, 40)]
        assert mw._compute_offset_ratios("left") == (-0.5, -0.5)

    def test_is_bucket_calibrated(self, qapp):
        mw = MainWindow(Config())
        b = BucketModel(id=1, description="t", code_x="X")
        assert mw._is_bucket_calibrated(b) is False
        b.hook_lug_x_offset_ratio_left = 0.0
        b.hook_lug_y_offset_ratio_left = 0.0
        assert mw._is_bucket_calibrated(b) is True

    def test_y_axis_uses_bottom_edge_not_center(self, qapp):
        mw = MainWindow(Config())
        # hook(30,30,20,40)：底边 y+h=70，中心 y=50；lug(40,50,20,20)：底边 70，中心 60
        # 底边对齐 → yr=0.0；若按中心算会是 (50-60)/20=-0.5
        mw.left_objects = [_obj("hook", 30, 30, w=20, h=40), _obj("lug", 40, 50)]
        assert mw._compute_offset_ratios("left") == (-0.5, 0.0)

    def test_compute_ratios_skips_lug_at_top(self, qapp):
        # 吊耳贴着画面上沿：采样无效，不参与标定
        mw = MainWindow(Config())
        mw.left_objects = [_obj("hook", 40, 0), _obj("lug", 40, 0)]
        assert mw._compute_offset_ratios("left") is None


class TestCalibrationFlow:
    def _make_mw(self, tmp_path):
        cfg_path = tmp_path / "config.ini"
        _write_minimal_config(cfg_path)
        config = load_config(str(cfg_path))
        mw = MainWindow(config)
        _fresh(mw)
        mw._config_path = str(cfg_path)
        mw.active_bucket = config.buckets[0]
        return mw, cfg_path

    def test_sample_and_finish_writes_offsets(self, qapp, tmp_path, monkeypatch):
        monkeypatch.setattr(mwmod.QMessageBox, "information", lambda *a, **k: None)
        mw, _ = self._make_mw(tmp_path)
        bucket = mw.active_bucket
        mw.left_objects = [_obj("hook", 30, 30), _obj("lug", 40, 40)]
        mw._start_calibration()
        for _ in range(5):
            mw._sample_calibration()
        assert bucket.hook_lug_x_offset_ratio_left is None
        mw._calib_until = time.time() - 1
        mw._sample_calibration()
        assert bucket.hook_lug_x_offset_ratio_left is not None
        assert bucket.hook_lug_x_offset_ratio_left == pytest.approx(-0.5, abs=1e-6)
        assert bucket.hook_lug_y_offset_ratio_left == pytest.approx(-0.5, abs=1e-6)

    def test_sample_requires_intersecting_pair(self, qapp, tmp_path, monkeypatch):
        monkeypatch.setattr(mwmod.QMessageBox, "information", lambda *a, **k: None)
        mw, _ = self._make_mw(tmp_path)
        bucket = mw.active_bucket
        mw.left_objects = [_obj("hook", 0, 0), _obj("lug", 200, 200)]
        mw._start_calibration()
        for _ in range(5):
            mw._sample_calibration()
        mw._calib_until = time.time() - 1
        mw._sample_calibration()
        assert bucket.hook_lug_x_offset_ratio_left is None
        assert mw._is_bucket_calibrated(bucket) is False

    def test_calibrated_bucket_reaches_ready(self, qapp, tmp_path, monkeypatch):
        monkeypatch.setattr(mwmod.QMessageBox, "information", lambda *a, **k: None)
        mw, _ = self._make_mw(tmp_path)
        # 先标定（中心对齐，偏移 0,0）
        mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
        mw._start_calibration()
        for _ in range(5):
            mw._sample_calibration()
        mw._calib_until = time.time() - 1
        mw._sample_calibration()
        assert mw.active_bucket.hook_lug_x_offset_ratio_left == pytest.approx(0.0, abs=1e-6)
        # 重新喂帧 → 标定值匹配 → READY
        _fresh(mw)
        for _ in range(3):
            mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
            mw.evaluate_hook_state()
        assert mw.hook_state == HookState.READY

    def test_reset_calibration(self, qapp, tmp_path):
        mw, _ = self._make_mw(tmp_path)
        bucket = mw.active_bucket
        bucket.hook_lug_x_offset_ratio_left = -0.5
        bucket.hook_lug_y_offset_ratio_left = -0.5
        assert mw._is_bucket_calibrated(bucket) is True
        mw._reset_calibration()
        assert mw._is_bucket_calibrated(bucket) is False


class TestCalibrationPersistence:
    def test_save_config_round_trip(self, qapp, tmp_path):
        cfg_path = tmp_path / "config.ini"
        _write_minimal_config(cfg_path)
        config = load_config(str(cfg_path))
        mw = MainWindow(config)
        mw._config_path = str(cfg_path)
        bucket = config.buckets[0]
        bucket.hook_lug_x_offset_ratio_left = -0.25
        bucket.hook_lug_y_offset_ratio_left = 0.125
        bucket.hook_lug_x_offset_ratio_right = 0.0
        bucket.hook_lug_y_offset_ratio_right = 0.0
        mw._save_config()
        reloaded = load_config(str(cfg_path))
        b = reloaded.buckets[0]
        assert b.hook_lug_x_offset_ratio_left == pytest.approx(-0.25, abs=1e-6)
        assert b.hook_lug_y_offset_ratio_left == pytest.approx(0.125, abs=1e-6)
        assert b.hook_lug_x_offset_ratio_right == pytest.approx(0.0, abs=1e-6)
        assert b.hook_lug_y_offset_ratio_right == pytest.approx(0.0, abs=1e-6)
        assert reloaded.ai.offset_ratio_tolerance == 0.25
        assert reloaded.ai.ready_debounce == 3

    def test_save_config_removes_uncalibrated_keys(self, qapp, tmp_path):
        cfg_path = tmp_path / "config.ini"
        _write_minimal_config(cfg_path)
        config = load_config(str(cfg_path))
        mw = MainWindow(config)
        mw._config_path = str(cfg_path)
        mw._save_config()
        reloaded = load_config(str(cfg_path))
        b = reloaded.buckets[0]
        assert b.hook_lug_x_offset_ratio_left is None
        assert b.hook_lug_y_offset_ratio_left is None


class TestCalibrationUI:
    def test_calib_button_exists(self, qapp):
        mw = MainWindow(Config())
        assert mw.calib_btn is not None
        assert mw.calib_btn.text() == "标定"

    def test_badge_refresh(self, qapp):
        bucket = BucketModel(id=1, description="t", code_x="X")
        config = Config()
        config.buckets = [bucket]
        mw = MainWindow(config)
        assert bucket.id in mw._calib_badge_labels
        mw.refresh_bucket_icons()
        assert mw._calib_badge_labels[bucket.id].styleSheet() == "background-color: #dcdfe6; border-radius: 2px;"
        assert "未标定" in mw._calib_badge_labels[bucket.id].toolTip()
        bucket.hook_lug_x_offset_ratio_left = 0.0
        bucket.hook_lug_y_offset_ratio_left = 0.0
        mw.refresh_bucket_icons()
        assert mw._calib_badge_labels[bucket.id].styleSheet() == "background-color: #67c23a; border-radius: 2px;"
        assert "已标定" in mw._calib_badge_labels[bucket.id].toolTip()


def _bucket_with_cameras(cfg_bucket):
    from hookvision.models import CameraModel, CameraPosition
    cfg_bucket.cameras = [
        CameraModel(position=CameraPosition.LEFT, rtsp_url="", topic="test/left"),
        CameraModel(position=CameraPosition.RIGHT, rtsp_url="", topic="test/right"),
    ]
    return cfg_bucket


class TestAutoCalibration:
    def _make_mw(self, qapp, tmp_path):
        cfg_path = tmp_path / "config.ini"
        _write_minimal_config(cfg_path)
        config = load_config(str(cfg_path))
        bucket = _bucket_with_cameras(config.buckets[0])
        mw = MainWindow(config)
        _fresh(mw)
        mw._config_path = str(cfg_path)
        mw.active_bucket = bucket
        return mw, bucket

    def _ready_trigger(self, mw, side="left"):
        for _ in range(3):
            mw._update_auto_calibration(side, True)

    def test_no_camera_side_not_learned(self, qapp, tmp_path):
        # 该侧无摄像头 → 不自动标定
        mw, bucket = self._make_mw(qapp, tmp_path)
        bucket.cameras = []
        mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
        self._ready_trigger(mw)
        assert bucket.hook_lug_x_offset_ratio_left is None

    def test_no_active_bucket_not_learned(self, qapp, tmp_path):
        # 无活动工位 → 不自动标定
        mw, bucket = self._make_mw(qapp, tmp_path)
        mw.active_bucket = None
        mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
        self._ready_trigger(mw)
        assert bucket.hook_lug_x_offset_ratio_left is None

    def test_three_ready_within_window_writes(self, qapp, tmp_path):
        mw, bucket = self._make_mw(qapp, tmp_path)
        mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
        self._ready_trigger(mw)
        assert bucket.hook_lug_x_offset_ratio_left == pytest.approx(0.0, abs=1e-6)
        assert bucket.hook_lug_y_offset_ratio_left == pytest.approx(0.0, abs=1e-6)

    def test_two_ready_not_enough(self, qapp, tmp_path):
        mw, bucket = self._make_mw(qapp, tmp_path)
        mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
        for _ in range(2):
            mw._update_auto_calibration("left", True)
        assert bucket.hook_lug_x_offset_ratio_left is None

    def test_interleaved_not_ready_tolerated(self, qapp, tmp_path):
        # 非连续：1 秒窗口内 3 帧 ready 即可触发，中间的非 ready 帧不清零
        mw, bucket = self._make_mw(qapp, tmp_path)
        mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
        mw._update_auto_calibration("left", True)
        mw._update_auto_calibration("left", False)
        mw._update_auto_calibration("left", True)
        mw._update_auto_calibration("left", True)
        assert bucket.hook_lug_x_offset_ratio_left == pytest.approx(0.0, abs=1e-6)

    def test_frames_outside_window_not_counted(self, qapp, tmp_path):
        # 超过窗口时长的 ready 帧互不累计
        mw, bucket = self._make_mw(qapp, tmp_path)
        mw.config.ai.auto_calib_duration_secs = 0.002
        mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
        for _ in range(3):
            mw._update_auto_calibration("left", True)
            time.sleep(0.005)
        assert bucket.hook_lug_x_offset_ratio_left is None

    def test_calibrated_side_not_overwritten(self, qapp, tmp_path):
        # 已标定侧 → 自动标定不应覆盖
        mw, bucket = self._make_mw(qapp, tmp_path)
        bucket.hook_lug_x_offset_ratio_left = -0.25
        bucket.hook_lug_y_offset_ratio_left = 0.125
        mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
        self._ready_trigger(mw)
        assert bucket.hook_lug_x_offset_ratio_left == pytest.approx(-0.25, abs=1e-6)
        assert bucket.hook_lug_y_offset_ratio_left == pytest.approx(0.125, abs=1e-6)

    def test_both_sides_calibrated_automatically(self, qapp, tmp_path):
        # 无需入队：左右未标定侧连续识别 ready 自动全部标定
        mw, bucket = self._make_mw(qapp, tmp_path)
        for side in ("left", "right"):
            mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
            mw.right_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
            self._ready_trigger(mw, side)
        assert bucket.hook_lug_x_offset_ratio_left is not None
        assert bucket.hook_lug_x_offset_ratio_right is not None
        assert bucket.id not in getattr(mw, "_auto_calib_pending", set())

    def test_reset_then_auto_recalibrates(self, qapp, tmp_path):
        # 重置后再次 ready 自动重新标定
        mw, bucket = self._make_mw(qapp, tmp_path)
        bucket.hook_lug_x_offset_ratio_left = -0.5
        bucket.hook_lug_y_offset_ratio_left = -0.5
        assert mw._is_bucket_calibrated(bucket) is True
        mw._reset_calibration()
        assert bucket.hook_lug_x_offset_ratio_left is None
        mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
        self._ready_trigger(mw)
        assert bucket.hook_lug_x_offset_ratio_left == pytest.approx(0.0, abs=1e-6)

    def test_reset_calib_button_exists(self, qapp):
        mw = MainWindow(Config())
        assert mw.reset_calib_btn is not None
        assert mw.reset_calib_btn.text() == "重置标定"

    def test_custom_ready_frames_honored(self, qapp, tmp_path):
        mw, bucket = self._make_mw(qapp, tmp_path)
        mw.config.ai.auto_calib_ready_frames = 2
        mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
        mw._update_auto_calibration("left", True)
        assert bucket.hook_lug_x_offset_ratio_left is None
        mw._update_auto_calibration("left", True)
        assert bucket.hook_lug_x_offset_ratio_left == pytest.approx(0.0, abs=1e-6)

    def test_custom_min_samples_not_written_too_early(self, qapp, tmp_path):
        mw, bucket = self._make_mw(qapp, tmp_path)
        mw.config.ai.auto_calib_min_samples = 5
        mw.left_objects = [_obj("hook", 40, 40), _obj("lug", 40, 40)]
        self._ready_trigger(mw)
        assert bucket.hook_lug_x_offset_ratio_left is None