test_pipeline.py

"""
Simple test script to verify Transformers pipeline integration
"""

from transformers import pipeline
import torch
from PIL import Image
import numpy as np

def test_transformers_pipeline():
    """Test if transformers depth estimation pipeline works"""
    print("🧪 Testing Transformers Depth Estimation Pipeline")
    print("=" * 50)
    
    try:
        # Initialize pipeline
        print("1. Initializing pipeline...")
        pipe = pipeline(
            "depth-estimation", 
            model="apple/DepthPro",
            device=-1,  # CPU
            torch_dtype=torch.float32
        )
        print("✅ Pipeline initialized successfully!")
        
        # Create test image
        print("2. Creating test image...")
        test_image = Image.new('RGB', (640, 480), color='blue')
        # Add some pattern
        pixels = np.array(test_image)
        for y in range(480):
            for x in range(640):
                intensity = int(255 * (1 - y / 480))
                pixels[y, x] = [intensity, intensity//2, intensity//3]
        test_image = Image.fromarray(pixels.astype(np.uint8))
        print("✅ Test image created!")
        
        # Test pipeline
        print("3. Running depth estimation...")
        result = pipe(test_image)
        print("✅ Pipeline executed successfully!")
        
        # Check result
        print("4. Checking results...")
        if isinstance(result, dict):
            if 'depth' in result:
                depth = result['depth']
                print(f"   Depth type: {type(depth)}")
                if hasattr(depth, 'shape'):
                    print(f"   Depth shape: {depth.shape}")
                elif hasattr(depth, 'size'):
                    print(f"   Depth size: {depth.size}")
                print("✅ Valid depth result obtained!")
                return True
            else:
                print(f"   Result keys: {result.keys()}")
                print("⚠️  No 'depth' key in result")
                return False
        else:
            print(f"   Result type: {type(result)}")
            if hasattr(result, 'depth'):
                print("✅ Result has depth attribute!")
                return True
            else:
                print("⚠️  Result format unexpected")
                return False
                
    except ImportError as e:
        print(f"❌ Import error: {e}")
        print("💡 Try: pip install transformers torch")
        return False
    except Exception as e:
        print(f"❌ Pipeline test failed: {e}")
        print("💡 This is expected if the model isn't available or if there are compatibility issues")
        return False

def test_fallback_dummy():
    """Test the dummy pipeline fallback"""
    print("\n🧪 Testing Dummy Pipeline Fallback")
    print("=" * 40)
    
    try:
        # Import dummy pipeline from our app
        import sys
        import os
        sys.path.append(os.path.dirname(os.path.abspath(__file__)))
        
        from app import DummyDepthPipeline
        
        dummy = DummyDepthPipeline()
        test_image = Image.new('RGB', (512, 384), color='green')
        
        result = dummy(test_image)
        
        if isinstance(result, dict) and 'depth' in result:
            depth = result['depth']
            print(f"✅ Dummy pipeline works! Depth shape: {depth.shape}")
            print(f"   Depth range: {np.min(depth):.2f} - {np.max(depth):.2f}")
            return True
        else:
            print(f"❌ Unexpected result format: {type(result)}")
            return False
            
    except Exception as e:
        print(f"❌ Dummy pipeline test failed: {e}")
        return False

if __name__ == "__main__":
    print("🚀 Testing Depth Pro Transformers Integration\n")
    
    # Test real pipeline
    pipeline_works = test_transformers_pipeline()
    
    # Test fallback
    dummy_works = test_fallback_dummy()
    
    print("\n" + "="*50)
    print("🏁 Test Summary:")
    print("="*50)
    print(f"Transformers Pipeline: {'✅ WORKS' if pipeline_works else '❌ FAILED (expected in some environments)'}")
    print(f"Dummy Pipeline Fallback: {'✅ WORKS' if dummy_works else '❌ FAILED'}")
    
    if dummy_works:
        print("\n🎉 The app should work with fallback even if the real model fails!")
    else:
        print("\n⚠️  There may be issues with the fallback implementation.")
    
    if pipeline_works:
        print("🌟 Real Depth Pro model should work perfectly!")
    else:
        print("💡 Real model may need specific environment setup or GPU.")