test_enhanced_field_detection.py

#!/usr/bin/env python3
"""
Test script for enhanced field detection combining:
1. Advanced boundary detection (contour analysis)
2. Pitch isolation (new crowd detection method)
3. Adaptive cropping for different camera angles
"""

import cv2
import numpy as np
import os
import sys
from pathlib import Path

# Add src to path
sys.path.append(str(Path(__file__).parent / "src"))

from src.preprocessing.video_loader import VideoLoader
from src.preprocessing.filters import FrameFilters
from src.preprocessing.field_detector import FieldDetector
from src.config import VIDEOS_DIR

def test_enhanced_detection():
    """Test the enhanced field detection on sample frames."""
    
    print("=== Enhanced Field Detection Test ===\n")
    
    # Initialize components
    try:
        video_loader = VideoLoader()
        filters = FrameFilters()
        field_detector = FieldDetector()
        
        print("✓ Components initialized successfully")
    except Exception as e:
        print(f"✗ Initialization failed: {e}")
        return
    
    # Read a few frames for testing
    frame_count = 0
    max_test_frames = 10
    
    print(f"\nTesting on up to {max_test_frames} frames...\n")
    
    while frame_count < max_test_frames and video_loader.is_opened():
        ret, frame = video_loader.read()
        if not ret:
            break
            
        frame_count += 1
        print(f"--- Frame {frame_count} ---")
        
        # Test 1: Basic green ratio (original method)
        is_pitch_basic, green_ratio = FrameFilters.is_pitch_frame(frame)
        print(f"Basic Green Ratio: {green_ratio:.3f} -> {'Pitch' if is_pitch_basic else 'No Pitch'}")
        
        # Test 2: Advanced field boundary detection
        field_bbox = field_detector.detect_field_boundary(frame)
        field_metrics = field_detector.get_field_metrics(frame)
        print(f"Field Boundary Detected: {'Yes' if field_bbox else 'No'}")
        if field_bbox:
            x, y, w, h = field_bbox
            print(f"  Bounding Box: ({x}, {y}, {w}, {h})")
            print(f"  Field Coverage: {field_metrics.get('field_coverage', 0):.3f}")
        
        # Test 3: Pitch isolation (new crowd detection method)
        try:
            isolated_frame, pitch_mask = FrameFilters.isolate_pitch(frame)
            pitch_isolation_success = pitch_mask is not None
            print(f"Pitch Isolation: {'Success' if pitch_isolation_success else 'Failed'}")
        except Exception as e:
            print(f"Pitch Isolation: Failed ({e})")
            pitch_isolation_success = False
        
        # Test 4: Combined enhanced detection
        is_pitch_enhanced, enhanced_metrics = filters.is_pitch_frame_advanced(frame)
        print(f"Enhanced Detection: {'Pitch' if is_pitch_enhanced else 'No Pitch'}")
        print(f"  Methods Used: {enhanced_metrics}")
        
        # Test 5: Enhanced cropping
        try:
            cropped_frame, crop_info = filters.crop_to_field_enhanced(frame)
            print(f"Enhanced Cropping: {crop_info['method']} - {'Applied' if crop_info['crop_applied'] else 'Not Applied'}")
        except Exception as e:
            print(f"Enhanced Cropping: Failed ({e})")
        
        # Test 6: Player detection
        num_players, has_ball = FrameFilters.detect_players_and_ball(frame)
        print(f"Players Detected: {num_players}, Ball: {'Yes' if has_ball else 'No'}")
        
        # Test 7: Live play detection
        is_live_play, live_metrics = filters.is_live_play_frame_advanced(
            frame, min_players=1, use_field_detection=True
        )
        print(f"Live Play Detection: {'Yes' if is_live_play else 'No'}")
        
        print()
    
    print(f"✓ Test completed on {frame_count} frames")

def compare_methods():
    """Compare different field detection methods side by side."""
    
    print("\n=== Method Comparison ===\n")
    
    try:
        video_loader = VideoLoader()
        filters = FrameFilters()
        
        # Read one frame for detailed comparison
        ret, frame = video_loader.read()
        if not ret:
            print("✗ Could not read frame")
            return
            
        print("Analyzing single frame with all methods...\n")
        
        # Method comparison
        methods = {
            "Basic Green Ratio": lambda f: FrameFilters.is_pitch_frame(f)[0],
            "Enhanced Detection": lambda f: filters.is_pitch_frame_advanced(f)[0],
        }
        
        for method_name, method_func in methods.items():
            try:
                result = method_func(frame)
                print(f"{method_name}: {'✓ Pitch' if result else '✗ No Pitch'}")
            except Exception as e:
                print(f"{method_name}: ✗ Error ({e})")
        
        # Cropping comparison
        print("\nCropping Methods:")
        try:
            cropped_enhanced, crop_info = filters.crop_to_field_enhanced(frame)
            print(f"Enhanced Cropping: {crop_info['method']} - Success")
        except Exception as e:
            print(f"Enhanced Cropping: Failed ({e})")
            
    except Exception as e:
        print(f"✗ Comparison failed: {e}")

def save_test_results():
    """Save sample results for visual inspection."""
    
    print("\n=== Saving Test Results ===\n")
    
    try:
        video_loader = VideoLoader()
        filters = FrameFilters()
        
        output_dir = Path("test_results")
        output_dir.mkdir(exist_ok=True)
        
        frame_count = 0
        max_save_frames = 3
        
        while frame_count < max_save_frames and video_loader.is_opened():
            ret, frame = video_loader.read()
            if not ret:
                break
                
            frame_count += 1
            
            # Save original frame
            cv2.imwrite(str(output_dir / f"frame_{frame_count}_original.jpg"), frame)
            
            # Try enhanced cropping
            try:
                cropped_frame, crop_info = filters.crop_to_field_enhanced(frame)
                if crop_info['crop_applied']:
                    cv2.imwrite(str(output_dir / f"frame_{frame_count}_cropped_{crop_info['method']}.jpg"), cropped_frame)
                    print(f"✓ Saved cropped frame {frame_count} using {crop_info['method']}")
                else:
                    print(f"✗ No crop applied to frame {frame_count}")
            except Exception as e:
                print(f"✗ Failed to crop frame {frame_count}: {e}")
        
        print(f"\n✓ Test results saved to {output_dir}/")
        
    except Exception as e:
        print(f"✗ Failed to save results: {e}")

if __name__ == "__main__":
    print("Enhanced Football Field Detection Test")
    print("=====================================\n")
    
    # Check if video files are available
    if not VIDEOS_DIR.exists():
        print(f"✗ Videos directory not found: {VIDEOS_DIR}")
        print("Please add video files to test the detection system.")
        sys.exit(1)
    
    video_files = list(VIDEOS_DIR.glob("*.mp4")) + list(VIDEOS_DIR.glob("*.avi")) + list(VIDEOS_DIR.glob("*.mkv"))
    if not video_files:
        print(f"✗ No video files found in {VIDEOS_DIR}")
        sys.exit(1)
    
    print(f"Found {len(video_files)} video file(s)")
    
    # Run tests
    test_enhanced_detection()
    compare_methods()
    save_test_results()
    
    print("\n=== Test Complete ===")
    print("\nThe enhanced system now combines:")
    print("1. Contour-based field boundary detection")
    print("2. Pitch isolation with convex hull masking")
    print("3. Adaptive cropping for different camera angles")
    print("4. Fallback mechanisms for robustness")
    print("\nThis should significantly improve:")
    print("- Crowd exclusion in horizontal camera angles")
    print("- Player preservation in close-up shots")
    print("- Overall field detection accuracy")