Update gully_drs_core/replay_utils.py

gully_drs_core/replay_utils.py

@@ -1,6 +1,5 @@
-# gully_drs_core/replay_utils.py
-
 import cv2
+import os
 
 def generate_replay(
     frames,
@@ -9,15 +8,22 @@ def generate_replay(
     impact_point,
     decision,
     stump_zone,
+    speed_kmh,
     output_path='output.mp4',
     fps=30
 ):
+    if not frames or not ball_path:
+        return None
+
+    sample = next((f for f in frames if f is not None), None)
+    if sample is None:
+        return None
+
+    height, width = sample.shape[:2]
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    height, width = frames[0].shape[:2]
     out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
 
     for i, frame in enumerate(frames):
-        # Draw ball trajectory
         if i < len(ball_path):
             for j in range(1, i):
                 if j < len(ball_path):
@@ -25,23 +31,19 @@ def generate_replay(
                     pt2 = ball_path[j]
                     cv2.line(frame, pt1, pt2, (0, 0, 255), 2)
 
-        # Draw bounce point
         if bounce_point:
             cv2.circle(frame, bounce_point, 8, (255, 255, 0), -1)
             cv2.putText(frame, "Bounce", (bounce_point[0]+5, bounce_point[1]-10),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1)
 
-        # Draw impact point
         if impact_point:
             cv2.circle(frame, impact_point, 8, (0, 255, 255), -1)
             cv2.putText(frame, "Impact", (impact_point[0]+5, impact_point[1]-10),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
 
-        # Draw stump zone
         x1, y1, x2, y2 = stump_zone
         cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
 
-        # Show decision at top
         cv2.putText(
             frame,
             f"Decision: {decision}",
@@ -52,6 +54,16 @@ def generate_replay(
             2
         )
 
+        cv2.putText(
+            frame,
+            f"Speed: {speed_kmh} km/h",
+            (20, 75),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.9,
+            (255, 255, 255),
+            2
+        )
+
         out.write(frame)
 
     out.release()