Update drs_engine.py

drs_engine.py

@@ -3,32 +3,43 @@ import os
 import tempfile
 import math
 from ultralytics import YOLO
+from pydub import AudioSegment
+import ffmpeg
 
-model = YOLO("best.pt")  # fallback to yolov8n.pt if needed
+model = YOLO("best.pt")
 
-# CONFIGS
-PITCH_LENGTH_METERS = 20.12  # meters
-FRAME_WIDTH = 1280  # assume standard
-FPS = 30            # assume 30fps if not extracted
+PITCH_LENGTH_METERS = 20.12
+FPS = 30
 
-# Distance conversion: use visual calibration if you want real scale
 def estimate_speed(p1, p2, fps):
     pixel_distance = math.hypot(p2[0] - p1[0], p2[1] - p1[1])
-    meters_per_pixel = 0.03  # approximate scale
+    meters_per_pixel = 0.03
     distance_m = pixel_distance * meters_per_pixel
     speed_mps = distance_m * fps
-    return speed_mps * 3.6  # to km/h
+    return speed_mps * 3.6
+
+def add_voice_to_video(video_path, verdict):
+    audio_file = "out.mp3" if verdict == "OUT" else "not_out.mp3"
+    temp_audio = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False).name
+    AudioSegment.from_file(audio_file).export(temp_audio, format="mp3")
+    final_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    (
+        ffmpeg
+        .input(video_path)
+        .output(final_output, vf='scale=trunc(iw/2)*2:trunc(ih/2)*2', audio=temp_audio, vcodec='copy', acodec='aac', strict='experimental')
+        .run(overwrite_output=True)
+    )
+    return final_output
 
 def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         raise Exception("Video open failed")
 
-    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps    = cap.get(cv2.CAP_PROP_FPS) or FPS
+    fps = cap.get(cv2.CAP_PROP_FPS) or FPS
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-
     temp_out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
     out = cv2.VideoWriter(temp_out, fourcc, fps, (width, height))
 
@@ -36,9 +47,9 @@ def process_video(video_path):
     bounce_frame = None
     verdict = "NOT OUT"
     bounce_detected = False
-
     prev_center = None
-    for f_idx in range(int(cap.get(cv2.CAP_PROP_FRAME_COUNT))):
+
+    for _ in range(int(cap.get(cv2.CAP_PROP_FRAME_COUNT))):
         ret, frame = cap.read()
         if not ret:
             break
@@ -47,43 +58,47 @@ def process_video(video_path):
         ball_box = None
         for box in results[0].boxes:
             cls = int(box.cls[0])
-            if cls == 0:  # cricket ball class
+            if cls == 0:
                 x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
                 cx, cy = int((x1 + x2)/2), int((y1 + y2)/2)
                 ball_box = (cx, cy)
-                trajectory.append((f_idx, cx, cy))
+                trajectory.append((cx, cy))
                 cv2.circle(frame, (cx, cy), 8, (0, 0, 255), -1)
-
-                # Simple bounce logic: sudden increase in Y
                 if prev_center and not bounce_detected:
                     dy = cy - prev_center[1]
                     if dy > 15:
-                        bounce_frame = f_idx
                         bounce_detected = True
                         cv2.putText(frame, "Bounce!", (cx, cy - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,255), 2)
                 prev_center = (cx, cy)
-                break  # take one ball only
+                break
 
-        # Draw trajectory line
         for i in range(1, len(trajectory)):
-            _, x1, y1 = trajectory[i-1]
-            _, x2, y2 = trajectory[i]
+            x1, y1 = trajectory[i-1]
+            x2, y2 = trajectory[i]
             cv2.line(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
 
-        # Add speed (estimate from last 2)
         if len(trajectory) >= 2:
-            _, x1, y1 = trajectory[-2]
-            _, x2, y2 = trajectory[-1]
+            x1, y1 = trajectory[-2]
+            x2, y2 = trajectory[-1]
             speed_kmh = estimate_speed((x1, y1), (x2, y2), fps)
             cv2.putText(frame, f"Speed: {int(speed_kmh)} km/h", (30, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,255), 2)
 
-        # OUT/NOT OUT logic (basic): Assume bounce before batsman, and impact near stumps
         if bounce_detected and trajectory:
-            last_x, last_y = trajectory[-1][1:]
-            if 550 < last_x < 750:  # center line assumed
+            last_x, last_y = trajectory[-1]
+            if 550 < last_x < 750:
                 verdict = "OUT"
 
-        # Add verdict text
+        # Draw pitch zones
+        zone_top = int(height * 0.55)
+        zone_bottom = int(height * 0.95)
+        zone_width = width // 3
+        cv2.rectangle(frame, (0, zone_top), (zone_width, zone_bottom), (255, 200, 200), 1)
+        cv2.putText(frame, "Leg", (zone_width//4, zone_top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 200, 200), 1)
+        cv2.rectangle(frame, (zone_width, zone_top), (zone_width*2, zone_bottom), (200, 255, 200), 1)
+        cv2.putText(frame, "Middle", (zone_width + 10, zone_top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200, 255, 200), 1)
+        cv2.rectangle(frame, (zone_width*2, zone_top), (width, zone_bottom), (200, 200, 255), 1)
+        cv2.putText(frame, "Off", (zone_width*2 + 10, zone_top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200, 200, 255), 1)
+
         cv2.putText(frame, f"Decision: {verdict}", (30, height - 30), cv2.FONT_HERSHEY_SIMPLEX, 1.2,
                     (0,255,0) if verdict == "NOT OUT" else (0,0,255), 3)
 
@@ -91,4 +106,5 @@ def process_video(video_path):
 
     cap.release()
     out.release()
-    return temp_out
+
+    return add_voice_to_video(temp_out, verdict)