Add unit tests for crop_generator and ds_utils utility functions

Author: wli51

tests/datasets/ds_engine/test_crop_generator.py

@@ -0,0 +1,89 @@
+"""
+Tests for crop_generator.py utility functions.
+"""
+
+import pytest
+
+from virtual_stain_flow.datasets.base_dataset import BaseImageDataset
+from virtual_stain_flow.datasets.ds_engine.crop_generator import (
+    _compute_center_crop,
+    generate_center_crops,
+)
+
+
+class TestComputeCenterCrop:
+    """Tests for _compute_center_crop function."""
+
+    def test_computes_correct_center_for_even_dimensions(self):
+        """Should compute correct top-left for centered crop."""
+        x, y = _compute_center_crop(100, 100, 50)
+        assert (x, y) == (25, 25)
+
+    def test_computes_correct_center_for_odd_remainder(self):
+        """Should floor division for odd remainder."""
+        x, y = _compute_center_crop(10, 10, 5)
+        assert (x, y) == (2, 2)
+
+    def test_computes_correct_center_for_rectangular_image(self):
+        """Should handle non-square images."""
+        x, y = _compute_center_crop(100, 50, 20)
+        assert (x, y) == (40, 15)
+
+    def test_raises_error_when_crop_exceeds_width(self):
+        """Should raise ValueError when crop_size > width."""
+        with pytest.raises(ValueError, match="exceeds image dimensions"):
+            _compute_center_crop(10, 100, 20)
+
+    def test_raises_error_when_crop_exceeds_height(self):
+        """Should raise ValueError when crop_size > height."""
+        with pytest.raises(ValueError, match="exceeds image dimensions"):
+            _compute_center_crop(100, 10, 20)
+
+
+class TestGenerateCenterCrops:
+    """Tests for generate_center_crops function."""
+
+    def test_generates_center_crops_for_all_samples(self, basic_dataset):
+        """Should generate one center crop per sample."""
+        crop_specs = generate_center_crops(basic_dataset, crop_size=4)
+        
+        # Should have 3 samples (from file_index fixture)
+        assert len(crop_specs) == 3
+        assert set(crop_specs.keys()) == {0, 1, 2}
+
+    def test_crop_specs_have_correct_format(self, basic_dataset):
+        """Should return crop specs in ((x, y), width, height) format."""
+        crop_specs = generate_center_crops(basic_dataset, crop_size=4)
+        
+        for idx, crops in crop_specs.items():
+            assert len(crops) == 1  # One center crop per sample
+            (x, y), w, h = crops[0]
+            assert w == 4
+            assert h == 4
+            # For 10x10 images with crop_size=4: center is at (3, 3)
+            assert (x, y) == (3, 3)
+
+    def test_raises_error_for_non_positive_crop_size(self, basic_dataset):
+        """Should raise ValueError for crop_size <= 0."""
+        with pytest.raises(ValueError, match="crop_size must be positive"):
+            generate_center_crops(basic_dataset, crop_size=0)
+        
+        with pytest.raises(ValueError, match="crop_size must be positive"):
+            generate_center_crops(basic_dataset, crop_size=-1)
+
+    def test_raises_error_when_no_active_channels(self, file_index):
+        """Should raise ValueError when no channels configured."""
+        dataset = BaseImageDataset(
+            file_index=file_index,
+            pil_image_mode="I;16",
+            input_channel_keys=None,
+            target_channel_keys=None,
+        )
+        with pytest.raises(ValueError, match="No active channels"):
+            generate_center_crops(dataset, crop_size=4)
+
+    def test_raises_error_when_crop_too_large(self, basic_dataset):
+        """Should raise ValueError when crop_size exceeds image dimensions."""
+        # Images are 10x10, so crop_size=20 should fail
+        with pytest.raises(ValueError, match="exceeds image dimensions"):
+            generate_center_crops(basic_dataset, crop_size=20)

tests/datasets/ds_engine/test_ds_utils.py

