dschandra commited on
Commit
03f7f25
·
verified ·
1 Parent(s): 3227154

Create visualizer.py

Browse files
Files changed (1) hide show
  1. visualizer.py +64 -0
visualizer.py ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+ import tempfile
4
+ import os
5
+ from moviepy.editor import ImageSequenceClip
6
+
7
+ def draw_virtual_stumps(frame, stump_zone=(280, 360), pitch_y=570):
8
+ """
9
+ Draw stumps and crease lines on the pitch.
10
+ """
11
+ for x in range(stump_zone[0], stump_zone[1] + 1, 20):
12
+ cv2.line(frame, (x, pitch_y), (x, pitch_y - 60), (255, 255, 255), 2)
13
+ cv2.line(frame, (stump_zone[0] - 40, pitch_y), (stump_zone[1] + 40, pitch_y), (255, 255, 255), 2)
14
+
15
+ def draw_trajectory(frame, trajectory_points):
16
+ """
17
+ Draws the predicted ball trajectory curve on frame.
18
+ """
19
+ for pt in trajectory_points:
20
+ x, y = pt
21
+ cv2.circle(frame, (int(x), int(y)), 5, (0, 255, 0), -1)
22
+
23
+ def add_decision_text(frame, decision):
24
+ """
25
+ Overlay OUT / NOT OUT decision on top of the frame.
26
+ """
27
+ color = (0, 0, 255) if decision == "OUT" else (0, 255, 0)
28
+ cv2.putText(frame, f"Decision: {decision}", (30, 50),
29
+ cv2.FONT_HERSHEY_SIMPLEX, 1.5, color, 3)
30
+
31
+ def generate_output_video(frames, detection_data, prediction_data, fps=25):
32
+ """
33
+ Generates replay video with overlays and returns its file path and decision.
34
+ """
35
+ impact_frame = detection_data["impact_frame"]
36
+ decision = prediction_data["decision"]
37
+ trajectory = prediction_data["trajectory_points"]
38
+
39
+ annotated_frames = []
40
+
41
+ for idx, frame in enumerate(frames):
42
+ frame_copy = frame.copy()
43
+ draw_virtual_stumps(frame_copy)
44
+
45
+ if idx <= impact_frame:
46
+ # Draw live ball trajectory up to impact
47
+ draw_trajectory(frame_copy, trajectory[:idx + 1])
48
+ else:
49
+ # Draw full trajectory post-impact
50
+ draw_trajectory(frame_copy, trajectory)
51
+
52
+ if idx == impact_frame:
53
+ cv2.rectangle(frame_copy, (0, 0), (frame.shape[1], frame.shape[0]), (255, 0, 0), 6)
54
+
55
+ add_decision_text(frame_copy, decision)
56
+ annotated_frames.append(cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB))
57
+
58
+ # Save as video using moviepy
59
+ temp_dir = tempfile.mkdtemp()
60
+ output_path = os.path.join(temp_dir, "lbw_replay.mp4")
61
+ clip = ImageSequenceClip(annotated_frames, fps=fps)
62
+ clip.write_videofile(output_path, codec="libx264", audio=False, verbose=False, logger=None)
63
+
64
+ return output_path, decision