Update app.py

app.py

@@ -1,87 +1,198 @@
-import gradio as gr
-import torch
 import cv2
 import numpy as np
-import matplotlib.pyplot as plt
-from pathlib import Path
-
-# Load the trained YOLOv5 model from a local path (directly using PyTorch)
-def load_model(model_path):
-    # Load model using torch.load, ensuring no downloading from Hugging Face
-    model = torch.load(model_path, map_location="cpu")  # Load model from the local path
-    model.eval()  # Set the model to evaluation mode
-    return model
-
-# Load the model directly from the local file
-model = load_model("best.pt")  # Ensure 'best.pt' is in the correct path
-
-# Function to process the video and calculate ball trajectory, speed, and visualize the pitch
-def process_video(video_file):
-    # Load video file using OpenCV
-    video = cv2.VideoCapture(video_file.name)
-    ball_positions = []
-    speed_data = []
-
-    frame_count = 0
-    last_position = None
-
-    while video.isOpened():
-        ret, frame = video.read()
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
+from ultralytics import YOLO
+import gradio as gr
+import os
+from scipy.interpolate import interp1d
+from scipy.optimize import curve_fit
+
+# Load YOLOv5 model
+model = YOLO("best.pt")  # Path to your best.pt
+
+# Cricket pitch dimensions (in meters)
+PITCH_LENGTH = 20.12  # Length of cricket pitch (stumps to stumps)
+PITCH_WIDTH = 3.05    # Width of pitch
+STUMP_HEIGHT = 0.71   # Stump height
+STUMP_WIDTH = 0.2286  # Stump width (including bails)
+
+# Function to process video and detect ball
+def process_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    frame_rate = cap.get(cv2.CAP_PROP_FPS)
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    positions = []
+    frame_numbers = []
+    bounce_frame = None
+    bounce_point = None
+
+    while cap.isOpened():
+        frame_num = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
+        ret, frame = cap.read()
         if not ret:
             break
