app.py

from flask import Flask, render_template, request, jsonify, send_from_directory
import os
import cv2
import numpy as np
import pandas as pd
import torch
from werkzeug.utils import secure_filename
import threading
import json
from datetime import datetime

# Import existing modules without modification
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)
import gc

app = Flask(__name__)

# Configuration
UPLOAD_FOLDER = 'uploads'
OUTPUT_FOLDER = 'output_videos'
ALLOWED_EXTENSIONS = {'mp4', 'avi', 'mov'}

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 500 * 1024 * 1024  # 500MB max file size

# Ensure folders exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# Global variable to track processing status
processing_status = {
    'is_processing': False,
    'progress': 0,
    'stage': 'idle',
    'message': 'Ready',
    'video_id': None,
    'error': None
}

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

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 update_status(stage, message, progress):
    """Update global processing status"""
    global processing_status
    processing_status['stage'] = stage
    processing_status['message'] = message
    processing_status['progress'] = progress
    print(f"[{stage}] {message} ({progress}%)")

def process_video(video_path, left_handed=True):
    """Process video using existing main.py logic without modification"""
    global processing_status

    try:
        processing_status['is_processing'] = True
        processing_status['error'] = None

        # Extract video ID
        video_filename = os.path.basename(video_path)
        video_id = os.path.splitext(video_filename)[0]
        processing_status['video_id'] = video_id

        update_status('reading', 'Reading video file...', 5)

        # Read video and create video ID
        video_frames, fps = read_video(video_path)
        frame_height = video_frames[0].shape[0]

        update_status('court_detection', 'Detecting tennis court...', 15)

        # === Court detection ===
        court_model_path = "models/keypoints_model.pth"
        court_line_detector = CourtLineDetector(court_model_path)

        COURT_DETECTION_INTERVAL = 30  # Detect court every 30 frames

        all_keypoints = []
        last_valid_keypoints = None

        total_frames = len(video_frames)
        for i, frame in enumerate(video_frames):
            if i % COURT_DETECTION_INTERVAL == 0 or last_valid_keypoints is None:
                raw_keypoints = court_line_detector.predict(frame)
                edges = court_line_detector.get_edges(frame)
                lines = court_line_detector.get_hough_lines(edges)
                court_keypoints = court_line_detector.snap_keypoints_to_lines(raw_keypoints, lines)
                last_valid_keypoints = court_keypoints
            else:
                court_keypoints = last_valid_keypoints

            all_keypoints.append(court_keypoints)

            # Update progress
            if i % 100 == 0:
                progress = 15 + int((i / total_frames) * 15)
                update_status('court_detection', f'Detecting court... ({i}/{total_frames} frames)', progress)

        gc.collect()

        update_status('court_zones', 'Dividing court into zones...', 35)

        # Divide opponent side of the court into 4 zones
        divided_img, opp_court_zones = court_line_detector.divide_opp_court(video_frames[0].copy(), court_keypoints)
        cv2.imwrite("output_videos/opponent_court_zones.jpg", divided_img)

        update_status('ball_tracking', 'Tracking tennis ball...', 40)

        # === TrackNet Ball tracking ===
        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()

        with torch.no_grad():
            ball_track, dists = infer_model(video_frames, ball_model, device)

        update_status('ball_tracking', 'Processing ball trajectory...', 55)

        ball_track = remove_outliers(ball_track, dists)
        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

        output_frames = write_track(video_frames, ball_track)

        update_status('bounce_detection', 'Detecting ball bounces...', 65)

        # 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)

        opponent_bounce_frames = [
            frame for frame, point in zip(opponent_bounce_frames, bounce_points)
            if point in valid_points
        ]

        print("Filtered opponent-side bounces frames:", opponent_bounce_frames)
        print("Filtered bounce zones:", bounce_zones)

        update_status('annotating', 'Annotating video...', 70)

        # Draw per-frame keypoints
        for i, (frame, keypoints) in enumerate(zip(output_frames, all_keypoints)):
            output_frames[i] = court_line_detector.draw_keypoints(frame, keypoints)
            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()

        update_status('saving_video', 'Saving annotated video...', 75)

        # Save video
        output_video_path = "output_videos/output_video.avi"

        try:
            save_video(output_frames, output_video_path)

            if os.path.exists(output_video_path):
                test_cap = cv2.VideoCapture(output_video_path)
                if test_cap.isOpened():
                    test_cap.release()
                    video_path_for_classification = output_video_path
                else:
                    video_path_for_classification = video_path
            else:
                video_path_for_classification = video_path

        except Exception as e:
            print(f"Error saving video: {e}")
            video_path_for_classification = video_path

        update_status('shot_classification', 'Classifying tennis shots...', 85)

        # Run shot classification with TensorFlow isolation
        rnn_model_path = "models/tennis_rnn.h5"
        detected_shots = run_shot_classification_safe(video_path_for_classification, rnn_model_path, left_handed)
        print("Detected shots:", detected_shots)

        update_status('saving_results', 'Saving results...', 95)

        # === 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)
        else:
            if not detected_shots:
                print("Warning: No shots detected.")
            if not opponent_bounce_frames:
                print("Warning: No opponent bounce frames detected.")
            if not bounce_zones:
                print("Warning: No bounce zones detected.")

        update_status('complete', f'Analysis complete! Results saved for {video_id}', 100)

        processing_status['is_processing'] = False

    except Exception as e:
        processing_status['error'] = str(e)
        processing_status['is_processing'] = False
        update_status('error', f'Error: {str(e)}', 0)
        print(f"Error processing video: {e}")
        raise

