tests/test_augment.py

from __future__ import annotations

from pathlib import Path
import random

from PIL import Image

import kneiff.datasets.augmentations.transforms as augmentation_transforms
import kneiff.datasets.augmentations.workflow as augmentation_workflow


def _write_image(path: Path, *, size: tuple[int, int] = (24, 18)) -> None:
    image = Image.new("RGB", size, color=(80, 100, 120))
    image.save(path)


def test_apply_transforms_is_repeatable_with_local_rng() -> None:
    image = Image.new("RGB", (24, 18), color=(80, 100, 120))

    first_image, first_loss = augmentation_transforms.apply_transforms(
        image,
        flip_lr=True,
        rng=random.Random(42),
    )
    second_image, second_loss = augmentation_transforms.apply_transforms(
        image,
        flip_lr=True,
        rng=random.Random(42),
    )

    assert first_loss == second_loss
    assert first_image.size == image.size
    assert first_image.tobytes() == second_image.tobytes()


def test_dataset_augmentation_writes_numbered_outputs_and_preserves_extension_filter(
    tmp_path: Path,
) -> None:
    input_dir = tmp_path / "input"
    output_dir = tmp_path / "output"
    input_dir.mkdir()
    _write_image(input_dir / "source.png")
    _write_image(input_dir / "ignored.tiff")

    augmentation_workflow.augment_dataset(
        input_dir=input_dir,
        output_dir=output_dir,
        num_aug=1,
        workers=1,
    )

    assert (output_dir / "source_aug_1.png").exists()
    assert not (output_dir / "ignored_aug_1.tiff").exists()


def test_augmentation_modules_expose_transform_and_workflow_entrypoints() -> None:
    assert callable(augmentation_transforms.apply_transforms)
    assert callable(augmentation_workflow.augment_dataset)