-        
-        frame_count += 1
-
-        # Run the model on the frame to detect the ball
-        results = model(frame)  # Using the pre-trained YOLOv5 model
-
-        # Extract the ball position (assuming class 0 = ball)
-        ball_detections = results.pandas().xywh
-        ball = ball_detections[ball_detections['class'] == 0]  # Class 0 for ball (adjust if needed)
-        
-        if not ball.empty:
-            ball_x = ball.iloc[0]['xmin'] + (ball.iloc[0]['xmax'] - ball.iloc[0]['xmin']) / 2
-            ball_y = ball.iloc[0]['ymin'] + (ball.iloc[0]['ymax'] - ball.iloc[0]['ymin']) / 2
-            ball_positions.append((frame_count, ball_x, ball_y))  # Track position in each frame
-
-            if last_position is not None:
-                # Calculate speed based on pixel displacement between frames
-                distance = np.sqrt((ball_x - last_position[1]) ** 2 + (ball_y - last_position[2]) ** 2)
-                fps = video.get(cv2.CAP_PROP_FPS)  # Frames per second of the video
-                speed = distance * fps  # Speed = distance / time (time between frames is 1/fps)
-                speed_data.append(speed)
-
-            last_position = (frame_count, ball_x, ball_y)  # Update last position
-
-    video.release()
-
-    # Ball trajectory plot
-    plot_trajectory(ball_positions)
+
+        # Run YOLOv5 detection
+        results = model(frame)
+        detections = results[0].boxes.xywh.cpu().numpy()  # [x_center, y_center, width, height]
+
+        for det in detections:
+            x_center, y_center, _, _ = det
+            positions.append((x_center, y_center))
+            frame_numbers.append(frame_num)
+
+            # Detect bounce (lowest y_center point)
+            if bounce_frame is None or y_center > positions[bounce_frame][1]:
+                bounce_frame = len(frame_numbers) - 1
+                bounce_point = (x_center, y_center)
+
+    cap.release()
+    return positions, frame_numbers, bounce_point, frame_rate, frame_width, frame_height
+
+# Polynomial function for trajectory fitting
+def poly_func(x, a, b, c):
+    return a * x**2 + b * x + c
+
+# Predict trajectory and LBW decision
+def predict_trajectory(positions, frame_numbers, frame_width, frame_height):
+    if len(positions) < 3:
+        return None, "Insufficient detections for trajectory prediction"
+
+    x_coords = [p[0] for p in positions]
+    y_coords = [p[1] for p in positions]
+    frames = np.array(frame_numbers)
+
+    # Fit polynomial to x and y coordinates
+    try:
+        popt_x, _ = curve_fit(poly_func, frames, x_coords)
+        popt_y, _ = curve_fit(poly_func, frames, y_coords)
+    except:
+        return None, "Failed to fit trajectory"
+
+    # Extrapolate to stumps (assume stumps at y=frame_height)
+    frame_max = max(frames) + 10  # Predict 10 frames ahead
+    future_frames = np.linspace(min(frames), frame_max, 100)
+    x_pred = poly_func(future_frames, *popt_x)
+    y_pred = poly_func(future_frames, *popt_y)
+
+    # Check if trajectory hits stumps
+    stump_x = frame_width / 2  # Assume stumps at center of frame
+    stump_y = frame_height     # Assume stumps at bottom of frame
+    stump_hit = False
+    for x, y in zip(x_pred, y_pred):
+        if abs(y - stump_y) < 50 and abs(x - stump_x) < STUMP_WIDTH * frame_width / PITCH_WIDTH:
+            stump_hit = True
+            break
+
+    lbw_decision = "OUT" if stump_hit else "NOT OUT"
+    return list(zip(future_frames, x_pred, y_pred)), lbw_decision
+
+# Map pitch location
+def map_pitch(bounce_point, frame_width, frame_height):
+    if bounce_point is None:
+        return None, "No bounce detected"
+
+    x, y = bounce_point
+    # Convert pixel coordinates to pitch coordinates
+    pitch_x = (x / frame_width) * PITCH_WIDTH - PITCH_WIDTH / 2  # Center at 0
+    pitch_y = (1 - y / frame_height) * PITCH_LENGTH  # Bottom of frame = 0
+    return pitch_x, pitch_y
+
+# Estimate ball speed
+def estimate_speed(positions, frame_numbers, frame_rate, frame_width):
+    if len(positions) < 2:
+        return None, "Insufficient detections for speed estimation"
+
+    # Calculate distance in pixels between consecutive detections
+    distances = []
+    for i in range(1, len(positions)):
+        x1, y1 = positions[i-1]
+        x2, y2 = positions[i]
+        pixel_dist = np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
+        distances.append(pixel_dist)
+
+    # Convert to meters (assume pitch length = frame height)
+    pixel_to_meter = PITCH_LENGTH / frame_width
+    distances_m = [d * pixel_to_meter for d in distances]
+
+    # Speed in m/s
+    time_interval = 1 / frame_rate
+    speeds = [d / time_interval for d in distances_m]
+    avg_speed_kmh = np.mean(speeds) * 3.6  # Convert m/s to km/h
+    return avg_speed_kmh, "Speed calculated successfully"
+
+# Create pitch map visualization
+def create_pitch_map(pitch_x, pitch_y):
+    fig = go.Figure()
+    # Draw pitch rectangle
+    fig.add_shape(
+        type="rect", x0=-PITCH_WIDTH/2, y0=0, x1=PITCH_WIDTH/2, y1=PITCH_LENGTH,
+        line=dict(color="Green"), fillcolor="Green", opacity=0.3
+    )
+    # Draw stumps
+    fig.add_shape(
+        type="rect", x0=-STUMP_WIDTH/2, y0=PITCH_LENGTH-0.1, x1=STUMP_WIDTH/2, y1=PITCH_LENGTH,
+        line=dict(color="Brown"), fillcolor="Brown"
+    )
+    # Plot bounce point
+    if pitch_x is not None and pitch_y is not None:
+        fig.add_trace(go.Scatter(x=[pitch_x], y=[pitch_y], mode="markers", marker=dict(size=10, color="Red"), name="Bounce Point"))
     
