tests/test_tensor_api.py

from __future__ import annotations

import pytest


def torch_or_skip():
    return pytest.importorskip("torch")


def test_tensor_to_pil_images_quantizes_rgb_range():
    torch = torch_or_skip()
    from uw_nriqa.tensor import tensor_to_pil_images

    images = torch.tensor([[[[0.0, 0.5], [1.0, 2.0]], [[-1.0, 0.25], [0.75, 1.0]], [[0.1, 0.2], [0.3, 0.4]]]])

    pil = tensor_to_pil_images(images)

    assert len(pil) == 1
    assert pil[0].mode == "RGB"
    assert pil[0].size == (2, 2)
    assert pil[0].getpixel((0, 0)) == (0, 0, 26)
    assert pil[0].getpixel((1, 0)) == (128, 64, 51)
    assert pil[0].getpixel((0, 1)) == (255, 191, 77)
    assert pil[0].getpixel((1, 1)) == (255, 255, 102)


def test_tensor_to_pil_images_rejects_non_bchw_rgb():
    torch = torch_or_skip()
    from uw_nriqa.tensor import tensor_to_pil_images

    with pytest.raises(ValueError, match="BCHW RGB"):
        tensor_to_pil_images(torch.zeros(1, 1, 8, 8))


def test_tensor_direct_rejects_unsupported_methods():
    torch = torch_or_skip()
    from uw_nriqa.tensor import score_tensor

    with pytest.raises(ValueError, match="unsupported"):
        score_tensor(torch.zeros(1, 3, 8, 8), methods=["tuda"], preload=False)
    with pytest.raises(ValueError, match="rsuia"):
        score_tensor(torch.zeros(1, 3, 8, 8), methods=["rsuia"], preload=False)


def test_tensor_direct_traditional_metrics_do_not_require_pil(monkeypatch):
    torch = torch_or_skip()
    from uw_nriqa.tensor import score_tensor

    def fail_pil(_images):
        raise AssertionError("PIL conversion should not be used for tensor metrics")

    monkeypatch.setattr("uw_nriqa.tensor.tensor_to_pil_images", fail_pil)
    images = torch.rand(4, 3, 64, 64)

    scores = score_tensor(images, methods=["uciqe", "uiqm"], preload=False)

    assert scores.shape == (4, 2)
    assert torch.isfinite(scores).all()


def test_tensor_uciqe_tracks_legacy_metric():
    torch = torch_or_skip()

    from uw_nriqa.metrics import calculate_uciqe
    from uw_nriqa.tensor import score_tensor, tensor_to_pil_images

    images = torch.rand(3, 3, 64, 64)
    tensor_scores = score_tensor(images, methods=["uciqe"], preload=False)[:, 0]
    legacy_scores = torch.tensor(
        [calculate_uciqe(pil) for pil in tensor_to_pil_images(images)],
        dtype=torch.float32,
    )

    assert torch.allclose(tensor_scores, legacy_scores, atol=0.02)


def test_tensor_uiqm_tracks_legacy_metric():
    torch = torch_or_skip()
    import numpy as np

    from uw_nriqa.metrics import calculate_uiqm
    from uw_nriqa.tensor import score_tensor, tensor_to_pil_images

    rng = np.random.default_rng(20260602)
    arrays = rng.integers(0, 256, size=(4, 72, 96, 3), dtype=np.uint8)
    images = torch.from_numpy(arrays.copy()).permute(0, 3, 1, 2).to(dtype=torch.float32) / 255.0
    tensor_scores = score_tensor(images, methods=["uiqm"], preload=False)[:, 0]
    legacy_scores = torch.tensor(
        [calculate_uiqm(pil) for pil in tensor_to_pil_images(images)],
        dtype=torch.float32,
    )

    assert torch.allclose(tensor_scores, legacy_scores, atol=1e-4, rtol=1e-4)


