Newer
Older
TIASJudgeModel / tests / segmentation / test_detection.py
"""舌の検出と枠の扱い(`segmentation/detection.py`)のテスト.

仕様(DeepTIAS 2.7 の再現.実装サマリーを仕様として扱う):

- スコアがしきい値(0.5)を**超える**中で最高の枠を採用する(該当なしは None)
- 正規化 `[ymin, xmin, ymax, xmax]` を画素の枠(`x1`・`y1` は含まない)に変換する.画素への
  変換は切り捨て.画像の範囲にクリップし,幅・高さが 0 以下なら None
- `mask_bounding_rect` は前景の外接矩形(前景が無ければ None)
- `OnnxDetector` は入出力名を `:0` の有無を問わずに探す.ファイルが無い・名前が無いと ValueError
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

from tias_judge_model.segmentation.detection import (
    OnnxDetector,
    Rect,
    full_rect,
    mask_bounding_rect,
    normalized_box_to_rect,
    select_best_box,
)

# --- Rect ---


def test_rect_width_height_and_full_rect() -> None:
    r = full_rect(30, 50)
    assert r == Rect(0, 0, 50, 30)
    assert (r.width, r.height) == (50, 30)


# --- mask_bounding_rect ---


def test_mask_bounding_rect_is_exclusive_on_right_and_bottom() -> None:
    mask = np.zeros((20, 30), bool)
    mask[5:9, 10:17] = True
    mask[12, 3] = True
    r = mask_bounding_rect(mask)
    assert r == Rect(3, 5, 17, 13)


def test_mask_bounding_rect_single_pixel_and_full() -> None:
    mask = np.zeros((4, 4), bool)
    mask[2, 1] = True
    assert mask_bounding_rect(mask) == Rect(1, 2, 2, 3)
    assert mask_bounding_rect(np.ones((4, 6), bool)) == Rect(0, 0, 6, 4)


def test_mask_bounding_rect_empty_is_none() -> None:
    assert mask_bounding_rect(np.zeros((5, 5), bool)) is None


# --- normalized_box_to_rect ---


def test_normalized_box_to_rect_basic() -> None:
    r = normalized_box_to_rect([0.1, 0.2, 0.5, 0.6], height=100, width=200)
    assert r == Rect(40, 10, 120, 50)


def test_normalized_box_to_rect_truncates() -> None:
    # 0.29 × 100 = 28.999999999999996 → 28(切り捨て),0.999 × 10 = 9.99 → 9
    r = normalized_box_to_rect([0.29, 0.0, 0.999, 0.999], height=100, width=10)
    assert r == Rect(0, 28, 9, 99)


def test_normalized_box_to_rect_clips_outside_values() -> None:
    r = normalized_box_to_rect([-0.2, -0.1, 1.3, 1.5], height=40, width=60)
    assert r == Rect(0, 0, 60, 40)


@pytest.mark.parametrize(
    "box",
    [
        [0.5, 0.5, 0.5, 0.8],  # 高さ 0
        [0.2, 0.5, 0.6, 0.5],  # 幅 0
        [0.6, 0.2, 0.3, 0.8],  # 上下が逆
        [0.2, 0.8, 0.6, 0.3],  # 左右が逆
        [1.2, 0.1, 1.5, 0.5],  # 全体が画像の外(下)
        [0.1, -0.5, 0.5, -0.1],  # 全体が画像の外(左)
        [0.10, 0.1, 0.105, 0.5],  # 切り捨てで高さ 0(10 → 10)
    ],
)
def test_normalized_box_to_rect_degenerate_is_none(box: list[float]) -> None:
    assert normalized_box_to_rect(box, height=100, width=100) is None


# --- select_best_box ---


def test_select_best_box_picks_highest_above_threshold() -> None:
    boxes = np.array([[0, 0, 0.5, 0.5], [0.1, 0.1, 0.9, 0.9], [0.2, 0.2, 0.3, 0.3]])
    scores = np.array([0.6, 0.95, 0.7])
    box, score = select_best_box(boxes, scores, 0.5)
    np.testing.assert_array_equal(box, boxes[1])
    assert score == pytest.approx(0.95)


def test_select_best_box_threshold_is_strict() -> None:
    boxes = np.array([[0, 0, 1, 1]])
    assert select_best_box(boxes, np.array([0.5]), 0.5) is None
    assert select_best_box(boxes, np.array([0.5000001]), 0.5) is not None


def test_select_best_box_none_when_all_below_or_empty() -> None:
    assert select_best_box(np.array([[0, 0, 1, 1]] * 2), np.array([0.1, 0.3]), 0.5) is None
    assert select_best_box(np.zeros((0, 4)), np.zeros(0), 0.5) is None


def test_select_best_box_accepts_batched_shapes() -> None:
    # 検出モデルの出力の形((1, N, 4)・(1, N))のままでも扱える
    boxes = np.array([[[0, 0, 0.5, 0.5], [0.1, 0.2, 0.3, 0.4]]])
    scores = np.array([[0.3, 0.8]])
    box, _ = select_best_box(boxes, scores, 0.5)
    np.testing.assert_allclose(box, [0.1, 0.2, 0.3, 0.4])


def test_select_best_box_tie_takes_first() -> None:
    boxes = np.array([[0, 0, 0.5, 0.5], [0.5, 0.5, 1, 1]])
    box, _ = select_best_box(boxes, np.array([0.9, 0.9]), 0.5)
    np.testing.assert_array_equal(box, boxes[0])


# --- OnnxDetector(ダミーの ONNX) ---


def _detector(path: Path, **kw) -> OnnxDetector:
    args = {
        "score_threshold": 0.5,
        "input_color": "rgb",
        "input_name": "image_tensor",
        "boxes_name": "detection_boxes",
        "scores_name": "detection_scores",
    }
    args.update(kw)
    return OnnxDetector(path, **args)


def _red_patch_bgr() -> np.ndarray:
    bgr = np.full((80, 120, 3), (150, 120, 110), np.uint8)  # 青みがかった背景(BGR)
    bgr[20:50, 30:90] = (105, 95, 200)  # 赤い領域(BGR)
    return bgr


def test_detector_finds_red_region_with_colon_zero_names(dummy_seg_models: dict) -> None:
    det = _detector(dummy_seg_models["detector"])
    result = det.detect(_red_patch_bgr())
    assert result.rect == Rect(30, 20, 90, 50)
    assert result.score == pytest.approx(0.9)


def test_detector_below_threshold_is_not_detected_but_records_score(
    dummy_seg_models: dict,
) -> None:
    det = _detector(dummy_seg_models["detector"], score_threshold=0.95)
    result = det.detect(_red_patch_bgr())
    assert result.rect is None
    assert result.score == pytest.approx(0.9)


def test_detector_bgr_input_color_changes_what_model_sees(dummy_seg_models: dict) -> None:
    # input_color: bgr ならそのまま渡す → ダミーは R−B を見るため,赤(BGR の B が小さい)を
    # 青と取り違え,背景(BGR の B が大きい)を前景とみなす
    det = _detector(dummy_seg_models["detector"], input_color="bgr")
    result = det.detect(_red_patch_bgr())
    assert result.rect != Rect(30, 20, 90, 50)


def test_detector_rejects_missing_file(tmp_path: Path) -> None:
    with pytest.raises(ValueError):
        _detector(tmp_path / "none.onnx")


@pytest.mark.parametrize(
    "kw", [{"input_name": "images"}, {"boxes_name": "boxes"}, {"scores_name": "scores"}]
)
def test_detector_rejects_unknown_names(dummy_seg_models: dict, kw: dict) -> None:
    with pytest.raises(ValueError):
        _detector(dummy_seg_models["detector"], **kw)