orisys-sdk/tests/test_stitch.py

"""Unit tests for orisys.stitch (synthetic tiles, no camera)."""

from __future__ import annotations

import numpy as np
import pytest

from orisys.frame_processors import CompoundEyeFrameProcessor
from orisys.stitch import (
    StitchStatus,
    get_blend_backend,
    paste_tile_clipped,
    resolve_blend_mode,
    tile_weight_map,
)
from orisys.stitch.backends.numpy_backend import blend_kernel_numpy
from orisys.stitch.config import numba_blend_available, validate_box_grid
from orisys.stitch.ops import compute_final_crop_rect, coverage_ratio
from orisys.Stitcher import Stitcher


def _solid_tile(val: int, h=32, w=32):
    return np.full((h, w, 3), val, dtype=np.uint8)


def test_tile_weight_map_edges_are_low():
    wmap = tile_weight_map(32, 32, margin=4)
    assert wmap.shape == (32, 32)
    assert wmap.dtype == np.float32
    assert float(wmap[0, 0]) < float(wmap[16, 16])
    assert float(wmap[16, 16]) == pytest.approx(1.0)


def test_paste_tile_clipped_and_coverage():
    canvas = np.zeros((40, 40, 3), dtype=np.uint8)
    cov = np.zeros((40, 40), dtype=np.uint8)
    tile = _solid_tile(200, 10, 10)
    paste_tile_clipped(canvas, tile, 5, 5, cov=cov)
    assert int(canvas[5, 5, 0]) == 200
    assert int(cov[5, 5]) == 1
    assert int(cov[0, 0]) == 0


def test_blend_kernel_numpy_first_paint():
    blend = np.zeros((4, 4, 3), dtype=np.float32)
    wsum = np.zeros((4, 4), dtype=np.float32)
    cov = np.zeros((4, 4), dtype=np.uint8)
    src = np.ones((4, 4, 3), dtype=np.float32) * 50
    weight = np.ones((4, 4), dtype=np.float32)
    blend_kernel_numpy(blend, wsum, cov, src, weight)
    assert np.allclose(blend, 50)
    assert np.all(cov == 1)


def test_numpy_backend_blend_tile():
    backend = get_blend_backend("numpy")
    blend = np.zeros((20, 20, 3), dtype=np.float32)
    wsum = np.zeros((20, 20), dtype=np.float32)
    cov = np.zeros((20, 20), dtype=np.uint8)
    src = np.full((8, 8, 3), 40, dtype=np.uint8)
    weight = tile_weight_map(8, 8, margin=2)
    backend.blend_tile_inplace(blend, wsum, cov, src, weight, 2, 3)
    assert float(blend[4, 5, 0]) > 0
    assert coverage_ratio(cov) > 0


@pytest.mark.skipif(not numba_blend_available(), reason="numba blend unavailable")
def test_numba_numpy_backend_parity():
    numpy_b = get_blend_backend("numpy")
    numba_b = get_blend_backend("numba")
    shape = (16, 16)
    src = np.random.randint(0, 255, (*shape, 3), dtype=np.uint8)
    weight = tile_weight_map(*shape, margin=3)

    def run(backend):
        blend = np.zeros((*shape, 3), dtype=np.float32)
        wsum = np.zeros(shape, dtype=np.float32)
        cov = np.zeros(shape, dtype=np.uint8)
        backend.blend_tile_inplace(blend, wsum, cov, src, weight, 0, 0)
        return blend, cov

    b0, c0 = run(numpy_b)
    b1, c1 = run(numba_b)
    assert np.allclose(b0, b1, atol=1e-2)
    assert np.array_equal(c0, c1)


def test_resolve_blend_mode_default(monkeypatch):
    monkeypatch.delenv("ORISYS_STITCH_BLEND", raising=False)
    monkeypatch.delenv("ORISYS_STITCH_BLEND_BACKEND", raising=False)
    mode, fell = resolve_blend_mode({})
    assert mode == "none"
    assert fell is False


def test_validate_box_grid():
    assert validate_box_grid([(0, 0, 1, 1)] * 9, 3, 3)
    assert not validate_box_grid([(0, 0, 1, 1)] * 8, 3, 3)


def test_compute_final_crop_rect():
    cov = np.zeros((20, 20), dtype=np.uint8)
    cov[4:16, 4:16] = 1
    rect = compute_final_crop_rect(cov)
    assert rect is not None
    y0, y1, x0, x1 = rect
    assert y1 > y0 and x1 > x0


def _mini_config():
    return {
        "min_area_frac": 0.01,
        "max_area_frac": 0.05,
        "max_aspect_ratio": 50,
        "index": [0, 1, 2, 3, 4, 5, 6, 7, 8],
        "margin": 10,
        "nominal_ov_pos": 100,
        "n": [3, 3],
        "row_tol": 0.6,
        "box_size": [32, 32],
        "blend": "numpy",
        "n_boxes_target": 9,
    }


def test_stitcher_status_and_rt_with_synthetic_tiles():
    stitcher = Stitcher(isstitch=True, config_stitcher=_mini_config(), verbose=False)
    # Preload tiles as if extract_tiles already ran.
    tiles = [_solid_tile(30 + i * 10) for i in range(9)]
    stitcher.tile_list = tiles
    stitcher.units = [(0, 0, t) for t in tiles]
    stitcher.boxes = [(i * 40, 0, i * 40 + 32, 32) for i in range(9)]

    status = stitcher.status
    assert isinstance(status, StitchStatus)
    assert status.blend_mode == "numpy"
    assert status.expected_boxes == 9

    # Feature matching on flat solid tiles usually fails → None is OK.
    out = stitcher.stitch_rt_optimized()
    assert out is None or (isinstance(out, np.ndarray) and out.ndim == 3)


def test_compound_eye_processor_calls_stitcher():
    class FakeStitcher:
        def __init__(self):
            self.calls = []

        def get_stitch_unit(self, frame):
            self.calls.append("get")

        def extract_tiles(self):
            self.calls.append("extract")

        def stitch_rt_optimized(self):
            self.calls.append("stitch")
            return np.zeros((8, 8, 3), dtype=np.uint8)

    fake = FakeStitcher()
    proc = CompoundEyeFrameProcessor(fake)
    out = proc.process(np.zeros((64, 64, 3), dtype=np.uint8))
    assert fake.calls == ["get", "extract", "stitch"]
    assert out.shape == (8, 8, 3)


def test_unknown_backend_raises():
    with pytest.raises(ValueError):
        get_blend_backend("torch")