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()