Update drs_engine.py

drs_engine.py

@@ -1,74 +1,3 @@
-import cv2
-import os
-import math
-import numpy as np
-import tempfile
-import matplotlib.pyplot as plt
-from ultralytics import YOLO
-from pydub import AudioSegment
-import ffmpeg
-from scipy.optimize import curve_fit
-from scipy.interpolate import UnivariateSpline
-
-model = YOLO("best.pt")
-FPS = 30
-
-def estimate_speed(p1, p2, fps):
-    dist = math.hypot(p2[0] - p1[0], p2[1] - p1[1])
-    meters_per_pixel = 0.03
-    mps = dist * meters_per_pixel * fps
-    return mps * 3.6
-
-def add_voice_to_video(video_path, verdict):
-    audio_file = "out.mp3" if verdict == "OUT" else "not_out.mp3"
-    audio_temp = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False).name
-    AudioSegment.from_file(audio_file).export(audio_temp, format="mp3")
-    final_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-    (
-        ffmpeg
-        .input(video_path)
-        .output(final_output, audio=audio_temp, vcodec='copy', acodec='aac', strict='experimental')
-        .run(overwrite_output=True)
-    )
-    return final_output
-
-def extend_trajectory_with_rotation(points, bounce_idx, final_x=20.12):
-    x_vals = [pt[0] for pt in points]
-    y_vals = [pt[1] for pt in points]
-    if bounce_idx is None or bounce_idx >= len(points) - 2:
-        return x_vals, y_vals
-    x_pre = x_vals[:bounce_idx]
-    y_pre = y_vals[:bounce_idx]
-    def poly2(x, a, b, c): return a*x**2 + b*x + c
-    try:
-        popt, _ = curve_fit(poly2, x_pre, y_pre)
-        x_post = np.linspace(x_vals[bounce_idx], final_x, 50)
-        curve_shift = np.linspace(0, 0.05, len(x_post))
-        y_post = poly2(x_post, *popt) + curve_shift
-        return x_vals + list(x_post), y_vals + list(y_post)
-    except:
-        return x_vals, y_vals
-
-def draw_top_down_trajectory(points, bounce_frame_idx, output_path):
-    if len(points) < 4:
-        return None
-    x_vals = [pt[0] for pt in points]
-    y_vals = [pt[1] for pt in points]
-    x_ext, y_ext = extend_trajectory_with_rotation(points, bounce_frame_idx)
-    plt.figure(figsize=(10, 3))
-    plt.plot(x_ext, y_ext, 'r-', label='Predicted Trajectory')
-    plt.scatter(x_vals, y_vals, c='blue', s=10, label='Detected Points')
-    plt.axvline(x=17.68, color='gray', linestyle='--', label='Stumps (17.68m)')
-    plt.title("Top-Down Predicted Ball Path")
-    plt.xlabel("Pitch Length (m)")
-    plt.ylabel("Lateral Movement (m)")
-    plt.grid(True)
-    plt.legend()
-    image_path = output_path.replace(".mp4", "_trajectory.png")
-    plt.savefig(image_path)
-    plt.close()
-    return image_path
-
 def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
@@ -81,7 +10,11 @@ def process_video(video_path):
     temp_out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
     out = cv2.VideoWriter(temp_out, fourcc, fps, (width, height))
 
-    # Auto homography
+    # ⏱️ Trim to 4 seconds max
+    max_frames = int(fps * 4)
+    frame_count = 0
+
+    # Homography
     src = np.array([
         [width * 0.3, height * 0.4],
         [width * 0.7, height * 0.4],
@@ -102,14 +35,14 @@ def process_video(video_path):
     trajectory, real_trajectory = [], []
     bounce_detected = False
     bounce_frame_idx = None
-    frame_index = 0
     prev_center = None
     verdict = "NOT OUT"
 
-    while True:
+    while frame_count < max_frames:
         ret, frame = cap.read()
         if not ret:
             break
+
         results = model(frame)
         for box in results[0].boxes:
             if int(box.cls[0]) == 0:
@@ -121,7 +54,7 @@ def process_video(video_path):
                 cv2.circle(frame, (cx, cy), 8, (0, 0, 255), -1)
                 if prev_center and not bounce_detected and cy - prev_center[1] > 15:
                     bounce_detected = True
-                    bounce_frame_idx = frame_index
+                    bounce_frame_idx = frame_count
                     cv2.putText(frame, "Bounce!", (cx, cy - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,255), 2)
                 prev_center = (cx, cy)
                 break
@@ -146,11 +79,18 @@ def process_video(video_path):
                     (0,255,0) if verdict == "NOT OUT" else (0,0,255), 3)
 
         out.write(frame)
-        frame_index += 1
+        frame_count += 1
 
     cap.release()
     out.release()
+
+    # Top-down chart
     topdown_image = draw_top_down_trajectory(real_trajectory, bounce_frame_idx, temp_out)
-    final_video = add_voice_to_video(temp_out, verdict)
-    return final_video, topdown_image
 
+    # Optional: voice-over
+    try:
+        final_video = add_voice_to_video(temp_out, verdict)
+    except:
+        final_video = temp_out
+
+    return final_video, topdown_image