#!/usr/bin/env python3 """ Test script for data loading modules (url_dataset.py, transforms.py, data_utils.py) Run from project root: python scripts/test/test_data_loading.py """ import sys import os from pathlib import Path import json import torch import numpy as np from PIL import Image import tempfile import shutil from typing import Dict, List import traceback import time # Add project root to path project_root = Path(__file__).parent.parent.parent sys.path.insert(0, str(project_root)) from src.data.url_dataset import URLDataset, CachedDataset, BaseDataset, create_subset_dataset from src.data.transforms import ( get_transforms_for_model, get_train_transforms, get_val_transforms, MixUpTransform, CutMixTransform, DeNormalize ) from src.data.data_utils import ( create_data_loaders, analyze_dataset_splits, prepare_data_for_training, get_dataset_statistics, create_weighted_sampler ) class TestRunner: """Test runner with colored output and summary""" def __init__(self): self.passed = 0 self.failed = 0 self.skipped = 0 self.errors = [] def run_test(self, test_name: str, test_func): """Run a single test with error handling""" print(f"\n{'=' * 60}") print(f"Running: {test_name}") print(f"{'=' * 60}") try: test_func() self.passed += 1 print(f"✅ PASSED: {test_name}") except FileNotFoundError as e: self.skipped += 1 print(f"⚠️ SKIPPED: {test_name} - {str(e)}") except Exception as e: self.failed += 1 self.errors.append({ 'test': test_name, 'error': str(e), 'traceback': traceback.format_exc() }) print(f"❌ FAILED: {test_name}") print(f"Error: {str(e)}") print(f"Traceback:\n{traceback.format_exc()}") def print_summary(self): """Print test summary""" total = self.passed + self.failed + self.skipped print(f"\n{'=' * 60}") print(f"TEST SUMMARY") print(f"{'=' * 60}") print(f"Total tests: {total}") print(f"✅ Passed: {self.passed}") print(f"❌ Failed: {self.failed}") print(f"⚠️ Skipped: {self.skipped}") if self.errors: print(f"\nFailed tests:") for error in self.errors: print(f" - {error['test']}: {error['error']}") def test_dataset_initialization(): """Test basic dataset initialization""" print("Testing dataset initialization...") data_dir = Path("data") split_file = data_dir / "splits" / "train.json" if not split_file.exists(): raise FileNotFoundError(f"Split file not found: {split_file}") # Test URLDataset initialization dataset = URLDataset( split_file=str(split_file), cache_dir=str(data_dir / "cache" / "images") ) print(f"✓ URLDataset initialized with {len(dataset)} samples") print(f"✓ Number of classes: {dataset.num_classes}") print(f"✓ Class names: {dataset.decades}") # Check label mapping assert dataset.num_classes == 5, f"Expected 5 classes, got {dataset.num_classes}" assert dataset.decades == ['1960s', '1970s', '1980s', '1990s', '2000s'], "Unexpected decade labels" # Test metadata extraction if len(dataset) > 0: metadata = dataset.get_metadata(0) print(f"✓ Sample metadata keys: {list(metadata.keys())}") assert 'id' in metadata, "Missing 'id' in metadata" assert 'decade' in metadata, "Missing 'decade' in metadata" def test_transforms(): """Test transform pipelines""" print("Testing transforms...") # Create dummy image dummy_image = Image.new('RGB', (300, 300), color='red') # Test different transform types with expected sizes transforms_to_test = [ ('train_transforms', get_train_transforms(224, 'medium'), (3, 224, 224)), ('val_transforms', get_val_transforms(224), (3, 224, 224)), ('efficientnet_transforms', get_transforms_for_model('efficientnet-b2', is_training=True), (3, 260, 260)) ] for name, transform, expected_shape in transforms_to_test: tensor = transform(dummy_image) print(f"✓ {name}: output shape = {tensor.shape}") assert tensor.shape == expected_shape, f"Unexpected tensor shape for {name}: expected {expected_shape}, got {tensor.shape}" # Test augmentation levels for level in ['light', 'medium', 'heavy']: transform = get_train_transforms(224, level) tensor = transform(dummy_image) print(f"✓ Augmentation level '{level}': shape = {tensor.shape}") # Test MixUp and CutMix batch_size = 4 images = torch.randn(batch_size, 3, 224, 224) labels = torch.tensor([0, 1, 2, 3]) mixup = MixUpTransform(alpha=1.0, num_classes=5) mixed_images, labels_a, labels_b, lam = mixup(images, labels) print(f"✓ MixUp: lambda = {lam:.3f}, output shape = {mixed_images.shape}") cutmix = CutMixTransform(alpha=1.0, num_classes=5) mixed_images, labels_a, labels_b, lam = cutmix(images, labels) print(f"✓ CutMix: lambda = {lam:.3f}, output shape = {mixed_images.shape}") def test_data_loading(): """Test actual data loading from dataset""" print("Testing data loading...") data_dir = Path("data") split_file = data_dir / "splits" / "train.json" if not split_file.exists(): raise FileNotFoundError(f"Split file not found: {split_file}") # Create small subset for testing dataset = URLDataset( split_file=str(split_file), cache_dir=str(data_dir / "cache" / "images"), transform=get_val_transforms(224), fallback_on_error=True ) # Create subset to test quickly subset = create_subset_dataset(dataset, fraction=0.001, seed=42) # 0.1% of data print(f"✓ Created subset with {len(subset)} samples") # Try loading a few samples successful_loads = 0 failed_loads = 0 for i in range(min(5, len(subset))): try: image, label, metadata = subset[i] print(f"✓ Loaded sample {i}: shape={image.shape}, label={label}, decade={metadata['decade']}") successful_loads += 1 # Validate tensor assert image.shape == (3, 224, 224), f"Unexpected image shape: {image.shape}" assert 0 <= label < 5, f"Invalid label: {label}" except Exception as e: print(f"✗ Failed to load sample {i}: {str(e)}") failed_loads += 1 print(f"\nLoading summary: {successful_loads} successful, {failed_loads} failed") # Test dataset statistics if hasattr(dataset, 'get_statistics'): stats = dataset.get_statistics() print(f"✓ Dataset statistics: {stats}") def test_data_loaders(): """Test DataLoader creation""" print("Testing DataLoader creation...") # Create test config config = { 'model_name': 'efficientnet-b2', 'batch_size': 4, 'num_workers': 0, # Use 0 for testing 'data_dir': 'data', 'use_cached': False, 'use_class_weights': True, 'use_weighted_sampling': False, 'augmentation_level': 'medium' } try: # Create data loaders with small subset train_loader, val_loader, class_weights, class_names = create_data_loaders( config, use_subset=True, subset_fraction=0.001 # Very small subset for testing ) print(f"✓ Train loader: {len(train_loader)} batches") print(f"✓ Val loader: {len(val_loader)} batches") print(f"✓ Class names: {class_names}") if class_weights is not None: print(f"✓ Class weights shape: {class_weights.shape}") print(f"✓ Class weights: {class_weights.numpy()}") # Test loading one batch if len(train_loader) > 0: for batch_idx, batch_data in enumerate(train_loader): # Handle different batch formats if len(batch_data) == 2: images, labels = batch_data metadata = None elif len(batch_data) == 3: images, labels, metadata = batch_data else: images = batch_data[0] labels = batch_data[1] metadata = batch_data[2:] if len(batch_data) > 2 else None print(f"✓ Batch {batch_idx}: images={images.shape}, labels={labels.shape}") # Handle metadata safely if metadata is not None: if isinstance(metadata, (list, tuple)) and len(metadata) > 0: print(f"✓ Sample metadata: {metadata[0]}") elif isinstance(metadata, dict): print(f"✓ Sample metadata: {metadata}") else: print(f"✓ Sample metadata type: {type(metadata)}") else: print(f"✓ Sample metadata: None") # Validate batch assert images.shape[0] <= config['batch_size'], "Batch size exceeded" # Update expected shape for EfficientNet expected_shape = (3, 260, 260) # EfficientNet-B2 uses 260x260 assert images.shape[1:] == expected_shape, f"Unexpected image shape: {images.shape}, expected {expected_shape}" assert labels.shape[0] == images.shape[0], "Label count mismatch" break # Just test first batch except FileNotFoundError as e: raise FileNotFoundError(f"Required data files not found. Please run data preparation first: {e}") def test_dataset_analysis(): """Test dataset analysis functions""" print("Testing dataset analysis...") data_dir = Path("data") try: # Analyze splits analysis = analyze_dataset_splits(data_dir) for split_name, stats in analysis.items(): if isinstance(stats, dict) and 'total_images' in stats: print(f"\n{split_name} split:") print(f" ✓ Total images: {stats['total_images']}") print(f" ✓ Unique products: {stats['unique_products']}") print(f" ✓ Decades distribution: {stats.get('decades', {})}") if 'images_per_product' in stats: img_stats = stats['images_per_product'] print(f" ✓ Images per product: mean={img_stats['mean']:.2f}, std={img_stats['std']:.2f}") # Check for data leakage if 'data_leakage' in analysis: print(f"\n⚠️ WARNING: Data leakage detected!") print(f" Train-Val overlap: {analysis['data_leakage']['train_val_overlap']} products") else: print(f"\n✓ No data leakage detected between train and val splits") except FileNotFoundError as e: raise FileNotFoundError(f"Split files not found. Please create splits first: {e}") def test_weighted_sampler(): """Test weighted sampler creation""" print("Testing weighted sampler...") # Create mock dataset with imbalanced classes class MockDataset: def __init__(self): # Create imbalanced labels (more samples for class 0) self.data = [{'decade': '1960s'} for _ in range(100)] self.data += [{'decade': '1970s'} for _ in range(50)] self.data += [{'decade': '1980s'} for _ in range(30)] self.data += [{'decade': '1990s'} for _ in range(20)] self.data += [{'decade': '2000s'} for _ in range(10)] self.decades = ['1960s', '1970s', '1980s', '1990s', '2000s'] self.label_to_idx = {d: i for i, d in enumerate(self.decades)} self.num_classes = 5 def get_labels(self): return [self.label_to_idx[item['decade']] for item in self.data] dataset = MockDataset() sampler = create_weighted_sampler(dataset) print(f"✓ Created weighted sampler with {len(sampler)} samples") # Sample and check distribution sampled_indices = list(sampler)[:1000] sampled_labels = [dataset.get_labels()[idx] for idx in sampled_indices] from collections import Counter label_counts = Counter(sampled_labels) print(f"✓ Sampled label distribution: {dict(label_counts)}") # Check if sampling is more balanced counts = list(label_counts.values()) if len(counts) > 0: balance_ratio = max(counts) / min(counts) print(f"✓ Balance ratio after sampling: {balance_ratio:.2f} (lower is better)") def test_cached_dataset(): """Test CachedDataset functionality""" print("Testing CachedDataset...") data_dir = Path("data") split_file = data_dir / "splits" / "val.json" cache_dir = data_dir / "cache" / "images" if not split_file.exists(): raise FileNotFoundError(f"Split file not found: {split_file}") try: dataset = CachedDataset( split_file=str(split_file), images_dir=str(cache_dir), transform=get_val_transforms(224), verify_images=True ) print(f"✓ CachedDataset initialized with {len(dataset)} valid cached images") # Try loading a sample if available if len(dataset) > 0: image, label, metadata = dataset[0] print(f"✓ Loaded cached sample: shape={image.shape}, label={label}") else: print("⚠️ No cached images available for testing") except Exception as e: print(f"⚠️ CachedDataset test partially failed: {str(e)}") print(" This is expected if images haven't been downloaded yet") def test_data_preparation(): """Test data preparation checker""" print("Testing data preparation...") config = { 'data_dir': 'data', 'use_cached': False } # Check if data is ready (don't download) is_ready = prepare_data_for_training( config, download_if_missing=False, verify_splits=True ) if is_ready: print("✓ Data is ready for training") else: print("⚠️ Data is not fully prepared") print(" Run the data preparation pipeline to download images") def test_model_specific_transforms(): """Test transforms for different models""" print("Testing model-specific transforms...") dummy_image = Image.new('RGB', (400, 400), color='blue') models_to_test = [ ('resnet50', 224), ('efficientnet-b2', 260), # Assuming this is configured ('convnext-tiny-384', 384) # Assuming this is configured ] for model_name, expected_size in models_to_test: try: transform = get_transforms_for_model(model_name, is_training=False) tensor = transform(dummy_image) print(f"✓ {model_name}: output shape = {tensor.shape}") # Note: The actual size might be 224 if not configured in model_configs if tensor.shape[1] != expected_size: print(f" ⚠️ Expected size {expected_size}, got {tensor.shape[1]}") print(f" This might be due to default config being used") except KeyError: print(f"⚠️ Model config not found for {model_name}, using default") def main(): """Run all tests""" print("=" * 60) print("VINTAGE PRODUCT CLASSIFICATION - DATA LOADING TEST SUITE") print("=" * 60) print(f"Project root: {project_root}") print(f"Current directory: {os.getcwd()}") # Change to project root for consistent paths os.chdir(project_root) print(f"Changed to project root: {os.getcwd()}") runner = TestRunner() # Define all tests tests = [ ("Dataset Initialization", test_dataset_initialization), ("Transform Pipelines", test_transforms), ("Data Loading", test_data_loading), ("DataLoader Creation", test_data_loaders), ("Dataset Analysis", test_dataset_analysis), ("Weighted Sampler", test_weighted_sampler), ("Cached Dataset", test_cached_dataset), ("Data Preparation Check", test_data_preparation), ("Model-Specific Transforms", test_model_specific_transforms), ] # Run all tests for test_name, test_func in tests: runner.run_test(test_name, test_func) # Print summary runner.print_summary() # Return exit code return 0 if runner.failed == 0 else 1 if __name__ == "__main__": exit_code = main() sys.exit(exit_code)