test_pretrained_detector.py

#!/usr/bin/env python3
"""
Test script for the Pretrained Damage Detector

This script demonstrates how to use the pretrained damage detection system
with sample images or test cases.

Author: AI Assistant
Date: June 2025
"""

import os
import sys
import json
from pathlib import Path
import logging

# Add the project root to the path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from pretrained_damage_detector import PretrainedDamageDetector

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def test_single_image():
    """Test detection on a single image."""
    print("\n" + "="*60)
    print("TESTING SINGLE IMAGE DETECTION")
    print("="*60)
    
    # Initialize detector
    detector = PretrainedDamageDetector(device='auto')
    
    # Create a sample test image path (you'll need to provide an actual image)
    test_image = "test_images/sample_poster.jpg"
    
    if not os.path.exists(test_image):
        print(f"⚠️  Test image not found: {test_image}")
        print("Please create a 'test_images' directory and add some poster images for testing.")
        return
    
    # Analyze the image
    result = detector.detect_damage(test_image)
    
    # Display results
    print(f"\n📊 ANALYSIS RESULTS for {test_image}")
    print("-" * 40)
    print(f"Damaged: {'YES' if result['is_damaged'] else 'NO'}")
    print(f"Confidence: {result['confidence']:.3f}")
    print(f"Primary Damage Type: {result.get('primary_damage_type', 'none')}")
    
    print(f"\n🔍 DAMAGE SCORES:")
    for damage_type, score in result['damage_scores'].items():
        if score > 0.1:  # Only show significant scores
            print(f"  {damage_type}: {score:.3f}")
    
    print(f"\n📝 BLIP ANALYSIS:")
    print(f"Caption: {result['detailed_analysis']['blip_caption']}")
    
    print(f"\n❓ Q&A RESPONSES:")
    for question, answer in result['detailed_analysis']['qa_responses'].items():
        print(f"  Q: {question}")
        print(f"  A: {answer}")


def test_batch_processing():
    """Test batch processing of multiple images."""
    print("\n" + "="*60)
    print("TESTING BATCH PROCESSING")
    print("="*60)
    
    # Initialize detector
    detector = PretrainedDamageDetector(device='auto')
    
    # Set up test directory
    test_dir = "test_images"
    if not os.path.exists(test_dir):
        print(f"⚠️  Test directory not found: {test_dir}")
        print("Please create a 'test_images' directory and add some poster images for testing.")
        return
    
    # Get all images in test directory
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
    image_paths = [
        os.path.join(test_dir, f) 
        for f in os.listdir(test_dir)
        if Path(f).suffix.lower() in image_extensions
    ]
    
    if not image_paths:
        print(f"⚠️  No images found in {test_dir}")
        return
    
    print(f"Found {len(image_paths)} images to process")
    
    # Process batch
    results = detector.batch_analyze(image_paths, "batch_results.json")
    
    # Generate report
    report = detector.generate_report(results)
    
    # Display summary
    print(f"\n📈 BATCH ANALYSIS SUMMARY")
    print("-" * 40)
    print(f"Total Images Processed: {report['summary']['total_images']}")
    print(f"Damaged Images Found: {report['summary']['damaged_images']}")
    print(f"Damage Rate: {report['summary']['damage_rate']:.2%}")
    print(f"Average Confidence: {report['summary']['average_confidence']:.3f}")
    
    print(f"\n🏷️  DAMAGE TYPE BREAKDOWN:")
    for damage_type, count in report['damage_breakdown'].items():
        if count > 0:
            print(f"  {damage_type}: {count} images")
    
    print(f"\n💾 Detailed results saved to: batch_results.json")


