Update gully_drs_core/replay_utils.py

gully_drs_core/replay_utils.py

@@ -1,61 +1,21 @@
+# gully_drs_core/replay_utils.py
 import cv2
 import numpy as np
-from scipy.interpolate import CubicSpline
 import os
-from gully_drs_core.video_utils import get_video_properties
 
-def generate_replay(video_path, ball_positions, stump_zone, decision, speed_kmh, bounce_point):
-    try:
-        # Get video properties
-        fps, width, height = get_video_properties(video_path)
-        
-        # Initialize video capture and output
-        cap = cv2.VideoCapture(video_path)
-        output_path = os.path.join("outputs", f"replay_{os.path.basename(video_path)}")
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-        
-        frame_idx = 0
-        
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-            
-            # Draw stump zone
-            cv2.polylines(frame, [np.array(stump_zone)], isClosed=True, color=(0, 0, 255), thickness=2)
-            
-            # Draw ball position
-            if frame_idx < len(ball_positions):
-                x, y = ball_positions[frame_idx]
-                cv2.circle(frame, (x, y), 5, (0, 255, 0), -1)
-            
-            # Draw bounce point
-            if bounce_point and frame_idx >= len(ball_positions) // 2:
-                cv2.circle(frame, bounce_point, 8, (255, 255, 0), -1)
-            
-            # Draw Bezier curve trajectory
-            if len(ball_positions) > 3 and frame_idx < len(ball_positions):
-                points = np.array(ball_positions[:frame_idx + 1])
-                if len(points) > 3:
-                    t = np.linspace(0, 1, len(points))
-                    cs_x = CubicSpline(t, points[:, 0])
-                    cs_y = CubicSpline(t, points[:, 1])
-                    t_fine = np.linspace(0, 1, min(100, len(points) * 10))
-                    curve_x = cs_x(t_fine).astype(int)
-                    curve_y = cs_y(t_fine).astype(int)
-                    for i in range(1, len(curve_x)):
-                        cv2.line(frame, (curve_x[i-1], curve_y[i-1]), (curve_x[i], curve_y[i]), (255, 0, 0), 2)
-            
-            # Add text overlays
-            cv2.putText(frame, f"Decision: {decision}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
-            cv2.putText(frame, f"Speed: {speed_kmh:.2f} km/h", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
-            
-            out.write(frame)
-            frame_idx += 1
-        
-        cap.release()
-        out.release()
-        return output_path
-    except Exception as e:
-        raise Exception(f"Error generating replay: {str(e)}")
+def generate_replay(frames, ball_path, output_path='output.mp4', fps=30):
+    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):
+        if i < len(ball_path):
+            for j in range(1, i):
+                if j < len(ball_path):
+                    pt1 = ball_path[j-1]
+                    pt2 = ball_path[j]
+                    cv2.line(frame, pt1, pt2, (0, 0, 255), 2)
+        out.write(frame)
+
+    out.release()
+    return output_path