batch.py

import os
import sys
# Add parent directory to sys.path to enable package-level imports
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import glob
import sqlite3
import subprocess
import json
import argparse
import datetime
import traceback
import football_analytics.config as config

def parse_args():
    parser = argparse.ArgumentParser(description="Football Analytics Batch Processor")
    parser.add_argument("--input_dir", type=str, required=True, help="Directory containing raw football videos")
    parser.add_argument("--output_dir", type=str, default=None, help="Directory for exports and DB")
    parser.add_argument("--skip_frames", type=int, default=2, help="Skip frames ratio (default 2 for batch speed)")
    parser.add_argument("--no_video", action="store_true", help="Skip writing output videos to save space and processing time")
    parser.add_argument("--force", action="store_true", help="Force reprocessing of already processed videos")
    return parser.parse_args()

def init_master_db(db_path):
    """
    Initializes the master tracking and statistics database.
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # 1. Processed videos tracking table
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS processed_videos (
            video_path TEXT PRIMARY KEY,
            video_name TEXT,
            processed_at TIMESTAMP,
            status TEXT,
            error_message TEXT
        )
    """)
    
    # 2. Aggregated Match Summaries table
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS batch_game_summaries (
            video_name TEXT PRIMARY KEY,
            possession_team_A REAL,
            possession_team_B REAL,
            total_distance_team_A_meters REAL,
            total_distance_team_B_meters REAL,
            passes_attempted_team_A INTEGER,
            passes_completed_team_A INTEGER,
            passes_attempted_team_B INTEGER,
            passes_completed_team_B INTEGER
        )
    """)

    # 3. Aggregated Player Summaries table
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS batch_player_summaries (
            video_name TEXT,
            player_id INTEGER,
            team_id INTEGER,
            distance_covered_meters REAL,
            possession_time_seconds REAL,
            successful_passes INTEGER,
            total_passes INTEGER,
            pass_accuracy REAL,
            average_x REAL,
            average_y REAL,
            PRIMARY KEY (video_name, player_id)
        )
    """)
    
    conn.commit()
    conn.close()

def log_video_status(db_path, video_path, status, error_message=None):
    """
    Updates the status of a video in the tracking table.
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    video_name = os.path.basename(video_path)
    now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    cursor.execute("""
        INSERT OR REPLACE INTO processed_videos (video_path, video_name, processed_at, status, error_message)
        VALUES (?, ?, ?, ?, ?)
    """, (video_path, video_name, now, status, error_message))
    
    conn.commit()
    conn.close()

def check_video_processed(db_path, video_path):
    """
    Returns True if the video was already successfully processed.
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute("SELECT status FROM processed_videos WHERE video_path = ?", (video_path,))
    row = cursor.fetchone()
    conn.close()
    
    if row and row[0] == "success":
        return True
    return False

def aggregate_stats_to_db(db_path, video_path, summary_json_path):
    """
    Extracts analytics from a video's summary.json and saves them in the master database.
    """
    if not os.path.exists(summary_json_path):
        print(f"[Batch] Warning: Summary JSON file not found at {summary_json_path}. Cannot aggregate stats.")
        return
        
    with open(summary_json_path, 'r') as f:
        summary_data = json.load(f)
        
    video_name = os.path.basename(video_path)
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # 1. Insert game summary
    ms = summary_data["match_statistics"]
    cursor.execute("""
        INSERT OR REPLACE INTO batch_game_summaries (
            video_name, possession_team_A, possession_team_B, 
            total_distance_team_A_meters, total_distance_team_B_meters,
            passes_attempted_team_A, passes_completed_team_A,
            passes_attempted_team_B, passes_completed_team_B
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
    """, (
        video_name,
        ms.get("possession_percent_team_A", 50.0),
        ms.get("possession_percent_team_B", 50.0),
        ms.get("total_distance_team_A_meters", 0.0),
        ms.get("total_distance_team_B_meters", 0.0),
        ms.get("passes_attempted_team_A", 0),
        ms.get("passes_completed_team_A", 0),
        ms.get("passes_attempted_team_B", 0),
        ms.get("passes_completed_team_B", 0)
    ))
    
    # 2. Insert player summaries
    for p in summary_data.get("players", []):
        cursor.execute("""
            INSERT OR REPLACE INTO batch_player_summaries (
                video_name, player_id, team_id, distance_covered_meters,
                possession_time_seconds, successful_passes, total_passes,
                pass_accuracy, average_x, average_y
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """, (
            video_name,
            p.get("player_id"),
            p.get("team_id"),
            p.get("distance_covered_meters", 0.0),
            p.get("possession_time_seconds", 0.0),
            p.get("successful_passes", 0),
            p.get("total_passes", 0),
            p.get("pass_accuracy", 0.0),
            p.get("average_x", 0.0),
            p.get("average_y", 0.0)
        ))
        
    conn.commit()
    conn.close()
    print(f"[Batch] Aggregated stats for {video_name} into master database.")