def test_tensor_niqe_tracks_legacy_metric():
    pytest.importorskip("pyiqa")
    torch = torch_or_skip()
    import numpy as np

    from uw_nriqa.metrics import calculate_niqe
    from uw_nriqa.tensor import score_tensor, tensor_to_pil_images

    rng = np.random.default_rng(20260602)
    arrays = rng.integers(0, 256, size=(3, 256, 256, 3), dtype=np.uint8)
    images = torch.from_numpy(arrays.copy()).permute(0, 3, 1, 2).to(dtype=torch.float32) / 255.0
    tensor_scores = score_tensor(images, methods=["niqe"], preload=False)[:, 0]
    legacy_scores = torch.tensor(
        [calculate_niqe(pil) for pil in tensor_to_pil_images(images)],
        dtype=torch.float32,
    )

    assert torch.allclose(tensor_scores, legacy_scores, atol=1e-4, rtol=1e-4)


def test_tensor_local_nuiq_tracks_image_batch_path():
    torch = torch_or_skip()
    import numpy as np
    from PIL import Image

    from uw_nriqa.nuiq import calculate_nuiq_local
    from uw_nriqa.tensor import score_tensor

    rng = np.random.default_rng(123)
    arrays = rng.integers(0, 256, size=(4, 36, 44, 3), dtype=np.uint8)
    images = [Image.fromarray(arr) for arr in arrays]
    tensor = torch.from_numpy(arrays).permute(0, 3, 1, 2).to(dtype=torch.float32) / 255.0

    image_scores = torch.tensor(calculate_nuiq_local(images), dtype=torch.float32)
    tensor_scores = score_tensor(tensor, methods=["nuiq_local"], preload=False)[:, 0]

    assert torch.all(torch.isfinite(tensor_scores))
    assert torch.allclose(tensor_scores, image_scores, atol=1e-3)


def test_predict_tensor_uses_tensor_scorer(monkeypatch):
    torch = torch_or_skip()
    from uw_nriqa import NRIQAEvaluator

    def fake_score_tensor_with_evaluator(evaluator, images):
        assert evaluator.methods == ["uranker", "uciqe"]
        assert tuple(images.shape) == (2, 3, 8, 8)
        return torch.tensor([[1.0, 0.5], [2.0, 0.7]], dtype=torch.float32)

    monkeypatch.setattr("uw_nriqa.tensor.score_tensor_with_evaluator", fake_score_tensor_with_evaluator)
    evaluator = NRIQAEvaluator(["uranker", "uciqe"], preload=False)

    scores = evaluator.predict_tensor(torch.zeros(2, 3, 8, 8), raise_errors=True)

    assert scores == {"uranker": [1.0, 2.0], "uciqe": [0.5, pytest.approx(0.7)]}


def test_uranker_histogram_matches_legacy_histc():
    torch = torch_or_skip()
    from uw_nriqa.vendor.UnderwaterRanker import utils

    images = torch.tensor(
        [
            [
                [[0.0, 0.1, 0.5], [1.0, -0.1, 1.2]],
                [[0.25, 0.25, 0.25], [0.75, 0.75, 0.75]],
                [[0.0, 1.0, 0.5], [0.25, 0.75, 0.33]],
            ],
            [
                [[0.2, 0.4, 0.6], [0.8, 1.0, 0.0]],
                [[0.9, 0.7, 0.5], [0.3, 0.1, 0.0]],
                [[1.1, -0.2, 0.4], [0.4, 0.4, 0.4]],
            ],
        ],
        dtype=torch.float32,
    )

    expected = []
    for image in images:
        channels = [torch.histc(channel, 64, min=0.0, max=1.0) for channel in image]
        expected.append(torch.cat(channels))
    expected = torch.stack(expected, dim=0).unsqueeze(1)

    actual = utils.build_historgram(images)

    assert actual.shape == (2, 1, 192)
    assert torch.equal(actual, expected)