Update app.py

app.py

@@ -31,7 +31,7 @@ BATCH_SIZE = 4  # Process 4 frames at a time
 
 # Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = attempt_load("yolov5s.pt")  # Load without map_location
+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
@@ -114,10 +114,10 @@ def process_video(video_path):
 def poly_func(x, a, b, c):
     return a * x**2 + b * x + c
 
-# Predict trajectory and LBW decision
+# Predict trajectory and wicket inline path
 def predict_trajectory(positions, frame_numbers, frame_width, frame_height):
     if len(positions) < 3:
-        return None, "Insufficient detections for trajectory prediction"
+        return None, None, "Insufficient detections for trajectory prediction"
 
     x_coords = [p[0] for p in positions]
     y_coords = [p[1] for p in positions]
@@ -128,7 +128,7 @@ def predict_trajectory(positions, frame_numbers, frame_width, frame_height):
         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"
+        return None, None, "Failed to fit trajectory"
 
     # Extrapolate to stumps
     frame_max = max(frames) + 10
@@ -136,9 +136,13 @@ def predict_trajectory(positions, frame_numbers, frame_width, frame_height):
     x_pred = poly_func(future_frames, *popt_x)
     y_pred = poly_func(future_frames, *popt_y)
 
-    # Check if trajectory hits stumps
+    # Wicket inline path (center line toward stumps)
     stump_x = frame_width / 2
     stump_y = frame_height
+    inline_x = np.linspace(min(x_coords), stump_x, 100)
+    inline_y = np.interp(inline_x, x_pred, y_pred)
+
+    # Check if trajectory hits stumps
     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:
@@ -146,7 +150,7 @@ def predict_trajectory(positions, frame_numbers, frame_width, frame_height):
             break
 
     lbw_decision = "OUT" if stump_hit else "NOT OUT"
-    return list(zip(future_frames, x_pred, y_pred)), lbw_decision
+    return list(zip(future_frames, x_pred, y_pred)), list(zip(inline_x, inline_y)), lbw_decision
 
 # Map pitch location
 def map_pitch(bounce_point, frame_width, frame_height):
@@ -177,27 +181,7 @@ def estimate_speed(positions, frame_numbers, frame_rate, frame_width):
     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"))
-    
-    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
+# Main Gradio function with video overlay
 def drs_analysis(video):
     # Video is a file path (string) in Hugging Face Spaces
     video_path = video if isinstance(video, str) else "temp_video.mp4"
@@ -205,27 +189,80 @@ def drs_analysis(video):
         with open(video_path, "wb") as f:
             f.write(video.read())
 
+    # Process video for detections
     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)
+    # Predict trajectory and wicket path
+    trajectory, inline_path, 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")
+    # Create output video with overlays
+    output_path = "output_video.mp4"
+    cap = cv2.VideoCapture(video_path)
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_path, fourcc, frame_rate, (frame_width, frame_height))
+
+    frame_count = 0
+    positions_dict = dict(zip(frame_numbers, positions))
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Skip frames for consistency with detection
+        if frame_count % FRAME_SKIP != 0:
+            frame_count += 1
+            continue
+
+        # Overlay ball trajectory (red) and wicket inline path (blue)
+        if frame_count in positions_dict:
+            cv2.circle(frame, (int(positions_dict[frame_count][0]), int(positions_dict[frame_count][1])), 5, (0, 0, 255), -1)  # Red dot
+        if trajectory:
+            traj_x = [int(t[1]) for t in trajectory if t[0] >= frame_count]
+            traj_y = [int(t[2]) for t in trajectory if t[0] >= frame_count]
+            if traj_x and traj_y:
+                for i in range(1, len(traj_x)):
+                    cv2.line(frame, (traj_x[i-1], traj_y[i-1]), (traj_x[i], traj_y[i]), (0, 0, 255), 2)  # Red line
+        if inline_path:
+            inline_x = [int(x) for x, _ in inline_path]
+            inline_y = [int(y) for _, y in inline_path]
+            if inline_x and inline_y:
+                for i in range(1, len(inline_x)):
+                    cv2.line(frame, (inline_x[i-1], inline_y[i-1]), (inline_x[i], inline_y[i]), (255, 0, 0), 2)  # Blue line
+
+        # Overlay pitch map in top-right corner
+        if pitch_x is not None and pitch_y is not None:
+            map_width = 200
+            map_height = int(map_width * PITCH_LENGTH / PITCH_WIDTH)
+            pitch_map = np.zeros((map_height, map_width, 3), dtype=np.uint8)
+            pitch_map[:] = (0, 255, 0)  # Green pitch
+            cv2.rectangle(pitch_map, (0, map_height-10), (map_width, map_height), (0, 51, 51), -1)  # Brown stumps
+            bounce_x = int((pitch_x + PITCH_WIDTH/2) / PITCH_WIDTH * map_width)
+            bounce_y = int((1 - pitch_y / PITCH_LENGTH) * map_height)
+            cv2.circle(pitch_map, (bounce_x, bounce_y), 5, (0, 0, 255), -1)  # Red bounce point
+            frame[0:map_height, frame_width-map_width:frame_width] = cv2.resize(pitch_map, (map_width, map_height))
+
+        # Add text annotations
+        text = f"LBW: {lbw_decision}\nSpeed: {speed_kmh:.2f} km/h"
+        cv2.putText(frame, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
+
+        out.write(frame)
+        frame_count += 1
 
-    fig_pitch = create_pitch_map(pitch_x, pitch_y)
+    cap.release()
+    out.release()
 
     if not isinstance(video, str):
         os.remove(video_path)
 
-    return fig_traj, fig_pitch, f"LBW Decision: {lbw_decision}\nSpeed: {speed_kmh:.2f} km/h", video_path
+    return None, None, None, output_path
 
 # Gradio interface
 with gr.Blocks() as demo: