Update app.py

app.py

@@ -1,196 +1,201 @@
-import gradio as gr
-import torch
-from ultralytics import YOLO
 import cv2
 import numpy as np
-from PIL import Image
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
+import torch
+import gradio as gr
 import os
-import matplotlib.pyplot as plt
-from scipy.interpolate import interp1d
-
-# Load YOLOv5 model
-model = YOLO("best.pt")
-
-class CentroidTracker:
-    def __init__(self, max_disappeared=50):
-        self.next_object_id = 0
-        self.objects = {}
-        self.disappeared = {}
-        self.max_disappeared = max_disappeared
-
-    def register(self, centroid):
-        self.objects[self.next_object_id] = centroid
-        self.disappeared[self.next_object_id] = 0
-        self.next_object_id += 1
-
-    def deregister(self, object_id):
-        del self.objects[object_id]
-        del self.disappeared[object_id]
-
-    def update(self, rects):
-        if len(rects) == 0:
-            for object_id in list(self.disappeared.keys()):
-                self.disappeared[object_id] += 1
-                if self.disappeared[object_id] > self.max_disappeared:
-                    self.deregister(object_id)
-            return self.objects
-
-        input_centroids = np.zeros((len(rects), 2), dtype="int")
-        for (i, (x1, y1, x2, y2)) in enumerate(rects):
-            cX = int((x1 + x2) / 2.0)
-            cY = int((y1 + y2) / 2.0)
-            input_centroids[i] = (cX, cY)
-
-        if len(self.objects) == 0:
-            for i in range(len(input_centroids)):
-                self.register(input_centroids[i])
-        else:
-            object_ids = list(self.objects.keys())
-            object_centroids = list(self.objects.values())
-            D = np.sqrt(((input_centroids[:, None] - object_centroids) ** 2).sum(axis=2))
-            rows = D.min(axis=1).argsort()
-            cols = D.argmin(axis=1)[rows]
-            used_rows = set()
-            used_cols = set()
-            for (row, col) in zip(rows, cols):
-                if row in used_rows or col in used_cols:
-                    continue
-                object_id = object_ids[col]
-                self.objects[object_id] = input_centroids[row]
-                self.disappeared[object_id] = 0
-                used_rows.add(row)
-                used_cols.add(col)
-            unused_rows = set(range(0, D.shape[0])).difference(used_rows)
-            unused_cols = set(range(0, D.shape[1])).difference(used_cols)
-            if D.shape[0] >= D.shape[1]:
-                for row in unused_rows:
-                    self.register(input_centroids[row])
-            else:
-                for col in unused_cols:
-                    object_id = object_ids[col]
-                    self.disappeared[object_id] += 1
-                    if self.disappeared[object_id] > self.max_disappeared:
-                        self.deregister(object_id)
-        return self.objects
-
-def detect_and_track_ball(video_path, conf_threshold=0.5, iou_threshold=0.5):
-    """
-    Detect and track ball in video, generate pitch map, and predict LBW outcome.
-    
-    Args:
-        video_path: Path to uploaded video
-        conf_threshold: Confidence threshold for detection
-        iou_threshold: IoU threshold for non-max suppression
-    
-    Returns:
-        Tuple of (annotated video path, pitch map image path, LBW decision)
-    """
-    # Initialize tracker
-    tracker = CentroidTracker(max_disappeared=10)
+from scipy.optimize import curve_fit
+import sys
+
+# Add yolov5 directory to sys.path
+sys.path.append(os.path.join(os.path.dirname(__file__), "yolov5"))
+
+# Import YOLOv5 modules
+from models.experimental import attempt_load
+from utils.general import non_max_suppression, xywh2xyxy
+
+# 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)
+
+# Load model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = attempt_load("best.pt")  # Load without map_location
+model.to(device).eval()  # Move model to device and set to evaluation mode
+
+# Function to process video and detect ball
+def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
-    output_path = "output_video.mp4"
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, 30.0, 
-                         (int(cap.get(3)), int(cap.get(4))))
-    
-    # Store ball centroids for trajectory
-    centroids = []
-    pitch_points = []
-    
+    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_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        results = model.predict(frame_rgb, conf=conf_threshold, iou=iou_threshold)
-        
-        rects = []
-        for box in results[0].boxes:
-            x1, y1, x2, y2 = map(int, box.xyxy[0])
-            conf = box.conf[0]
-            label = f"Ball: {conf:.2f}"
-            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
-            rects.append((x1, y1, x2, y2))
-        
-        # Update tracker
-        objects = tracker.update(rects)
-        for object_id, centroid in objects.items():
-            cv2.circle(frame, centroid, 5, (0, 0, 255), -1)
-            centroids.append(centroid)
-        
-        out.write(frame)
-    
+
+        # Preprocess frame for YOLOv5
+        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        img = torch.from_numpy(img).to(device).float() / 255.0
+        img = img.permute(2, 0, 1).unsqueeze(0)  # [1, 3, H, W]
+
+        # Run inference
+        with torch.no_grad():
+            pred = model(img)[0]
+        pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)
+
+        # Process detections
+        for det in pred:
+            if det is not None and len(det):
+                det = xywh2xyxy(det)  # Convert to [x1, y1, x2, y2]
+                for *xyxy, conf, cls in det:
+                    x_center = (xyxy[0] + xyxy[2]) / 2
+                    y_center = (xyxy[1] + xyxy[3]) / 2
+                    positions.append((x_center.item(), y_center.item()))
+                    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.item(), y_center.item())
+
     cap.release()
