tests/test_image_processor.py

import math
import unittest

import torch

from data.image_processor import (
    convert_image_to_patches,
    get_aspect_ratio_preserving_size,
    pad_along_first_dim,
    patches_merge,
)


def reference_patchify(image: torch.Tensor, patch_size: int) -> torch.Tensor:
    """
    Independent, loop-based re-implementation of `convert_image_to_patches`, used as
    an oracle so the tests do not merely restate the implementation under test.

    Parameters:
     * image (torch.Tensor) : a (C, H, W) image tensor; requires H and W to be
         divisible by patch_size

     * patch_size (int) : length, in pixels, of one side of a patch; requires
         patch_size >= 1

    Returns:
        A (num_patches, patch_size * patch_size * C) tensor whose rows enumerate
        patches in row-major (top-to-bottom, left-to-right) order, and where each
        row flattens its patch in (row, col, channel) order.
    """
    num_channels, height, width = image.shape
    num_patches_height = height // patch_size
    num_patches_width = width // patch_size

    rows = []
    for i in range(num_patches_height):           # iterate patch rows top -> bottom
        for j in range(num_patches_width):         # iterate patch cols left -> right
            patch = image[:, i * patch_size:(i + 1) * patch_size,
                          j * patch_size:(j + 1) * patch_size]   # (C, ps, ps)
            # Flatten in (row, col, channel) order to match the source permute (1,3,2,4,0)
            rows.append(patch.permute(1, 2, 0).reshape(-1))
    return torch.stack(rows)


class TestGetAspectRatioPreservingSize(unittest.TestCase):
    """Behavioral spec for `get_aspect_ratio_preserving_size`."""

    # Parameters matching the worked example documented in image_processor.py
    TEACHER_PATCH_SIZE = 16
    MAX_TEACHER_PATCHES = 1500
    POOLING_KERNEL_SIZE = 3

    def _size(self, height, width):
        # Thin wrapper that fixes the patch-budget parameters for the common case.
        return get_aspect_ratio_preserving_size(
            height=height,
            width=width,
            teacher_patch_size=self.TEACHER_PATCH_SIZE,
            max_teacher_patches=self.MAX_TEACHER_PATCHES,
            pooling_kernel_size=self.POOLING_KERNEL_SIZE,
        )

    @property
    def _side_mult(self):
        # A valid output side must be a whole number of model patches wide/tall.
        return self.POOLING_KERNEL_SIZE * self.TEACHER_PATCH_SIZE

    def test_documented_worked_example(self):
        # The docstring walks through 500x600 -> 528x672; pin that exact contract.
        self.assertEqual(self._size(500, 600), (528, 672))

    def test_dimensions_divisible_by_model_patch_size(self):
        # Every returned side must be a multiple of pooling_kernel_size * teacher_patch_size,
        # otherwise the merge step could not tile the image into model patches.
        for height, width in [(500, 600), (800, 800), (333, 1024), (1080, 1920), (97, 640)]:
            target_h, target_w = self._size(height, width)
            self.assertEqual(target_h % self._side_mult, 0, msg=f"h for {height}x{width}")
            self.assertEqual(target_w % self._side_mult, 0, msg=f"w for {height}x{width}")

    def test_never_exceeds_patch_budget(self):
        # The whole point of the function: the patchified result fits the teacher budget.
        max_px = self.MAX_TEACHER_PATCHES * self.TEACHER_PATCH_SIZE ** 2
        for height, width in [(500, 600), (800, 800), (333, 1024), (1080, 1920), (97, 640)]:
            target_h, target_w = self._size(height, width)
            num_patches = (target_h // self.TEACHER_PATCH_SIZE) * (target_w // self.TEACHER_PATCH_SIZE)
            self.assertLessEqual(num_patches, self.MAX_TEACHER_PATCHES, msg=f"{height}x{width}")
            self.assertLessEqual(target_h * target_w, max_px, msg=f"{height}x{width}")

    def test_square_image_stays_square(self):
        # Aspect ratio 1:1 in must yield aspect ratio 1:1 out.
        target_h, target_w = self._size(640, 640)
        self.assertEqual(target_h, target_w)

    def test_transpose_symmetry(self):
        # Swapping height and width must transpose the result: the function treats the
        # two axes symmetrically, so f(h, w) == reverse(f(w, h)).
        target_h, target_w = self._size(500, 600)
        swapped_h, swapped_w = self._size(600, 500)
        self.assertEqual((target_h, target_w), (swapped_w, swapped_h))

    def test_raises_when_both_dimensions_round_to_zero(self):
        # A tiny image against a tiny budget rounds both ideal sides below one model
        # patch; the function must refuse rather than emit a 0x0 image.
        with self.assertRaises(ValueError):
            get_aspect_ratio_preserving_size(
                height=10, width=10, teacher_patch_size=16,
                max_teacher_patches=1, pooling_kernel_size=3,
            )

    def test_extreme_aspect_ratio_uses_fallback(self):
        # When one ideal side rounds to zero but the other does not, the degenerate side
        # is clamped up to a single model patch and the long side is capped by the budget.
        # 20x2000 -> (48, 96): height collapses to side_mult=48, width capped at
        # max_side_length = (20 // 3**2) * 48 = 96. The transpose maps to (96, 48).
        wide = get_aspect_ratio_preserving_size(
            height=20, width=2000, teacher_patch_size=16,
            max_teacher_patches=20, pooling_kernel_size=3,
        )
        tall = get_aspect_ratio_preserving_size(
            height=2000, width=20, teacher_patch_size=16,
            max_teacher_patches=20, pooling_kernel_size=3,
        )
        self.assertEqual(wide, (48, 96))
        self.assertEqual(tall, (96, 48))


class TestConvertImageToPatches(unittest.TestCase):
    """Behavioral spec for `convert_image_to_patches`."""

    def test_output_shape(self):
        # (C, H, W) -> (num_patches_h * num_patches_w, patch_size**2 * C).
        image = torch.rand(3, 32, 48)
        patches = convert_image_to_patches(image, patch_size=16)
        # 32/16 * 48/16 = 2 * 3 = 6 patches, each 16*16*3 = 768 long.
        self.assertEqual(patches.shape, (6, 768))

    def test_row_major_patch_ordering(self):
        # With patch_size=1 each pixel is its own patch, so the flattened output reveals
        # the iteration order directly: top-left, top-right, bottom-left, bottom-right.
        image = torch.tensor([[[10.0, 20.0], [30.0, 40.0]]])  # (1, 2, 2)
        patches = convert_image_to_patches(image, patch_size=1)
        self.assertTrue(torch.equal(patches.reshape(-1), torch.tensor([10.0, 20.0, 30.0, 40.0])))

    def test_matches_reference_implementation(self):
        # Cross-check against an independent loop-based patchifier for a multi-channel,
        # multi-patch image. This guards both patch order and within-patch flattening.
        image = torch.arange(3 * 4 * 6).float().reshape(3, 4, 6)
        patches = convert_image_to_patches(image, patch_size=2)
        self.assertTrue(torch.equal(patches, reference_patchify(image, patch_size=2)))

    def test_patchification_is_invertible(self):
        # No pixel may be dropped or duplicated: reversing the documented reshape/permute
        # must reconstruct the original image exactly.
        num_channels, height, width, patch_size = 3, 8, 12, 4
        image = torch.arange(num_channels * height * width).float().reshape(num_channels, height, width)
        patches = convert_image_to_patches(image, patch_size)

        nph, npw = height // patch_size, width // patch_size
        # Undo the (num_patches, ps*ps*C) flattening back to spatial layout.
        restored = (
            patches.reshape(nph, npw, patch_size, patch_size, num_channels)
            .permute(4, 0, 2, 1, 3)               # (C, nph, ps, npw, ps)
            .reshape(num_channels, height, width)
        )
        self.assertTrue(torch.equal(restored, image))


class TestPadAlongFirstDim(unittest.TestCase):
    """Behavioral spec for `pad_along_first_dim`."""

    def _make(self):
        # Three real patches with distinct positions; small flat dim for readability.
        flat_patches = torch.arange(12).float().reshape(3, 4)
        positions = torch.tensor([[0, 0], [0, 1], [1, 0]])
        return flat_patches, positions

    def test_output_shapes_after_padding(self):
        flat_patches, positions = self._make()
        padded_patches, padded_positions = pad_along_first_dim(flat_patches, positions, target_length=6)
        self.assertEqual(padded_patches.shape, (6, 4))
        self.assertEqual(padded_positions.shape, (6, 2))

    def test_pad_values_and_originals_preserved(self):
        flat_patches, positions = self._make()
        padded_patches, padded_positions = pad_along_first_dim(flat_patches, positions, target_length=6)
        # Original rows are untouched.
        self.assertTrue(torch.equal(padded_patches[:3], flat_patches))
        self.assertTrue(torch.equal(padded_positions[:3], positions))
        # Patch padding is the zero vector; position padding is the [-1, -1] sentinel.
        self.assertTrue(torch.equal(padded_patches[3:], torch.zeros(3, 4)))
        self.assertTrue(torch.equal(padded_positions[3:], torch.full((3, 2), -1)))

    def test_noop_when_target_equals_current_length(self):
        # Asking for exactly the current length must leave both tensors unchanged.
        flat_patches, positions = self._make()
        padded_patches, padded_positions = pad_along_first_dim(flat_patches, positions, target_length=3)
        self.assertTrue(torch.equal(padded_patches, flat_patches))
        self.assertTrue(torch.equal(padded_positions, positions))

    def test_no_truncation_when_target_smaller(self):
        # padding_length goes negative, so the padding branch is skipped: the function
        # never truncates, it only ever appends.
        flat_patches, positions = self._make()
        padded_patches, padded_positions = pad_along_first_dim(flat_patches, positions, target_length=1)
        self.assertEqual(padded_patches.shape, (3, 4))
        self.assertTrue(torch.equal(padded_patches, flat_patches))
        self.assertTrue(torch.equal(padded_positions, positions))


class TestPatchesMerge(unittest.TestCase):
    """Behavioral spec for `patches_merge`."""

    def _grid_positions(self, num_patches_width, num_patches_height):
        # Build the same (num_patches, 2) xy positions that ImageProcessor.preprocess feeds in.
        grid = torch.meshgrid(
            torch.arange(num_patches_width),
            torch.arange(num_patches_height),
            indexing="xy",
        )
        return torch.stack(grid, dim=-1).reshape(num_patches_width * num_patches_height, 2)

    def test_output_shapes(self):
        # 4x4 teacher patches (patch_size=2, D=12) merged with k=2 -> 2x2 = 4 model patches,
        # each of dimension (k*ps)**2 * 3 = 4**2 * 3 = 48 = k**2 * D.
        image = torch.arange(3 * 8 * 8).float().reshape(3, 8, 8)
        teacher_patches = convert_image_to_patches(image, patch_size=2)
        positions = self._grid_positions(4, 4)
        length = teacher_patches.shape[0] // (2 * 2)

        merged, merged_positions = patches_merge(
            teacher_patches.unsqueeze(0), positions.unsqueeze(0), length
        )
        self.assertEqual(merged.shape, (1, 4, 48))
        self.assertEqual(merged_positions.shape, (1, 4, 2))

    def test_merged_positions_are_kernel_grid(self):
        # Each merged patch covers a k x k block of teacher patches; its new position is
        # the floor-divided minimum corner, i.e. the coordinate in the coarse model grid.
        image = torch.arange(3 * 8 * 8).float().reshape(3, 8, 8)
        teacher_patches = convert_image_to_patches(image, patch_size=2)
        positions = self._grid_positions(4, 4)

        _, merged_positions = patches_merge(
            teacher_patches.unsqueeze(0), positions.unsqueeze(0), 4
        )
        # Row-major model grid in [x, y]: (0,0), (1,0), (0,1), (1,1).
        expected = torch.tensor([[[0, 0], [1, 0], [0, 1], [1, 1]]])
        self.assertTrue(torch.equal(merged_positions, expected))

    def test_merge_equivalent_to_direct_patchify(self):
        # The defining invariant: merging k x k teacher patches must reproduce, byte for
        # byte, the result of patchifying the same image directly with patch_size = k * ps.
        patch_size, k = 2, 2
        image = torch.arange(3 * 8 * 8).float().reshape(3, 8, 8)
        teacher_patches = convert_image_to_patches(image, patch_size)
        positions = self._grid_positions(8 // patch_size, 8 // patch_size)
        length = teacher_patches.shape[0] // (k * k)

        merged, _ = patches_merge(teacher_patches.unsqueeze(0), positions.unsqueeze(0), length)
        direct = convert_image_to_patches(image, k * patch_size)
        self.assertTrue(torch.equal(merged.squeeze(0), direct))

    def test_raises_on_invalid_patch_dimension(self):
        # A last dim that is not patch_size**2 * 3 cannot describe square RGB patches.
        bad_patches = torch.zeros(1, 4, 10)   # 10 != p**2 * 3 for any integer p
        positions = torch.zeros(1, 4, 2)
        with self.assertRaises(ValueError):
            patches_merge(bad_patches, positions, length=1)

    def test_raises_when_length_incompatible(self):
        # Requires L == length * k**2; length=5 against L=16 has no integer k.
        patches = torch.zeros(1, 16, 12)      # 16 patches of valid dim (p=2)
        positions = torch.zeros(1, 16, 2)
        with self.assertRaises(ValueError):
            patches_merge(patches, positions, length=5)


if __name__ == "__main__":
    unittest.main()