test_demo.py

#!/usr/bin/env python3
"""
Test script for TRIQA Hugging Face demo
Run this to test the demo locally before uploading
"""

import os
import sys

def test_demo():
    """Test the demo components."""
    print("Testing TRIQA Hugging Face Demo...")
    
    # Check if required files exist
    required_files = [
        'app.py',
        'convnext_original.py',
        'convnext_finetune.py',
        'requirements.txt',
        'package.json'
    ]
    
    missing_files = []
    for file in required_files:
        if not os.path.exists(file):
            missing_files.append(file)
    
    if missing_files:
        print(f"❌ Missing files: {missing_files}")
        return False
    
    print("✅ All required files present")
    
    # Check if model files exist
    model_files = [
        'feature_models/convnext_tiny_22k_224.pth',
        'feature_models/triqa_quality_aware.pth',
        'Regression_Models/KonIQ_scaler.save',
        'Regression_Models/KonIQ_TRIQA.save'
    ]
    
    missing_models = []
    for file in model_files:
        if not os.path.exists(file):
            missing_models.append(file)
    
    if missing_models:
        print(f"⚠️  Missing model files: {missing_models}")
        print("   Download from Box: https://utexas.box.com/s/8aw6axc2lofouja65uc726lca8b1cduf")
    else:
        print("✅ All model files present")
    
    # Check if sample images exist
    sample_dir = 'sample_image'
    if os.path.exists(sample_dir):
        sample_files = os.listdir(sample_dir)
        if sample_files:
            print(f"✅ Sample images present: {len(sample_files)} files")
        else:
            print("⚠️  No sample images found")
    else:
        print("⚠️  Sample image directory not found")
    
    print("\nDemo test complete!")
    print("If all files are present, you can upload to Hugging Face Spaces.")
    
    return len(missing_files) == 0

if __name__ == "__main__":
    test_demo()