-    # Return results
-    avg_speed = np.mean(speed_data) if speed_data else 0
-    return f"Average Ball Speed: {avg_speed:.2f} pixels per second"
-
-# Function to plot ball trajectory using matplotlib
-def plot_trajectory(ball_positions):
-    x_positions = [pos[1] for pos in ball_positions]
-    y_positions = [pos[2] for pos in ball_positions]
-
-    plt.figure(figsize=(10, 6))
-    plt.plot(x_positions, y_positions, label="Ball Trajectory", color='b')
-    plt.title("Ball Trajectory on Pitch")
-    plt.xlabel("X Position (pitch width)")
-    plt.ylabel("Y Position (pitch length)")
-    plt.grid(True)
-    plt.legend()
-    plt.show()
-
-# Gradio interface for the app
-iface = gr.Interface(
-    fn=process_video,  # Function to call when video is uploaded
-    inputs=gr.inputs.File(label="Upload a Video File"),  # File input (video)
-    outputs="text",  # Output the result as text
-    live=True  # Keep the interface live
-)
-
-iface.launch(debug=True)
+    fig.update_layout(
+        title="Pitch Map", xaxis_title="Width (m)", yaxis_title="Length (m)",
+        xaxis_range=[-PITCH_WIDTH/2, PITCH_WIDTH/2], yaxis_range=[0, PITCH_LENGTH]
+    )
+    return fig
+
+# Main Gradio function
+def drs_analysis(video):
+    # Save uploaded video temporarily
+    video_path = "temp_video.mp4"
+    with open(video_path, "wb") as f:
+        f.write(video.read())
+
+    # Process video
+    positions, frame_numbers, bounce_point, frame_rate, frame_width, frame_height = process_video(video_path)
+    if not positions:
+        return None, None, "No ball detected in video", None
+
+    # Predict trajectory
+    trajectory, lbw_decision = predict_trajectory(positions, frame_numbers, frame_width, frame_height)
+    if trajectory is None:
+        return None, None, lbw_decision, None
+
+    # Map pitch
+    pitch_x, pitch_y = map_pitch(bounce_point, frame_width, frame_height)
+
+    # Estimate speed
+    speed_kmh, speed_status = estimate_speed(positions, frame_numbers, frame_rate, frame_width)
+
+    # Create trajectory plot
+    trajectory_df = pd.DataFrame(trajectory, columns=["Frame", "X", "Y"])
+    fig_traj = px.line(trajectory_df, x="X", y="Y", title="Ball Trajectory (Pixel Coordinates)")
+    fig_traj.update_yaxes(autorange="reversed")  # Invert y-axis to match video frame
+
+    # Create pitch map
+    fig_pitch = create_pitch_map(pitch_x, pitch_y)
+
+    # Clean up
+    os.remove(video_path)
+
+    return fig_traj, fig_pitch, f"LBW Decision: {lbw_decision}\nSpeed: {speed_kmh:.2f} km/h", video_path
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("## Cricket DRS Analysis")
+    video_input = gr.Video(label="Upload Video Clip")
+    btn = gr.Button("Analyze")
+    trajectory_output = gr.Plot(label="Ball Trajectory")
+    pitch_output = gr.Plot(label="Pitch Map")
+    text_output = gr.Textbox(label="Analysis Results")
+    video_output = gr.Video(label="Processed Video")
+    btn.click(drs_analysis, inputs=video_input, outputs=[trajectory_output, pitch_output, text_output, video_output])
+
+if __name__ == "__main__":
+    demo.launch()