def create_sample_test_structure():
    """Create sample directory structure for testing."""
    print("\n" + "="*60)
    print("CREATING SAMPLE TEST STRUCTURE")
    print("="*60)
    
    # Create test directories
    test_dir = Path("test_images")
    test_dir.mkdir(exist_ok=True)
    
    # Create subdirectories for different damage types
    damage_dirs = [
        "undamaged", "torn", "stained", "faded", 
        "creased", "water_damaged", "burned"
    ]
    
    for damage_dir in damage_dirs:
        (test_dir / damage_dir).mkdir(exist_ok=True)
    
    # Create a README file with instructions
    readme_content = """# Test Images Directory

This directory structure is for testing the Pretrained Damage Detector.

## Directory Structure:
- `undamaged/` - Place clean, undamaged poster images here
- `torn/` - Place images of posters with tears here
- `stained/` - Place images of stained posters here
- `faded/` - Place images of faded posters here
- `creased/` - Place images of creased/wrinkled posters here
- `water_damaged/` - Place images of water-damaged posters here
- `burned/` - Place images of burned posters here

## Usage:
1. Add poster images to the appropriate subdirectories
2. Run the test script: `python test_pretrained_detector.py`
3. Or test individual features using the main detector

## Supported Formats:
- JPEG (.jpg, .jpeg)
- PNG (.png)
- BMP (.bmp)
- TIFF (.tiff)
- WebP (.webp)

## Note:
The detector can analyze any poster image and will automatically
determine the damage type and severity.
"""
    
    with open(test_dir / "README.md", "w") as f:
        f.write(readme_content)
    
    print(f"✅ Created test directory structure in: {test_dir}")
    print("📝 Please add poster images to the subdirectories for testing")
    print(f"📖 See {test_dir}/README.md for detailed instructions")


def demonstrate_features():
    """Demonstrate key features of the detector without requiring images."""
    print("\n" + "="*60)
    print("PRETRAINED DAMAGE DETECTOR FEATURES")
    print("="*60)
    
    print("🤖 LOADED MODELS:")
    print("  • BLIP2 - Visual understanding and captioning")
    print("  • ViT (Vision Transformer) - Defect classification") 
    print("  • DinoV2 - Advanced anomaly detection")
    print("  • DETR - Object detection for structural analysis")
    
    print("\n🔍 DETECTION CAPABILITIES:")
    damage_types = [
        'tear', 'crease', 'stain', 'fade', 'burn', 
        'water_damage', 'scratches', 'holes', 'wrinkles'
    ]
    for i, damage_type in enumerate(damage_types, 1):
        print(f"  {i:2d}. {damage_type.replace('_', ' ').title()}")
    
    print("\n📊 ANALYSIS METHODS:")
    print("  • Visual Question Answering (VQA) using BLIP2")
    print("  • Anomaly detection with pretrained classifiers")
    print("  • Structural integrity analysis")
    print("  • Visual feature extraction (edges, texture, color)")
    print("  • Rules-based damage classification")
    
    print("\n💡 KEY ADVANTAGES:")
    print("  ✅ No training required - uses pretrained models")
    print("  ✅ Multi-modal analysis (visual + text understanding)")
    print("  ✅ Detailed damage type classification")
    print("  ✅ Confidence scoring for each prediction")
    print("  ✅ Batch processing capabilities")
    print("  ✅ Comprehensive reporting features")
    
    print("\n🚀 USAGE EXAMPLES:")
    print("  # Single image analysis:")
    print("  python pretrained_damage_detector.py image.jpg")
    print()
    print("  # Batch processing with report:")
    print("  python pretrained_damage_detector.py test_images/ --batch --report")
    print()
    print("  # Save results to JSON:")
    print("  python pretrained_damage_detector.py image.jpg --output results.json")


def main():
    """Main test function."""
    print("🎯 PRETRAINED DAMAGE DETECTOR - TEST SUITE")
    
    # Check if test images exist
    test_dir = Path("test_images")
    has_test_images = test_dir.exists() and any(test_dir.iterdir())
    
    if not has_test_images:
        print("\n⚠️  No test images found. Setting up test structure...")
        create_sample_test_structure()
        print("\n📌 Please add some poster images to the test directories and run again.")
        demonstrate_features()
        return
    
    try:
        # Demonstrate features first
        demonstrate_features()
        
        # Test single image detection
        test_single_image()
        
        # Test batch processing
        test_batch_processing()
        
        print("\n" + "="*60)
        print("✅ ALL TESTS COMPLETED SUCCESSFULLY!")
        print("="*60)
        
    except Exception as e:
        logger.error(f"Test failed with error: {e}")
        print(f"\n❌ Test failed: {e}")
        print("\nPossible solutions:")
        print("1. Install missing dependencies: pip install -r requirements.txt")
        print("2. Ensure you have internet connection for model downloads")
        print("3. Check that you have sufficient disk space (~2GB for models)")
        print("4. Try running on CPU if GPU issues: --device cpu")


if __name__ == "__main__":
    main()