tests/test_datagen_plate.py

"""Synthetic, data-free tests for fpgm.datagen.plate.

Tiny videos are written with cv2 (MJPG, matching tests/test_frames.py's
approach) into ``tmp_path`` -- no real DROID mp4 needed. No GPU, no network.
"""

from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np
import pytest

from fpgm.datagen.cache import StageCache
from fpgm.datagen.plate import (
    NonStaticCameraWarning,
    PlateStage,
    compute_background_plate,
)
from fpgm.types import DataError

_SIZE = (40, 30)  # (width, height)


def _background_frame(width: int, height: int, x_offset: int = 0) -> np.ndarray:
    """A spatially-varying BGR pattern, all three channels depending on ``x``
    (offset by ``x_offset``) -- unlike a flat colour, this is distinguishable
    per pixel (catching the median silently blurring detail) and, because
    every channel moves together under a horizontal shift, a panning camera
    shows up strongly in the corner-patch check regardless of which channel
    happens to dominate that check's median.
    """
    xx, _yy = np.meshgrid(np.arange(width) - x_offset, np.arange(height))
    b = (xx * 5) % 256
    g = (xx * 5 + 85) % 256
    r = (xx * 5 + 170) % 256
    return np.stack([b, g, r], axis=-1).astype(np.uint8)


def _write_static_video_with_moving_blob(path: Path, n_frames: int = 16) -> np.ndarray:
    """Fixed background + a small bright blob that visits a different location
    every frame (kept away from the corner-check region), so no single pixel
    is covered by the blob more than once -- the temporal median at every
    pixel must therefore equal the background.

    Returns the ground-truth background frame (BGR) for the caller to compare
    against.
    """
    width, height = _SIZE
    bg = _background_frame(width, height)
    writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"MJPG"), 10.0, (width, height))
    try:
        for i in range(n_frames):
            frame = bg.copy()
            bx = 10 + (3 * i) % (width - 14)
            by = 10 + (5 * i) % (height - 14)
            frame[by : by + 4, bx : bx + 4] = (255, 255, 255)
            writer.write(frame)
    finally:
        writer.release()
    return bg


def _write_panning_video(path: Path, n_frames: int = 16) -> None:
    """A camera pan: the whole background pattern shifts a few pixels every
    frame, so a fixed pixel location sees genuinely different scene content
    over time -- the failure mode compute_background_plate must detect."""
    width, height = _SIZE
    writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"MJPG"), 10.0, (width, height))
    try:
        for i in range(n_frames):
            writer.write(_background_frame(width, height, x_offset=4 * i))
    finally:
        writer.release()


class TestComputeBackgroundPlateStaticCamera:
    def test_median_plate_recovers_background(self, tmp_path: Path):
        video_path = tmp_path / "video.avi"
        bg = _write_static_video_with_moving_blob(video_path)

        plate_rgb, stats = compute_background_plate(video_path)

        bg_rgb = cv2.cvtColor(bg, cv2.COLOR_BGR2RGB)
        # MJPG is a lossy, block-based codec: a handful of pixels right next
        # to the sharp white blob can pick up a one-frame DCT ringing
        # artifact that median-of-16 doesn't fully wash out, so this checks
        # the *typical* pixel (mean), not a bit-exact match. "No blob pixel
        # anywhere" is checked separately and more strictly below.
        diff = np.abs(plate_rgb.astype(int) - bg_rgb.astype(int))
        assert diff.mean() <= 5.0
        assert stats.static_camera_ok
        assert stats.n_frames == 16

    def test_blob_never_dominates_any_single_pixel(self, tmp_path: Path):
        video_path = tmp_path / "video.avi"
        _write_static_video_with_moving_blob(video_path)
        plate_rgb, _ = compute_background_plate(video_path)
        # If the blob had leaked through, some pixel would read near-white.
        assert not (plate_rgb > 240).all(axis=-1).any()

    def test_no_warning_raised_for_a_static_camera(self, tmp_path: Path, recwarn):
        video_path = tmp_path / "video.avi"
        _write_static_video_with_moving_blob(video_path)
        compute_background_plate(video_path)
        assert not any(issubclass(w.category, NonStaticCameraWarning) for w in recwarn.list)


class TestComputeBackgroundPlatePanning:
    def test_panning_camera_triggers_static_camera_warning(self, tmp_path: Path):
        video_path = tmp_path / "panning.avi"
        _write_panning_video(video_path)

        with pytest.warns(NonStaticCameraWarning):
            _plate_rgb, stats = compute_background_plate(video_path)

        assert not stats.static_camera_ok
        assert stats.corner_median_abs_diff > 6.0


class TestComputeBackgroundPlateErrors:
    def test_undecodable_video_raises_data_error(self, tmp_path: Path):
        bogus = tmp_path / "not_a_video.avi"
        bogus.write_bytes(b"not a real video file")
        with pytest.raises(DataError):
            compute_background_plate(bogus)


class TestPlateStage:
    def test_writes_plate_png_and_meta(self, tmp_path: Path):
        video_path = tmp_path / "video.avi"
        _write_static_video_with_moving_blob(video_path)
        cache = StageCache(tmp_path / "cache_root")

        stage = PlateStage()
        plate_rgb, stats = stage.run(video_path=video_path, cache=cache)

        plate_path = cache.root / "master" / "plate.png"
        assert plate_path.exists()
        assert plate_rgb.shape == (_SIZE[1], _SIZE[0], 3)
        assert stats.static_camera_ok

    def test_rerun_with_same_video_reuses_cache(self, tmp_path: Path, monkeypatch):
        video_path = tmp_path / "video.avi"
        _write_static_video_with_moving_blob(video_path)
        cache = StageCache(tmp_path / "cache_root")
        stage = PlateStage()

        stage.run(video_path=video_path, cache=cache)
        plate_path = cache.root / "master" / "plate.png"
        first_mtime = plate_path.stat().st_mtime_ns

        # Poison compute_background_plate so a second run would be detectably
        # wrong if the cache were NOT hit -- this proves reuse, not just "ran
        # again and got a similar-looking answer".
        import fpgm.datagen.plate as plate_mod

        def _boom(*args, **kwargs):
            raise AssertionError("compute_background_plate should not run on a cache hit")

        monkeypatch.setattr(plate_mod, "compute_background_plate", _boom)

        stage.run(video_path=video_path, cache=cache)
        assert plate_path.stat().st_mtime_ns == first_mtime

    def test_different_video_invalidates_cache(self, tmp_path: Path):
        video_path = tmp_path / "video.avi"
        _write_static_video_with_moving_blob(video_path)
        cache = StageCache(tmp_path / "cache_root")
        stage = PlateStage()
        stage.run(video_path=video_path, cache=cache)

        # A different file at the same path (different mtime/size) must not
        # be treated as fresh.
        video_path2 = tmp_path / "video2.avi"
        _write_panning_video(video_path2)
        _plate_rgb, stats = stage.run(video_path=video_path2, cache=cache)
        assert not stats.static_camera_ok