@@ -0,0 +1,70 @@
+"""
+Tests for ds_utils.py utility functions.
+"""
+
+import pytest
+
+from virtual_stain_flow.datasets.ds_engine.ds_utils import (
+    _get_active_channels,
+    _validate_same_dimensions_across_channels,
+)
+
+
+class TestGetActiveChannels:
+    """Tests for _get_active_channels function."""
+
+    def test_returns_union_of_input_and_target_channels(self, basic_dataset):
+        """Should return all unique channel keys from input and target."""
+        result = _get_active_channels(basic_dataset)
+        assert result == ["input_ch1", "input_ch2", "target_ch1"]
+
+    def test_removes_duplicates_preserving_order(self, file_index):
+        """Should remove duplicates while preserving order."""
+        from virtual_stain_flow.datasets.base_dataset import BaseImageDataset
+
+        # Create dataset with overlapping channel keys
+        dataset = BaseImageDataset(
+            file_index=file_index,
+            pil_image_mode="I;16",
+            input_channel_keys=["input_ch1", "input_ch2"],
+            target_channel_keys=["input_ch1"],  # Duplicate
+        )
+        result = _get_active_channels(dataset)
+        assert result == ["input_ch1", "input_ch2"]
+
+
+class TestValidateSameDimensionsAcrossChannels:
+    """Tests for _validate_same_dimensions_across_channels function."""
+
+    def test_returns_common_dimension_when_all_match(self):
+        """Should return the common dimension when all channels match."""
+        dims = ((10, 10), (10, 10), (10, 10))
+        channels = ["ch1", "ch2", "ch3"]
+        result = _validate_same_dimensions_across_channels(dims, channels, idx=0)
+        assert result == (10, 10)
+
+    def test_raises_error_on_dimension_mismatch(self):
+        """Should raise ValueError when dimensions don't match."""
+        dims = ((10, 10), (20, 20))
+        channels = ["ch1", "ch2"]
+        with pytest.raises(ValueError, match="Dimension mismatch"):
+            _validate_same_dimensions_across_channels(dims, channels, idx=0)
+
+    def test_raises_error_on_empty_dims(self):
+        """Should raise ValueError when dims is empty."""
+        with pytest.raises(ValueError, match="No dimensions returned"):
+            _validate_same_dimensions_across_channels((), [], idx=0)
+
+    def test_raises_error_when_all_files_missing(self):
+        """Should raise ValueError when all channel files are missing (None)."""
+        dims = (None, None)
+        channels = ["ch1", "ch2"]
+        with pytest.raises(ValueError, match="All channel files missing"):
+            _validate_same_dimensions_across_channels(dims, channels, idx=0)
+
+    def test_ignores_none_values_for_missing_files(self):
+        """Should skip None values and validate remaining dimensions."""
+        dims = ((10, 10), None, (10, 10))
+        channels = ["ch1", "ch2", "ch3"]
+        result = _validate_same_dimensions_across_channels(dims, channels, idx=0)
+        assert result == (10, 10)

tests/datasets/test_crop_dataset.py

@@ -336,3 +336,28 @@ def test_dataloader_with_shuffle(self, crop_dataset):
         for inp_batch, tar_batch in loader:
             assert inp_batch.shape[0] <= 2
             assert tar_batch.shape[0] <= 2
+
+
+class TestCropImageDatasetFromBaseDataset:
+    """Test suite for CropImageDataset.from_base_dataset class method."""
+
+    def test_from_base_dataset_default_center_crops(self, basic_dataset):
+        """Test from_base_dataset creates CropImageDataset with default center crops."""
+        crop_ds = CropImageDataset.from_base_dataset(basic_dataset, crop_size=4)
+        
+        # Should have one center crop per image (3 images)
+        assert len(crop_ds) == 3
+        
+        # Should preserve channel keys
+        assert crop_ds.input_channel_keys == basic_dataset.input_channel_keys
+        assert crop_ds.target_channel_keys == basic_dataset.target_channel_keys
+        assert crop_ds.pil_image_mode == basic_dataset.pil_image_mode
+        
+        # Verify crop shape and values
+        inp_np, tar_np = crop_ds.get_raw_item(0)
+        assert inp_np.shape == (2, 4, 4)
+        assert tar_np.shape == (1, 4, 4)
+        
+        # Center crop of 10x10 image with crop_size=4 starts at (3, 3)
+        assert crop_ds.crop_info.x == 3
+        assert crop_ds.crop_info.y == 3