Spaces:
Sleeping
Sleeping
| """ | |
| Unit tests for data preprocessing and dataset loading. | |
| """ | |
| import pytest | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| import sys | |
| sys.path.insert(0, '.') | |
| from src.data.preprocessing import ( | |
| preprocess_image, | |
| handle_channels, | |
| get_optical_transform, | |
| get_sar_transform | |
| ) | |
| from src.data.dataset import ( | |
| CrisisLandMarkDataset, | |
| create_splits | |
| ) | |
| class TestPreprocessing: | |
| """Test preprocessing functions.""" | |
| def test_optical_preprocessing(self): | |
| """Test optical RGB preprocessing.""" | |
| dummy_image = Image.fromarray( | |
| np.random.randint(0, 255, (256, 256, 3), dtype=np.uint8) | |
| ) | |
| result = preprocess_image(dummy_image, "optical", size=224) | |
| assert result.shape == (3, 224, 224) | |
| assert result.dtype == torch.float32 | |
| def test_sar_preprocessing(self): | |
| """Test SAR preprocessing.""" | |
| dummy_image = Image.fromarray( | |
| np.random.randint(0, 255, (256, 256, 2), dtype=np.uint8) | |
| ) | |
| result = preprocess_image(dummy_image, "sar", size=224) | |
| assert result.shape == (2, 224, 224) | |
| assert result.dtype == torch.float32 | |
| def test_multispectral_preprocessing(self): | |
| """Test multispectral preprocessing.""" | |
| # Create 12-channel image | |
| dummy_array = np.random.randint(0, 255, (256, 256, 12), dtype=np.uint8) | |
| dummy_image = Image.fromarray(dummy_array[..., :3]) # PIL only supports 3 channels | |
| # For 12-channel, we'd need custom handling | |
| # This test verifies the function accepts the modality | |
| result = preprocess_image(dummy_image, "optical", size=224) | |
| assert result.shape[0] == 3 # Should be 3 channels from RGB | |
| def test_channel_handling(self): | |
| """Test channel mismatch handling.""" | |
| # Create 4-channel image (should be trimmed to 3 for optical) | |
| image_4ch = np.random.randint(0, 255, (64, 64, 4), dtype=np.uint8) | |
| result = handle_channels(image_4ch, 3, "optical") | |
| assert result.shape[-1] == 3 | |
| def test_invalid_modality(self): | |
| """Test invalid modality raises error.""" | |
| dummy_image = Image.fromarray( | |
| np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8) | |
| ) | |
| with pytest.raises(ValueError): | |
| preprocess_image(dummy_image, "invalid_modality") | |
| class TestDataset: | |
| """Test dataset class.""" | |
| def test_dataset_creation(self): | |
| """Test dataset can be created.""" | |
| dataset = CrisisLandMarkDataset(modality="optical") | |
| assert len(dataset) > 0 | |
| def test_dataset_getitem(self): | |
| """Test dataset returns correct format.""" | |
| dataset = CrisisLandMarkDataset(modality="optical") | |
| image, modality_label, class_label = dataset[0] | |
| assert isinstance(image, torch.Tensor) | |
| assert isinstance(modality_label, int) | |
| assert isinstance(class_label, int) | |
| def test_modality_labels(self): | |
| """Test modality labels are correct.""" | |
| # Test optical and SAR (multispectral requires 12-channel images) | |
| for modality, expected_label in [("optical", 0), ("sar", 1)]: | |
| dataset = CrisisLandMarkDataset(modality=modality) | |
| _, modality_label, _ = dataset[0] | |
| assert modality_label == expected_label | |
| class TestSplitting: | |
| """Test data splitting.""" | |
| def test_split_no_overlap(self): | |
| """Test query/gallery split has no overlap.""" | |
| dataset = CrisisLandMarkDataset(modality="optical") | |
| query_idx, gallery_idx = create_splits(dataset, query_ratio=0.2) | |
| # Check no overlap | |
| assert len(set(query_idx) & set(gallery_idx)) == 0 | |
| def test_split_ratio(self): | |
| """Test split maintains approximate ratio.""" | |
| dataset = CrisisLandMarkDataset(modality="optical") | |
| query_idx, gallery_idx = create_splits(dataset, query_ratio=0.2) | |
| total = len(query_idx) + len(gallery_idx) | |
| query_ratio = len(query_idx) / total | |
| # Allow some tolerance | |
| assert 0.15 <= query_ratio <= 0.25 | |
| def test_split_reproducibility(self): | |
| """Test split is reproducible with same seed.""" | |
| dataset = CrisisLandMarkDataset(modality="optical") | |
| query1, gallery1 = create_splits(dataset, seed=42) | |
| query2, gallery2 = create_splits(dataset, seed=42) | |
| assert query1 == query2 | |
| assert gallery1 == gallery2 | |
| if __name__ == "__main__": | |
| pytest.main([__file__, "-v"]) |