viral_velocity/test_gemini_enhancement.py

#!/usr/bin/env python3
"""
Test script for Gemini 2.0 Flash Preview Image Enhancement
Tests the new Gemini-based image enhancement functionality
"""

import os
import sys
import logging
from image_enhancer import ImageEnhancer
from PIL import Image
import io

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def create_test_image():
    """Create a simple test image for testing"""
    # Create a simple test image
    img = Image.new('RGB', (200, 200), color='lightblue')
    
    # Save to file
    test_image_path = "test_enhancement_image.jpg"
    img.save(test_image_path, 'JPEG')
    
    return test_image_path

def test_gemini_enhancer():
    """Test the Gemini image enhancer"""
    print("🧪 Testing Gemini Image Enhancement")
    print("=" * 50)
    
    # Check Gemini API key
    api_key = os.getenv('GEMINI_API_KEY')
    if not api_key:
        print("❌ GEMINI_API_KEY not found in environment variables")
        print("   To enable Gemini enhancement:")
        print("   1. Get API key from: https://makersuite.google.com/app/apikey")
        print("   2. Set GEMINI_API_KEY environment variable")
        return False
    
    print(f"✅ Gemini API key found")
    
    # Initialize enhancer
    try:
        print("1. Initializing Gemini Image Enhancer...")
        enhancer = ImageEnhancer()
        print("   ✅ Enhancer initialized successfully")
    except Exception as e:
        print(f"   ❌ Failed to initialize enhancer: {e}")
        return False
    
    # Create test image
    print("\n2. Creating test image...")
    test_image_path = create_test_image()
    print(f"   ✅ Test image created: {test_image_path}")
    
    # Test enhancement
    print("\n3. Testing image enhancement...")
    try:
        # Test with user preferences
        user_preferences = {
            "aesthetic": "Minimalist",
            "niche": "Business Professional",
            "target_audience": "Professionals"
        }
        
        print("   📸 Starting enhancement with user preferences...")
        result = enhancer.enhance_image(test_image_path, user_preferences)
        
        if result['status'] == 'success':
            print(f"   ✅ Enhancement completed successfully!")
            print(f"   📊 Generated {result['total_generated']} enhanced images")
            
            # Show enhancement details
            for i, enhanced_img in enumerate(result['enhanced_images'], 1):
                print(f"     Version {i}: {enhanced_img['image_path']}")
                print(f"       Method: {enhanced_img.get('generation_method', 'unknown')}")
                print(f"       Prompt: {enhanced_img['prompt'][:100]}...")
            
            # Show original analysis
            if 'original_image' in result and 'analysis' in result['original_image']:
                analysis = result['original_image']['analysis']
                print(f"\n   🔍 Original Image Analysis:")
                print(f"     Issues found: {analysis.get('issues_found', [])}")
                print(f"     Quality: {analysis.get('overall_quality', 'unknown')}")
                print(f"     Main subject: {analysis.get('main_subject', 'unknown')}")
            
            return True
        else:
            print(f"   ❌ Enhancement failed: {result.get('error', 'Unknown error')}")
            return False
            
    except Exception as e:
        print(f"   ❌ Enhancement test failed: {e}")
        import traceback
        traceback.print_exc()
        return False
    
    finally:
        # Clean up test image
        if os.path.exists(test_image_path):
            os.remove(test_image_path)
            print(f"\n   🧹 Cleaned up test image: {test_image_path}")

def test_enhancement_prompts():
    """Test the enhancement prompt generation"""
    print("\n4. Testing Enhancement Prompt Generation...")
    
    try:
        enhancer = ImageEnhancer()
        
        # Test analysis data
        test_analysis = {
            "issues_found": ["blurry areas", "poor lighting", "closed eyes"],
            "enhancement_priorities": ["fix blur", "improve lighting", "open eyes"],
            "overall_quality": "poor",
            "main_subject": "person",
            "background": "office",
            "lighting": "dim"
        }
        
        # Test user preferences
        user_preferences = {
            "aesthetic": "Y2K",
            "niche": "Fashion Influencer",
            "target_audience": "Gen Z"
        }
        
        prompts = enhancer._generate_enhancement_prompts(test_analysis, user_preferences)
        
        print(f"   ✅ Generated {len(prompts)} enhancement prompts")
        
        for i, prompt in enumerate(prompts, 1):
            print(f"\n   Prompt {i}:")
            print(f"     {prompt[:150]}...")
            
            # Check for key requirements
            if "Do NOT alter personal appearance" in prompt:
                print(f"     ✅ Contains personal appearance protection")
            if "fix" in prompt.lower() or "improve" in prompt.lower():
                print(f"     ✅ Focuses on fixing imperfections")
        
        return True
        
    except Exception as e:
        print(f"   ❌ Prompt generation test failed: {e}")
        return False

def main():
    """Main test function"""
    print("🚀 Viral Velocity - Gemini Image Enhancement Test")
    print("=" * 60)
    
    # Test enhancement functionality
    enhancement_ok = test_gemini_enhancer()
    
    # Test prompt generation
    prompt_ok = test_enhancement_prompts()
    
    print("\n" + "=" * 60)
    if enhancement_ok and prompt_ok:
        print("✅ All tests completed successfully!")
        print("🎉 Gemini 2.0 Flash Preview Image Enhancement is working correctly")
        print("📝 Key features verified:")
        print("   - Gemini API integration")
        print("   - Image analysis for enhancement opportunities")
        print("   - 5 different enhancement prompts")
        print("   - Personal appearance protection")
        print("   - Focus on fixing imperfections (blur, lighting, objects)")
        print("   - User preference integration")
    else:
        print("❌ Some tests failed - check the logs above")
        print("📝 To fix issues:")
        print("   1. Ensure GEMINI_API_KEY is set correctly")
        print("   2. Check internet connection")
        print("   3. Verify Gemini API access")

if __name__ == "__main__":
    main()