test_model.py

"""
Test script for the CIFAR-100 ResNet model in Hugging Face Space
"""
import torch
from app import ResNet18, CIFAR100_CLASSES
import requests
from PIL import Image
from io import BytesIO

def test_model():
    """Test if the model loads and makes predictions correctly"""
    print("🧪 Testing CIFAR-100 ResNet-18 model...")
    
    # Create model
    model = ResNet18(num_classes=100)
    print(f"✅ Model created with {sum(p.numel() for p in model.parameters())} parameters")
    
    # Test with a dummy input
    dummy_input = torch.randn(1, 3, 32, 32)
    with torch.no_grad():
        output = model(dummy_input)
    
    print(f"✅ Model forward pass successful, output shape: {output.shape}")
    print(f"✅ Number of classes: {len(CIFAR100_CLASSES)}")
    print(f"✅ Sample classes: {CIFAR100_CLASSES[:10]}")
    
    # Test predictions
    probabilities = torch.softmax(output, dim=1)
    top5_prob, top5_idx = torch.topk(probabilities, 5, dim=1)
    
    print("\n🎯 Sample prediction (random input):")
    for i in range(5):
        class_idx = top5_idx[0][i].item()
        class_name = CIFAR100_CLASSES[class_idx]
        confidence = top5_prob[0][i].item()
        print(f"  {i+1}. {class_name}: {confidence:.3f}")
    
    print("\n✅ All tests passed! Model ready for deployment.")

if __name__ == "__main__":
    test_model()