batch_processing.py

import os
import json
import time
from datetime import datetime
import cv2
import numpy as np
import pandas as pd
import torch
from utils import (read_video, save_video, match_shots_with_bounces, save_to_csv)
from court_line_detector import CourtLineDetector
from TrackNet import BallTrackerNet, infer_model, remove_outliers, split_track, interpolation, write_track, interpolate_full_track, BounceDetector

# DON'T import tennis_shot_recognition here - we'll import it only when needed

def run_shot_classification_safe(video_path, model_path, left_handed=True):
    """Safely run shot classification with TensorFlow isolation"""
    
    # Force TensorFlow to CPU before import
    original_cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES', '0')
    os.environ['CUDA_VISIBLE_DEVICES'] = ''  # Hide GPU from TensorFlow
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # Reduce TF logging
    
    try:
        # Import TensorFlow-based shot classification only when needed
        from tennis_shot_recognition import run_shot_classification
        
        # Run shot classification on CPU
        result = run_shot_classification(video_path, model_path, left_handed)
        
        return result
        
    except Exception as e:
        print(f"Error in shot classification: {e}")
        return []
    
    finally:
        # Restore CUDA visibility for PyTorch
        os.environ['CUDA_VISIBLE_DEVICES'] = original_cuda_visible

def rename_output_files(video_id):
    """Rename output files to include video ID without modifying main script"""
    
    files_to_rename = [
        ("output_videos/output_video.avi", f"output_videos/output_video_{video_id}.avi"),
        ("output_videos/opponent_court_zones.jpg", f"output_videos/opponent_court_zones_{video_id}.jpg"),
        ("output_with_shots.mp4", f"output_videos/output_with_shots_{video_id}.mp4")  # Also move to output_videos
    ]
    
    for old_path, new_path in files_to_rename:
        if os.path.exists(old_path):
            try:
                # Make sure output_videos directory exists
                os.makedirs("output_videos", exist_ok=True)
                
                # Rename/move the file
                os.rename(old_path, new_path)
                print(f"✅ Renamed: {old_path} → {new_path}")
            except Exception as e:
                print(f"⚠️ Failed to rename {old_path}: {e}")
        else:
            print(f"⚠️ File not found: {old_path}")

