test_oscd_fixes.py

# %% [markdown]
# # Test OSCD Fixes
# 
# This script tests the fixes made to the OSCD training pipeline.

# %%
import torch
import numpy as np
from pathlib import Path
import sys

# Add current directory to path for imports
sys.path.append('.')

from models.vision_transformer import create_oscd_model
from data_loader_oscd import OSCDDataset, create_oscd_dataloaders

# %%
def test_model_creation():
    """Test that the OSCD model can be created and run forward pass."""
    print("Testing OSCD model creation...")
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Create model
    model = create_oscd_model(
        model_size="tiny",  # Use tiny for faster testing
        use_pretrained=False
    ).to(device)
    
    # Create dummy input
    batch_size = 2
    image_pair = torch.randn(batch_size, 2, 13, 64, 64).to(device)
    
    # Forward pass
    with torch.no_grad():
        output = model(image_pair)
        print(f"Input shape: {image_pair.shape}")
        print(f"Output shape: {output.shape}")
        print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
    
    assert output.shape == (batch_size, 1), f"Expected (batch_size, 1), got {output.shape}"
    print("✓ Model creation and forward pass successful!")

# %%
def test_data_loader():
    """Test that the data loader works correctly."""
    print("\nTesting data loader...")
    
    # Check if data directory exists
    data_dir = Path("./oscd_npz")
    if not data_dir.exists():
        print("⚠ Data directory not found, skipping data loader test")
        return
    
    try:
        # Create dataset with limited samples
        dataset = OSCDDataset(
            data_dir=str(data_dir),
            max_samples_per_city=10  # Limit for testing
        )
        
        print(f"Dataset size: {len(dataset)}")
        print(f"City names: {dataset.get_city_names()}")
        
        # Test getting a sample
        if len(dataset) > 0:
            image_pair, label = dataset[0]
            print(f"Sample 0 - Image pair shape: {image_pair.shape}")
            print(f"Sample 0 - Label shape: {label.shape}")
            print(f"Sample 0 - Label type: {type(label)}")
            
            # Debug: Check the raw data structure
            print(f"Raw data shape for city 0: {dataset.X[0].shape}")
            print(f"Raw sample 0 shape: {dataset.X[0][0].shape}")
            print(f"Expected image pair shape: (2, 13, H, W)")
            
            # Test that label is properly indexed (not all labels for the city)
            assert label.shape != (dataset.X[0].shape[0],), "Label indexing bug not fixed!"
            
            # Test that we get an image pair
            if image_pair.shape[0] == 2:
                print("✓ Image pair structure correct!")
                print(f"  Before image shape: {image_pair[0].shape}")
                print(f"  After image shape: {image_pair[1].shape}")
            else:
                print(f"⚠ Expected image pair with 2 images, got {image_pair.shape[0]} images")
            
            # Test multiple samples to check consistency
            print("\nTesting multiple samples...")
            for i in range(min(3, len(dataset))):
                img_pair, lbl = dataset[i]
                print(f"  Sample {i}: image_pair={img_pair.shape}, label={lbl.shape}")
            
            print("✓ Data loader sample retrieval successful!")
        
    except Exception as e:
        print(f"⚠ Data loader test failed: {e}")
        import traceback
        traceback.print_exc()

# %%
def test_training_components():
    """Test that training components work together."""
    print("\nTesting training components...")
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Create model
    model = create_oscd_model(model_size="tiny").to(device)
    
    # Create dummy data with fixed size (model expects 64x64)
    batch_size = 2
    image_pair = torch.randn(batch_size, 2, 13, 64, 64).to(device)
    
    # Create 2D mask label to test (like OSCD data)
    label = torch.randn(64, 64).to(device)  # 2D mask
    
    # Test loss function
    criterion = torch.nn.BCEWithLogitsLoss()
    
    print(f"Testing with 2D mask label: {label.shape}")
    
    # Forward pass
    output = model(image_pair)  # Shape: (batch_size, 1)
    
    # Process label (same logic as in training)
    binary_label = (label > 0).float().mean()  # Average change percentage
    binary_label = (binary_label > 0.5).float()  # Threshold to binary
    
    # Create batch of labels to match output shape
    binary_labels = binary_label.repeat(batch_size)  # Shape: (batch_size,)
    
    # Calculate loss
    loss = criterion(output.squeeze(), binary_labels)
    
    # Get prediction
    prediction = (torch.sigmoid(output.squeeze()) > 0.5).float()
    
    print(f"  Output shape: {output.shape}")
    print(f"  Binary label: {binary_label.item():.3f}")
    print(f"  Binary labels batch: {binary_labels.shape}")
    print(f"  Loss: {loss.item():.4f}")
    print(f"  Predictions: {prediction.shape}")
    
    print("✓ Training components test successful!")

# %%
def main():
    """Run all tests."""
    print("=" * 50)
    print("OSCD FIXES TESTING")
    print("=" * 50)
    
    # Test model creation
    test_model_creation()
    
    # Test data loader
    test_data_loader()
    
    # Test training components
    test_training_components()
    
    print("\n" + "=" * 50)
    print("ALL TESTS COMPLETED!")
    print("=" * 50)
    
    print("\n" + "=" * 50)
    print("IMPORTANT NOTES:")
    print("=" * 50)
    print("1. OSCD data structure: X=(2, 13, H, W) - 2 samples per city, each with 13 bands")
    print("2. Labels: y=(H, W) - 2D change mask per city")
    print("3. Data loader creates image pairs from consecutive samples")
    print("4. Images are resized to 64x64 for ViT compatibility")
    print("5. 2D masks are converted to binary labels for change detection")
    print("=" * 50)

# %%
if __name__ == "__main__":
    main()