def main():
    args = parse_args()
    
    if not os.path.exists(args.input_dir):
        print(f"Error: Input directory '{args.input_dir}' does not exist.")
        return
        
    output_dir = args.output_dir or config.OUTPUT_DIR
    os.makedirs(output_dir, exist_ok=True)
    
    master_db_path = os.path.join(output_dir, "batch_analytics.sqlite")
    init_master_db(master_db_path)
    
    # Supported video formats
    video_extensions = ["*.mp4", "*.avi", "*.mkv", "*.mov", "*.MP4", "*.MOV"]
    video_files = []
    for ext in video_extensions:
        video_files.extend(glob.glob(os.path.join(args.input_dir, ext)))
        
    # Remove duplicates if any
    video_files = list(set(video_files))
    
    print("\n" + "="*60)
    print(f"FOOTBALL ANALYTICS BATCH PROCESSOR")
    print(f"Input Directory: {args.input_dir}")
    print(f"Output Directory: {output_dir}")
    print(f"Master Database: {master_db_path}")
    print(f"Videos Found: {len(video_files)}")
    print("="*60 + "\n")
    
    if not video_files:
        print("No video files found matching typical extensions.")
        return
        
    for i, video_path in enumerate(video_files):
        video_name = os.path.basename(video_path)
        video_name_no_ext = os.path.splitext(video_name)[0]
        
        print(f"\n[{i+1}/{len(video_files)}] Processing: {video_name}")
        
        # Check tracking to skip if already done
        if not args.force and check_video_processed(master_db_path, video_path):
            print(f"  -> Skipping. Already successfully processed. (Use --force to reprocess)")
            continue
            
        # Create output subdirectory for this specific video
        video_out_dir = os.path.join(output_dir, video_name_no_ext)
        os.makedirs(video_out_dir, exist_ok=True)
        
        # Mark running
        log_video_status(master_db_path, video_path, "running")
        
        # Build command line for subprocess
        # Running via subprocess guarantees GPU memory is cleared completely between video runs
        cmd = [
            "python3", "pipeline.py",
            "--video", video_path,
            "--output_dir", video_out_dir,
            "--skip_frames", str(args.skip_frames)
        ]
        if args.no_video:
            cmd.append("--no_video")
            
        print(f"  -> Launching subprocess pipeline...")
        print(f"  -> Command: {' '.join(cmd)}")
        
        try:
            # Run the pipeline
            result = subprocess.run(
                cmd, 
                cwd=config.BASE_DIR, 
                check=True, 
                stdout=subprocess.PIPE, 
                stderr=subprocess.PIPE,
                text=True
            )
            
            # Print last few lines of output
            print("  -> Pipeline finished successfully.")
            
            # Log success
            log_video_status(master_db_path, video_path, "success")
            
            # Aggregate stats to master DB
            summary_path = os.path.join(video_out_dir, "summary.json")
            aggregate_stats_to_db(master_db_path, video_path, summary_path)
            
        except subprocess.CalledProcessError as e:
            print(f"  -> ERROR: Pipeline subprocess failed for {video_name}.")
            print("--- Subprocess Error Stdout ---")
            print(e.stdout[-1000:] if e.stdout else "No stdout output")
            print("--- Subprocess Error Stderr ---")
            print(e.stderr[-1000:] if e.stderr else "No stderr output")
            
            log_video_status(master_db_path, video_path, "failed", error_message=e.stderr)
            
        except Exception as e:
            print(f"  -> ERROR: Unexpected failure for {video_name}: {e}")
            log_video_status(master_db_path, video_path, "failed", error_message=str(e))
            
    print("\n" + "="*60)
    print("BATCH PROCESSING TASK COMPLETE!")
    print(f"Check results database at: {master_db_path}")
    print("="*60 + "\n")

if __name__ == "__main__":
    main()