def process_single_video(video_path, progress_info):
    """Process a single video - modified version of your main() function"""
    
    video_filename = os.path.basename(video_path)
    video_id = os.path.splitext(video_filename)[0]
    
    print(f"\n{'='*60}")
    print(f"Processing: {video_filename} (ID: {video_id})")
    print(f"Progress: {progress_info}")
    print(f"{'='*60}")
    
    start_time = time.time()
    
    try:
        # Read video
        print(f"📹 Reading video: {video_filename}")
        video_frames, fps = read_video(video_path)
        frame_height = video_frames[0].shape[0]
        print(f"✅ Video loaded: {len(video_frames)} frames, {fps} FPS")

        # === Court detection ===
        print(f"🎾 Starting court detection...")
        court_model_path = "models/keypoints_model.pth"
        court_line_detector = CourtLineDetector(court_model_path)

        all_keypoints = []
        for i, frame in enumerate(video_frames):
            if i % 50 == 0:  # Progress update every 50 frames
                print(f"   Court detection progress: {i}/{len(video_frames)} frames")
            
            # Step 1: Predict raw keypoints
            raw_keypoints = court_line_detector.predict(frame)

            # Step 2: Detect court lines
            edges = court_line_detector.get_edges(frame)
            lines = court_line_detector.get_hough_lines(edges)

            # Step 3: Snap keypoints to nearest lines
            court_keypoints = court_line_detector.snap_keypoints_to_lines(raw_keypoints, lines)

            all_keypoints.append(court_keypoints)

        print(f"✅ Court detection completed")

        # Make sure output directory exists
        os.makedirs("output_videos", exist_ok=True)

        # Divide opponent side of the court into 4 zones using predicted keypoints
        divided_img, opp_court_zones = court_line_detector.divide_opp_court(video_frames[0].copy(), court_keypoints)
        cv2.imwrite(f"output_videos/opponent_court_zones_{video_id}.jpg", divided_img)

        # === TrackNet Ball tracking ===
        print(f"🏓 Starting ball tracking...")
        print(f"PyTorch CUDA available: {torch.cuda.is_available()}")
        print(f"PyTorch GPU count: {torch.cuda.device_count()}")
        
        ball_model = BallTrackerNet()

        device = 'cuda' if torch.cuda.is_available() and torch.cuda.device_count() > 0 else 'cpu'
        print(f"Using device: {device}")
        
        ball_model.load_state_dict(torch.load('models/TrackNet.pt', map_location=device))
        ball_model = ball_model.to(device)
        ball_model.eval()
        
        # Verify model is on correct device
        print(f"Model device: {next(ball_model.parameters()).device}")
        
        ball_track, dists = infer_model(video_frames, ball_model, device)
        ball_track = remove_outliers(ball_track, dists)

        # Interpolate missing values for entire sequence
        ball_track = interpolate_full_track(ball_track) 

        extrapolate = True   
        
        if extrapolate == True:
            subtracks = split_track(ball_track)
            for r in subtracks:
                ball_subtrack = ball_track[r[0]:r[1]]
                ball_subtrack = interpolation(ball_subtrack)
                ball_track[r[0]:r[1]] = ball_subtrack

        print(f"✅ Ball tracking completed")
        
        output_frames = write_track(video_frames, ball_track)

        # Bounce Detection
        print(f"⚡ Starting bounce detection...")
        path_bounce_model = "models/ctb_regr_bounce.cbm"

        bounce_detector = BounceDetector(path_bounce_model)
        x_ball = [x[0] for x in ball_track]
        y_ball = [x[1] for x in ball_track]

        opponent_bounce_frames = bounce_detector.predict_far_side_only(x_ball, y_ball, frame_height)

        # Zone bounce detection 
        bounce_points = [(ball_track[f][0], ball_track[f][1]) for f in opponent_bounce_frames]
        valid_points, bounce_zones = court_line_detector.get_bounce_zones(bounce_points, opp_court_zones)

        # Also filter the frames to match only valid points
        opponent_bounce_frames = [
            frame for frame, point in zip(opponent_bounce_frames, bounce_points) 
            if point in valid_points
        ]

        print(f"✅ Found {len(opponent_bounce_frames)} opponent-side bounces")
        print("Filtered bounce zones:", bounce_zones)

        # Draw per-frame keypoints
        print(f"🎨 Drawing annotations...")
        for i, (frame, keypoints) in enumerate(zip(output_frames, all_keypoints)):
            output_frames[i] = court_line_detector.draw_keypoints(frame, keypoints)
            
            # Write the frame number in top left of frames
            cv2.putText(output_frames[i], f"Frame: {i}", (10,30),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        # Clean up GPU memory before TensorFlow usage
        if device == 'cuda':
            del ball_model
            torch.cuda.empty_cache()
            print(f"GPU memory after PyTorch cleanup: {torch.cuda.memory_allocated(0)/1024**2:.2f} MB")

        # Save video with video-specific name
        output_video_path = f"output_videos/output_video.avi"  # Keep original name for now
        print(f"💾 Saving video to {output_video_path}...")
        
        try:
            save_video(output_frames, output_video_path)
            print(f"✅ Video saved successfully")
            
            # Verify the file was created and can be read
            if os.path.exists(output_video_path):
                test_cap = cv2.VideoCapture(output_video_path)
                if test_cap.isOpened():
                    print("✅ Video file verification: SUCCESS")
                    test_cap.release()
                    video_path_for_classification = output_video_path
                else:
                    print("⚠️ Video file verification: FAILED - using original video")
                    video_path_for_classification = video_path
            else:
                print("⚠️ Video file not found - using original video")
                video_path_for_classification = video_path
                
        except Exception as e:
            print(f"❌ Error saving video: {e}")
            print("Using original video for shot classification...")
            video_path_for_classification = video_path
            
        # Run shot classification with TensorFlow isolation
        print(f"🎯 Running shot classification...")
        rnn_model_path = "models/tennis_rnn.h5"
        
        # Use the safe function that isolates TensorFlow
        detected_shots = run_shot_classification_safe(video_path_for_classification, rnn_model_path, left_handed=True)
        print(f"✅ Detected {len(detected_shots)} shots")

        # Rename output files to include video ID (without modifying main script)
        print(f"📁 Renaming output files with video ID...")
        rename_output_files(video_id)

        # === Match shots with bounces and save to CSV ===
        if detected_shots and opponent_bounce_frames and bounce_zones:
            matched_data = match_shots_with_bounces(detected_shots, opponent_bounce_frames, bounce_zones)
            save_to_csv(matched_data, video_id)
            print(f"✅ CSV data saved for video {video_id}")
        else:
            warnings = []
            if not detected_shots:
                warnings.append("No shots detected")
            if not opponent_bounce_frames:
                warnings.append("No opponent bounce frames detected")  
            if not bounce_zones:
                warnings.append("No bounce zones detected")
            
            print(f"⚠️ Warnings: {', '.join(warnings)}")
            print("Nothing to save to CSV for this video.")

        # Calculate processing time
        processing_time = time.time() - start_time
        print(f"✅ Video {video_id} completed in {processing_time:.1f} seconds")
        
        return True, f"Success in {processing_time:.1f}s"
        
    except Exception as e:
        processing_time = time.time() - start_time
        error_msg = f"Failed after {processing_time:.1f}s: {str(e)}"
        print(f"❌ ERROR processing {video_id}: {str(e)}")
        return False, error_msg

def load_progress():
    """Load progress from JSON file"""
    progress_file = "batch_progress.json"
    if os.path.exists(progress_file):
        with open(progress_file, 'r') as f:
            return json.load(f)
    return {"processed": [], "failed": [], "start_time": None}

def save_progress(progress_data):
    """Save progress to JSON file"""
    progress_file = "batch_progress.json"
    with open(progress_file, 'w') as f:
        json.dump(progress_data, f, indent=2)

def batch_process_videos():
    """Main batch processing function"""
    
    print("🚀 Starting Batch Video Processing")
    print(f"Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("="*60)
    
    # Load previous progress
    progress_data = load_progress()
    
    # Find all MP4 files in fyp_videos folder
    video_folder = "fyp_videos"
    if not os.path.exists(video_folder):
        print(f"❌ Error: {video_folder} folder not found!")
        return
    
    # Get all mp4 files and sort them numerically
    all_files = [f for f in os.listdir(video_folder) if f.endswith('.mp4')]
    
    # Sort by numerical order (1.mp4, 2.mp4, etc.)
    def get_video_number(filename):
        try:
            return int(os.path.splitext(filename)[0])
        except ValueError:
            return float('inf')  # Put non-numeric files at the end
    
    all_files.sort(key=get_video_number)
    
    print(f"📁 Found {len(all_files)} MP4 files in {video_folder}/")
    
    # Filter out already processed videos
    remaining_files = [f for f in all_files if f not in progress_data["processed"]]
    
    if not remaining_files:
        print("✅ All videos have been processed!")
        print(f"Processed: {len(progress_data['processed'])}")
        print(f"Failed: {len(progress_data['failed'])}")
        return
    
    print(f"📋 Videos to process: {len(remaining_files)}")
    print(f"Already processed: {len(progress_data['processed'])}")
    print(f"Previously failed: {len(progress_data['failed'])}")
    
    if progress_data["start_time"]:
        print(f"🔄 Resuming from previous session started: {progress_data['start_time']}")
    else:
        progress_data["start_time"] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    
    # Process each remaining video
    for i, filename in enumerate(remaining_files, 1):
        video_path = os.path.join(video_folder, filename)
        progress_info = f"{i}/{len(remaining_files)} (Total: {len(progress_data['processed']) + i}/{len(all_files)})"
        
        success, result_msg = process_single_video(video_path, progress_info)
        
        if success:
            progress_data["processed"].append(filename)
            print(f"✅ {filename}: {result_msg}")
        else:
            progress_data["failed"].append({"filename": filename, "error": result_msg, "timestamp": datetime.now().strftime('%Y-%m-%d %H:%M:%S')})
            print(f"❌ {filename}: {result_msg}")
        
        # Save progress after each video
        save_progress(progress_data)
        
        # Brief pause between videos to let GPU cool down
        if i < len(remaining_files):
            print("⏸️ Brief pause before next video...")
            time.sleep(2)
    
    # Final summary
    print("\n" + "="*60)
    print("🏁 BATCH PROCESSING COMPLETE")
    print("="*60)
    print(f"✅ Successfully processed: {len(progress_data['processed'])}")
    print(f"❌ Failed: {len(progress_data['failed'])}")
    
    if progress_data["failed"]:
        print("\n❌ Failed videos:")
        for fail in progress_data["failed"]:
            print(f"   - {fail['filename']}: {fail['error']}")
    
    total_time = time.time()
    print(f"\n⏱️ Session completed at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

if __name__ == "__main__":
    batch_process_videos()