Update visualizer.py

visualizer.py

@@ -1,50 +1,41 @@
+# visualizer.py
 import cv2
-import numpy as np
-import tempfile
 import os
+import uuid
+from utils import save_video
 
-def draw_virtual_stumps(frame, stump_zone=(280, 360), pitch_y=570):
-    for x in range(stump_zone[0], stump_zone[1] + 1, 20):
-        cv2.line(frame, (x, pitch_y), (x, pitch_y - 60), (255, 255, 255), 2)
-    cv2.line(frame, (stump_zone[0] - 40, pitch_y), (stump_zone[1] + 40, pitch_y), (255, 255, 255), 2)
+OUTPUT_DIR = "output_videos"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-def draw_trajectory(frame, trajectory_points):
-    for pt in trajectory_points:
-        x, y = pt
-        cv2.circle(frame, (int(x), int(y)), 5, (0, 255, 0), -1)
+def draw_visuals(frames, ball_positions, trajectory, impact_frame_idx, decision):
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    output_frames = []
 
-def add_decision_text(frame, decision):
-    color = (0, 0, 255) if decision == "OUT" else (0, 255, 0)
-    cv2.putText(frame, f"Decision: {decision}", (30, 50),
-                cv2.FONT_HERSHEY_SIMPLEX, 1.5, color, 3)
+    for i, frame in enumerate(frames):
+        frame_draw = frame.copy()
 
-def generate_output_video(frames, detection_data, prediction_data, fps=25):
-    impact_frame = detection_data["impact_frame"]
-    decision = prediction_data["decision"]
-    trajectory = prediction_data["trajectory_points"]
+        # Draw past ball positions
+        for _, x, y in ball_positions:
+            cv2.circle(frame_draw, (x, y), 5, (0, 255, 255), -1)
 
-    height, width, _ = frames[0].shape
-    temp_dir = tempfile.mkdtemp()
-    output_path = os.path.join(temp_dir, "lbw_replay.mp4")
+        # Draw trajectory after impact
+        if i == impact_frame_idx:
+            for x, y in trajectory:
+                cv2.circle(frame_draw, (x, y), 4, (0, 0, 255), -1)
 
-    # Define video writer
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+        # Highlight impact frame
+        if i == impact_frame_idx:
+            cv2.putText(frame_draw, "IMPACT!", (50, 50), font, 1.0, (0, 0, 255), 2)
+            cv2.rectangle(frame_draw, (10, 10), (630, 470), (0, 0, 255), 4)
 
-    for idx, frame in enumerate(frames):
-        frame_copy = frame.copy()
-        draw_virtual_stumps(frame_copy)
+        # Final decision on last frame
+        if i == len(frames) - 1:
+            cv2.putText(frame_draw, f"Decision: {decision}", (100, 420), font, 1.2, (255, 0, 0), 3)
 
-        if idx <= impact_frame:
-            draw_trajectory(frame_copy, trajectory[:idx + 1])
-        else:
-            draw_trajectory(frame_copy, trajectory)
+        output_frames.append(frame_draw)
 
-        if idx == impact_frame:
-            cv2.rectangle(frame_copy, (0, 0), (width, height), (255, 0, 0), 6)
-
-        add_decision_text(frame_copy, decision)
-        out.write(frame_copy)
-
-    out.release()
-    return output_path, decision
+    # Save annotated video
+    video_id = str(uuid.uuid4())
+    out_path = os.path.join(OUTPUT_DIR, f"{video_id}_annotated.mp4")
+    save_video(output_frames, out_path)
+    return out_path