WIP

tests/shared/test_dataset_augmenter.py

@@ -0,0 +1,293 @@
+"""
+Tests for DatasetAugmenter.
+
+TDD Phase 1: RED - Write tests first, then implement to pass.
+"""
+
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import pytest
+from PIL import Image
+
+
+class TestDatasetAugmenter:
+    """Tests for DatasetAugmenter class."""
+
+    @pytest.fixture
+    def sample_dataset(self, tmp_path: Path) -> Path:
+        """Create a sample YOLO dataset structure."""
+        dataset_dir = tmp_path / "dataset"
+
+        # Create directory structure
+        for split in ["train", "val", "test"]:
+            (dataset_dir / "images" / split).mkdir(parents=True)
+            (dataset_dir / "labels" / split).mkdir(parents=True)
+
+        # Create sample images and labels
+        for i in range(3):
+            # Create 100x100 white image
+            img = Image.new("RGB", (100, 100), color="white")
+            img_path = dataset_dir / "images" / "train" / f"doc_{i}.png"
+            img.save(img_path)
+
+            # Create label with 2 bboxes
+            # Format: class_id x_center y_center width height
+            label_content = "0 0.5 0.3 0.2 0.1\n1 0.7 0.6 0.15 0.2\n"
+            label_path = dataset_dir / "labels" / "train" / f"doc_{i}.txt"
+            label_path.write_text(label_content)
+
+        # Create data.yaml
+        data_yaml = dataset_dir / "data.yaml"
+        data_yaml.write_text(
+            "path: .\n"
+            "train: images/train\n"
+            "val: images/val\n"
+            "test: images/test\n"
+            "nc: 10\n"
+            "names: [class0, class1, class2, class3, class4, class5, class6, class7, class8, class9]\n"
+        )
+
+        return dataset_dir
+
+    @pytest.fixture
+    def augmentation_config(self) -> dict:
+        """Create a sample augmentation config."""
+        return {
+            "gaussian_noise": {
+                "enabled": True,
+                "probability": 1.0,
+                "params": {"std": 10},
+            },
+            "gaussian_blur": {
+                "enabled": True,
+                "probability": 1.0,
+                "params": {"kernel_size": 3},
+            },
+        }
+
+    def test_augmenter_creates_additional_images(
+        self, sample_dataset: Path, augmentation_config: dict
+    ):
+        """Test that augmenter creates new augmented images."""
+        from shared.augmentation.dataset_augmenter import DatasetAugmenter
+
+        augmenter = DatasetAugmenter(augmentation_config)
+
+        # Count original images
+        original_count = len(list((sample_dataset / "images" / "train").glob("*.png")))
+        assert original_count == 3
+
+        # Apply augmentation with multiplier=2
+        result = augmenter.augment_dataset(sample_dataset, multiplier=2)
+
+        # Should now have original + 2x augmented = 3 + 6 = 9 images
+        new_count = len(list((sample_dataset / "images" / "train").glob("*.png")))
+        assert new_count == 9
+        assert result["augmented_images"] == 6
+
+    def test_augmenter_creates_matching_labels(
+        self, sample_dataset: Path, augmentation_config: dict
+    ):
+        """Test that augmenter creates label files for each augmented image."""
+        from shared.augmentation.dataset_augmenter import DatasetAugmenter
+
+        augmenter = DatasetAugmenter(augmentation_config)
+        augmenter.augment_dataset(sample_dataset, multiplier=2)
+
+        # Check that each image has a matching label file
+        images = list((sample_dataset / "images" / "train").glob("*.png"))
+        labels = list((sample_dataset / "labels" / "train").glob("*.txt"))
+
+        assert len(images) == len(labels)
+
+        # Check that augmented images have corresponding labels
+        for img_path in images:
+            label_path = sample_dataset / "labels" / "train" / f"{img_path.stem}.txt"
+            assert label_path.exists(), f"Missing label for {img_path.name}"
+
+    def test_augmented_labels_have_valid_format(
+        self, sample_dataset: Path, augmentation_config: dict
+    ):
+        """Test that augmented label files have valid YOLO format."""
+        from shared.augmentation.dataset_augmenter import DatasetAugmenter
+
+        augmenter = DatasetAugmenter(augmentation_config)
+        augmenter.augment_dataset(sample_dataset, multiplier=1)
+
+        # Check all label files
+        for label_path in (sample_dataset / "labels" / "train").glob("*.txt"):
+            content = label_path.read_text().strip()
+            if not content:
+                continue  # Empty labels are valid (background images)
+
+            for line in content.split("\n"):
+                parts = line.split()
+                assert len(parts) == 5, f"Invalid label format in {label_path.name}"
+
+                class_id = int(parts[0])
+                x_center = float(parts[1])
+                y_center = float(parts[2])
+                width = float(parts[3])
+                height = float(parts[4])
+
+                # Check values are in valid range
+                assert 0 <= class_id < 100, f"Invalid class_id: {class_id}"
+                assert 0 <= x_center <= 1, f"Invalid x_center: {x_center}"
+                assert 0 <= y_center <= 1, f"Invalid y_center: {y_center}"
+                assert 0 <= width <= 1, f"Invalid width: {width}"
+                assert 0 <= height <= 1, f"Invalid height: {height}"
+
+    def test_augmented_images_are_different(
+        self, sample_dataset: Path, augmentation_config: dict
+    ):
+        """Test that augmented images are actually different from originals."""
+        from shared.augmentation.dataset_augmenter import DatasetAugmenter
+
+        # Load original image
+        original_path = sample_dataset / "images" / "train" / "doc_0.png"
+        original_img = np.array(Image.open(original_path))
+
+        augmenter = DatasetAugmenter(augmentation_config)
+        augmenter.augment_dataset(sample_dataset, multiplier=1)
+
+        # Find augmented version
+        aug_path = sample_dataset / "images" / "train" / "doc_0_aug0.png"
+        assert aug_path.exists()
+
+        aug_img = np.array(Image.open(aug_path))
+
+        # Images should be different (due to noise/blur)
+        assert not np.array_equal(original_img, aug_img)
+
+    def test_augmented_images_same_size(
+        self, sample_dataset: Path, augmentation_config: dict
+    ):
+        """Test that augmented images have same size as originals."""
+        from shared.augmentation.dataset_augmenter import DatasetAugmenter
+
+        # Get original size
+        original_path = sample_dataset / "images" / "train" / "doc_0.png"
+        original_img = Image.open(original_path)
+        original_size = original_img.size
+
+        augmenter = DatasetAugmenter(augmentation_config)
+        augmenter.augment_dataset(sample_dataset, multiplier=1)
+
+        # Check all augmented images have same size
+        for img_path in (sample_dataset / "images" / "train").glob("*_aug*.png"):
+            img = Image.open(img_path)
+            assert img.size == original_size, f"{img_path.name} has wrong size"
+
+    def test_perspective_warp_updates_bboxes(self, sample_dataset: Path):
+        """Test that perspective_warp augmentation updates bbox coordinates."""
+        from shared.augmentation.dataset_augmenter import DatasetAugmenter
+
+        config = {
+            "perspective_warp": {
+                "enabled": True,
+                "probability": 1.0,
+                "params": {"max_warp": 0.05},  # Use larger warp for visible difference
+            },
+        }
+
+        # Read original label
+        original_label = (sample_dataset / "labels" / "train" / "doc_0.txt").read_text()
+        original_bboxes = [line.split() for line in original_label.strip().split("\n")]
+
+        augmenter = DatasetAugmenter(config)
+        augmenter.augment_dataset(sample_dataset, multiplier=1)
+
+        # Read augmented label
+        aug_label = (sample_dataset / "labels" / "train" / "doc_0_aug0.txt").read_text()
+        aug_bboxes = [line.split() for line in aug_label.strip().split("\n")]
+
+        # Same number of bboxes
+        assert len(original_bboxes) == len(aug_bboxes)
+
+        # At least one bbox should have different coordinates
+        # (perspective warp changes geometry)
+        differences_found = False
+        for orig, aug in zip(original_bboxes, aug_bboxes):
+            # Class ID should be same
+            assert orig[0] == aug[0]
+            # Coordinates might differ
+            if orig[1:] != aug[1:]:
+                differences_found = True
+
+        assert differences_found, "Perspective warp should change bbox coordinates"
+
+    def test_augmenter_only_processes_train_split(
+        self, sample_dataset: Path, augmentation_config: dict
+    ):
+        """Test that augmenter only processes train split by default."""
+        from shared.augmentation.dataset_augmenter import DatasetAugmenter
+
+        # Add a val image
+        val_img = Image.new("RGB", (100, 100), color="white")
+        val_img.save(sample_dataset / "images" / "val" / "val_doc.png")
+        (sample_dataset / "labels" / "val" / "val_doc.txt").write_text("0 0.5 0.5 0.1 0.1\n")
+
+        augmenter = DatasetAugmenter(augmentation_config)
+        augmenter.augment_dataset(sample_dataset, multiplier=2)
+
+        # Val should still have only 1 image
+        val_count = len(list((sample_dataset / "images" / "val").glob("*.png")))
+        assert val_count == 1
+
+    def test_augmenter_with_multiplier_zero_does_nothing(
+        self, sample_dataset: Path, augmentation_config: dict
+    ):
+        """Test that multiplier=0 creates no augmented images."""
+        from shared.augmentation.dataset_augmenter import DatasetAugmenter
+
+        original_count = len(list((sample_dataset / "images" / "train").glob("*.png")))
+
+        augmenter = DatasetAugmenter(augmentation_config)
+        result = augmenter.augment_dataset(sample_dataset, multiplier=0)
+
+        new_count = len(list((sample_dataset / "images" / "train").glob("*.png")))
+        assert new_count == original_count
+        assert result["augmented_images"] == 0
+
+    def test_augmenter_with_seed_is_reproducible(
+        self, sample_dataset: Path, augmentation_config: dict
+    ):
+        """Test that same seed produces same augmentation results."""
+        from shared.augmentation.dataset_augmenter import DatasetAugmenter
+
+        # Create two separate datasets
+        import shutil
+        dataset1 = sample_dataset
+        dataset2 = sample_dataset.parent / "dataset2"
+        shutil.copytree(dataset1, dataset2)
+
+        # Augment both with same seed
+        augmenter1 = DatasetAugmenter(augmentation_config, seed=42)
+        augmenter1.augment_dataset(dataset1, multiplier=1)
+
+        augmenter2 = DatasetAugmenter(augmentation_config, seed=42)
+        augmenter2.augment_dataset(dataset2, multiplier=1)
+
+        # Compare augmented images
+        aug1 = np.array(Image.open(dataset1 / "images" / "train" / "doc_0_aug0.png"))
+        aug2 = np.array(Image.open(dataset2 / "images" / "train" / "doc_0_aug0.png"))
+
+        assert np.array_equal(aug1, aug2), "Same seed should produce same augmentation"
+
+    def test_augmenter_returns_summary(
+        self, sample_dataset: Path, augmentation_config: dict
+    ):
+        """Test that augmenter returns a summary of what was done."""
+        from shared.augmentation.dataset_augmenter import DatasetAugmenter
+
+        augmenter = DatasetAugmenter(augmentation_config)
+        result = augmenter.augment_dataset(sample_dataset, multiplier=2)
+
+        assert "original_images" in result
+        assert "augmented_images" in result
+        assert "total_images" in result
+        assert result["original_images"] == 3
+        assert result["augmented_images"] == 6
+        assert result["total_images"] == 9