colab_enhanced_extraction.py

import os
import sys
import importlib
from pathlib import Path

# 1. Setup paths
PROJECT_ROOT = "/content/drive/MyDrive/football-tracking-data/"
if PROJECT_ROOT not in sys.path:
    sys.path.append(os.path.abspath(PROJECT_ROOT))

# 2. Force reload EVERYTHING (Updated with SoccerNet utils)
import src.config
import src.preprocessing.logger
import src.preprocessing.frame_extractor
import src.preprocessing.soccernet_utils
import src.preprocessing.filters
importlib.reload(src.config)
importlib.reload(src.preprocessing.logger)
importlib.reload(src.preprocessing.frame_extractor)
importlib.reload(src.preprocessing.soccernet_utils)
importlib.reload(src.preprocessing.filters)

from src.preprocessing.video_loader import VideoLoader
from src.preprocessing.frame_extractor import FrameExtractor
from src.preprocessing.reconstructor import VideoReconstructor
from src.preprocessing.soccernet_utils import SoccerNetReplayDetector
from src.config import FRAMES_DIR, RECONSTRUCTED_DIR, SOCCERNET_MODEL_PATH

# --- INITIALIZE THE EXPERT ---
# Load the model weights once at the start
print("Loading SoccerNet Expert...")
try:
    detector = SoccerNetReplayDetector(model_path=str(SOCCERNET_MODEL_PATH))
    print("✓ SoccerNet Expert loaded successfully")
except Exception as e:
    print(f"⚠ SoccerNet loading failed: {e}")
    print("Continuing without SoccerNet...")
    detector = None
# -----------------------------

video_folder = os.path.join(PROJECT_ROOT, "data/raw-videos/long")

# Get a video file (using index 1 as in your original script)
video_files = [f for f in os.listdir(video_folder) if f.endswith(('.mp4', '.avi', '.mov', '.mkv'))]

if not video_files:
    print("No video files found!")
    sys.exit(1)

filename = video_files[1]

# Construct paths
video_name = os.path.splitext(filename)[0]
current_video_dir = (FRAMES_DIR / video_name).resolve()
os.makedirs(str(current_video_dir), exist_ok=True)

print(f"\n" + "="*50)
print(f"PROCESSING: {filename}")
print(f"TARGET DIRECTORY: {current_video_dir}")
print("="*50)

try:
    vid_loader = VideoLoader(os.path.join(video_folder, filename))
    video_info = vid_loader.get_video_info()
    fps = video_info['fps']
    
    print(f"Video Info: {video_info}")

    # ENHANCED: Initialize FrameExtractor with advanced field detection
    extractor = FrameExtractor(
        vid_loader,
        target_fps=fps,
        target_width=vid_loader.width,
        target_height=vid_loader.height,
        use_advanced_field_detection=True,  # ← Enable enhanced field detection
        apply_field_cropping=True,          # ← Enable crowd masking
    )

    # UPDATED: Enhanced extraction with all features
    extractor.extract_frames(
        start_minute=0,
        end_minute=3,
        frame_limit=20000,
        output_dir=str(current_video_dir),
        display=False,
        filter_blurry=True,
        filter_crowd=True,
        filter_transitions=True,
        filter_brightness=True,
        filter_replays=True,
        use_soccernet=True if detector else False,  # ← Enable if detector loaded
        soccernet_detector=detector,              # ← Pass detector if available
        use_pitch_isolation=True,    # ← Enable your new crowd detection method
    )

    # Reconstruction
    output_path = RECONSTRUCTED_DIR / f"clean_{video_name}.mp4"
    reconstructor = VideoReconstructor(output_fps=fps)
    reconstructor.reconstruct_from_frames(str(current_video_dir), str(output_path))
    print(f"Successfully created: {output_path}")

except Exception as e:
    print(f"CRITICAL ERROR on {filename}: {e}")
    import traceback
    traceback.print_exc()

print("\n=== Enhanced Features Used ===")
print("✓ Advanced field boundary detection")
print("✓ Pitch isolation with convex hull masking")
print("✓ Adaptive cropping (preserves players, excludes crowds)")
print("✓ SoccerNet deep learning replay detection" if detector else "⚠ SoccerNet not available")
print("✓ Enhanced crowd exclusion for horizontal camera angles")