@app.route('/')
def index():
    """Serve the main page"""
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    """Handle video upload and start processing"""
    global processing_status

    if processing_status['is_processing']:
        return jsonify({'error': 'A video is already being processed'}), 400

    if 'video' not in request.files:
        return jsonify({'error': 'No video file provided'}), 400

    file = request.files['video']
    left_handed = request.form.get('left_handed', 'false').lower() == 'true'

    if file.filename == '':
        return jsonify({'error': 'No file selected'}), 400

    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)

        # Start processing in a separate thread
        thread = threading.Thread(target=process_video, args=(filepath, left_handed))
        thread.start()

        return jsonify({
            'message': 'Video uploaded successfully. Processing started.',
            'filename': filename
        })

    return jsonify({'error': 'Invalid file type. Please upload MP4, AVI, or MOV'}), 400

@app.route('/status')
def get_status():
    """Get current processing status"""
    return jsonify(processing_status)

@app.route('/results/<filename>')
def download_result(filename):
    """Download result files"""
    return send_from_directory(OUTPUT_FOLDER, filename)

@app.route('/results')
def list_results():
    """List available result files"""
    try:
        files = []
        if os.path.exists(OUTPUT_FOLDER):
            for filename in os.listdir(OUTPUT_FOLDER):
                filepath = os.path.join(OUTPUT_FOLDER, filename)
                if os.path.isfile(filepath):
                    files.append({
                        'name': filename,
                        'size': os.path.getsize(filepath),
                        'modified': datetime.fromtimestamp(os.path.getmtime(filepath)).isoformat()
                    })

        # Also check for CSV results
        if os.path.exists('tennis_analysis_results.csv'):
            files.append({
                'name': 'tennis_analysis_results.csv',
                'size': os.path.getsize('tennis_analysis_results.csv'),
                'modified': datetime.fromtimestamp(os.path.getmtime('tennis_analysis_results.csv')).isoformat()
            })

        return jsonify({'files': files})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    print("="*60)
    print("\nStarting web server...")
    print("\nAccess the application at: http://localhost:5001")
    print("\nPress Ctrl+C to stop the server")
    print("="*60)

    app.run(debug=False, host='127.0.0.1', port=5001, threaded=True)