Update app.py

app.py

@@ -3,14 +3,15 @@ import numpy as np
 import pandas as pd
 import plotly.express as px
 import plotly.graph_objects as go
-from ultralytics import YOLO
+import torch
 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
+# Load YOLOv5 model from yolov5 repository
+from yolov5.models.experimental import attempt_load
+from yolov5.utils.general import non_max_suppression, xywh2xyxy
 
 # Cricket pitch dimensions (in meters)
 PITCH_LENGTH = 20.12  # Length of cricket pitch (stumps to stumps)
@@ -18,6 +19,11 @@ 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", map_location=device)
+model.eval()
+
 # Function to process video and detect ball
 def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
@@ -35,19 +41,30 @@ def process_video(video_path):
         if not ret:
             break
 
-        # 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)
+        # 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()
     return positions, frame_numbers, bounce_point, frame_rate, frame_width, frame_height
@@ -72,15 +89,15 @@ def predict_trajectory(positions, frame_numbers, frame_width, frame_height):
     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
+    # 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  # Assume stumps at center of frame
-    stump_y = frame_height     # Assume stumps at bottom of frame
+    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:
@@ -96,9 +113,8 @@ def map_pitch(bounce_point, frame_width, frame_height):
         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
+    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
@@ -106,7 +122,6 @@ 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]
@@ -114,30 +129,24 @@ def estimate_speed(positions, frame_numbers, frame_rate, frame_width):
         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
+    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()
-    # 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"))
     
@@ -149,36 +158,27 @@ def create_pitch_map(pitch_x, pitch_y):
 
 # 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
+    fig_traj.update_yaxes(autorange="reversed")
 
-    # 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