-    out.release()
-    
-    # Generate pitch map
-    pitch_map_path = "pitch_map.png"
-    fig, ax = plt.subplots(figsize=(8, 4))
-    ax.set_xlim(0, 22)  # Pitch length in meters (approx)
-    ax.set_ylim(-1.5, 1.5)  # Pitch width (approx)
-    ax.set_xlabel("Length (m)")
-    ax.set_ylabel("Width (m)")
-    ax.set_title("Pitch Map with Ball Trajectory")
-    
-    # Plot stumps
-    ax.plot([20.12, 20.12], [-0.135, 0.135], 'k-', lw=5)  # Stumps at bowling end
-    ax.plot([0, 0], [-0.135, 0.135], 'k-', lw=5)  # Stumps at batting end
-    ax.plot([0, 20.12], [0, 0], 'k--')  # Pitch center line
-    
-    # Map centroids to pitch coordinates (simplified scaling)
-    if centroids:
-        x_coords = [20.12 - (c[1] / cap.get(4)) * 20.12 for c in centroids]  # Scale y to pitch length
-        y_coords = [(c[0] / cap.get(3)) * 2.7 - 1.35 for c in centroids]  # Scale x to pitch width
-        ax.plot(x_coords, y_coords, 'ro-', label="Ball Trajectory")
-        pitch_points = list(zip(x_coords, y_coords))
-    
-    ax.legend()
-    plt.savefig(pitch_map_path)
-    plt.close()
+    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
+    frame_max = max(frames) + 10
+    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
+    stump_y = frame_height
+    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
+    pitch_x = (x / frame_width) * PITCH_WIDTH - PITCH_WIDTH / 2
+    pitch_y = (1 - y / frame_height) * PITCH_LENGTH
+    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"
+
+    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)
+
+    pixel_to_meter = PITCH_LENGTH / frame_width
+    distances_m = [d * pixel_to_meter for d in distances]
+    time_interval = 1 / frame_rate
+    speeds = [d / time_interval for d in distances_m]
+    avg_speed_kmh = np.mean(speeds) * 3.6
+    return avg_speed_kmh, "Speed calculated successfully"
+
+# Create pitch map visualization
+def create_pitch_map(pitch_x, pitch_y):
+    fig = go.Figure()
+    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
+    )
+    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"
+    )
+    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"))
     
-    # LBW Decision (simplified physics-based model)
-    lbw_decision = "Not Out"
-    if pitch_points:
-        # Check pitching, impact, and wickets
-        pitching = any(0 <= x <= 20.12 and -0.135 <= y <= 0.135 for x, y in pitch_points[:len(pitch_points)//2])
-        impact = any(18 <= x <= 20.12 for x, y in pitch_points[len(pitch_points)//2:])
-        
-        # Fit a quadratic curve to predict trajectory post-impact
-        if len(x_coords) > 2:
-            t = np.linspace(0, 1, len(x_coords))
-            f_x = interp1d(t, x_coords, kind='quadratic', fill_value="extrapolate")
-            f_y = interp1d(t, y_coords, kind='quadratic', fill_value="extrapolate")
-            t_future = np.array([1.5])  # Predict beyond impact
-            x_future = f_x(t_future)[0]
-            y_future = f_y(t_future)[0]
-            wickets = (18 <= x_future <= 20.12) and (-0.135 <= y_future <= 0.135)
-            
-            if pitching and impact and wickets:
-                lbw_decision = "Out"
-            elif pitching and impact:
-                lbw_decision = "Umpire's Call"  # Marginal case
-            
-    return output_path, pitch_map_path, lbw_decision
+    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):
+    video_path = "temp_video.mp4"
+    with open(video_path, "wb") as f:
+        f.write(video.read())
+
+    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
+
+    trajectory, lbw_decision = predict_trajectory(positions, frame_numbers, frame_width, frame_height)
+    if trajectory is None:
+        return None, None, lbw_decision, None
+
+    pitch_x, pitch_y = map_pitch(bounce_point, frame_width, frame_height)
+    speed_kmh, speed_status = estimate_speed(positions, frame_numbers, frame_rate, frame_width)
+
+    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")
+
+    fig_pitch = create_pitch_map(pitch_x, pitch_y)
+
+    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("# DRS Review System for Cricket")
-    gr.Markdown("Upload a cricket video to analyze ball tracking, pitch mapping, and LBW review. Adjust thresholds for detection accuracy.")
-    
-    video_input = gr.Video(label="Upload Cricket Video")
-    conf_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5, label="Confidence Threshold")
-    iou_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5, label="IoU Threshold")
-    output_video = gr.Video(label="Annotated Video with Ball Tracking")
-    output_image = gr.Image(label="Pitch Map")
-    output_text = gr.Textbox(label="LBW Decision")
-    submit_button = gr.Button("Analyze DRS")
-    
-    submit_button.click(
-        fn=detect_and_track_ball,
-        inputs=[video_input, conf_slider, iou_slider],
-        outputs=[output_video, output_image, output_text]
-    )
+    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])
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()