tests/transforms/test_physics.py

"""Tests for physics-aware augmentations."""

from __future__ import annotations

import numpy as np
import pytest

from dipauglib.transforms.physics import (
    CastShadow,
    ColourFade,
    ColourTempShift,
    DefocusBlur,
    DustOverlay,
    IlluminationGradient,
    MotionBlur,
    SensorNoise,
)


TRANSFORMS = [
    IlluminationGradient,
    CastShadow,
    MotionBlur,
    DefocusBlur,
    ColourTempShift,
    ColourFade,
    DustOverlay,
    SensorNoise,
]


@pytest.mark.parametrize("transform_cls", TRANSFORMS)
def test_transform_preserves_shape_and_mask(transform_cls):
    image = np.full((32, 32, 3), 128, dtype=np.uint8)
    mask = np.zeros((32, 32), dtype=np.uint8)
    mask[8:20, 10:22] = 1
    transform = transform_cls(intensity=0.5, p=1.0)
    output = transform(image=image, mask=mask)
    assert output["image"].shape == image.shape
    assert output["mask"].shape == mask.shape
    assert np.array_equal(output["mask"], mask)


@pytest.mark.parametrize("transform_cls", TRANSFORMS)
def test_transform_handles_all_black_image(transform_cls):
    image = np.zeros((24, 24, 3), dtype=np.uint8)
    mask = np.zeros((24, 24), dtype=np.uint8)
    transform = transform_cls(intensity=1.0, p=1.0)
    output = transform(image=image, mask=mask)
    assert output["image"].shape == image.shape
    assert output["mask"].shape == mask.shape


@pytest.mark.parametrize("transform_cls", TRANSFORMS)
def test_transform_handles_single_pixel_image(transform_cls):
    image = np.array([[[255, 0, 0]]], dtype=np.uint8)
    mask = np.array([[1]], dtype=np.uint8)
    transform = transform_cls(intensity=0.2, p=1.0)
    output = transform(image=image, mask=mask)
    assert output["image"].shape == image.shape
    assert output["mask"].shape == mask.shape