Spaces:
Sleeping
Sleeping
Update video_processor.py
Browse files- video_processor.py +29 -3
video_processor.py
CHANGED
@@ -5,6 +5,7 @@ from detector import LBWDetector
|
|
5 |
from utils import draw_boxes, overlay_decision_text
|
6 |
|
7 |
def process_video(video_path, output_path="output.mp4"):
|
|
|
8 |
detector = LBWDetector()
|
9 |
cap = cv2.VideoCapture(video_path)
|
10 |
|
@@ -14,17 +15,42 @@ def process_video(video_path, output_path="output.mp4"):
|
|
14 |
|
15 |
out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
|
16 |
|
|
|
|
|
|
|
|
|
17 |
while cap.isOpened():
|
18 |
ret, frame = cap.read()
|
19 |
if not ret:
|
20 |
break
|
|
|
|
|
21 |
detections, class_names = detector.detect_objects(frame)
|
22 |
frame = draw_boxes(frame, detections, class_names)
|
23 |
-
|
24 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
25 |
frame = overlay_decision_text(frame, decision)
|
26 |
out.write(frame)
|
27 |
|
28 |
cap.release()
|
29 |
out.release()
|
30 |
-
return output_path
|
|
|
5 |
from utils import draw_boxes, overlay_decision_text
|
6 |
|
7 |
def process_video(video_path, output_path="output.mp4"):
|
8 |
+
"""Process a video, detect objects, make LBW decisions, and return output path and decision."""
|
9 |
detector = LBWDetector()
|
10 |
cap = cv2.VideoCapture(video_path)
|
11 |
|
|
|
15 |
|
16 |
out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
|
17 |
|
18 |
+
# Assume stumps are in the center third of the frame (adjust as needed)
|
19 |
+
stump_zone = (width // 3, 2 * width // 3)
|
20 |
+
final_decision = "Not Out" # Default decision
|
21 |
+
|
22 |
while cap.isOpened():
|
23 |
ret, frame = cap.read()
|
24 |
if not ret:
|
25 |
break
|
26 |
+
|
27 |
+
# Detect objects (ball, player/pads)
|
28 |
detections, class_names = detector.detect_objects(frame)
|
29 |
frame = draw_boxes(frame, detections, class_names)
|
30 |
+
|
31 |
+
# Simple LBW decision logic
|
32 |
+
decision = "Not Out" # Frame-level decision
|
33 |
+
ball_bbox = None
|
34 |
+
player_bbox = None
|
35 |
+
|
36 |
+
for x1, y1, x2, y2, conf, cls_id in detections:
|
37 |
+
label = class_names[int(cls_id)]
|
38 |
+
if label == "ball":
|
39 |
+
ball_bbox = (x1, y1, x2, y2)
|
40 |
+
elif label in ["player", "pads"]: # Adjust based on your model's class names
|
41 |
+
player_bbox = (x1, y1, x2, y2)
|
42 |
+
|
43 |
+
if ball_bbox and player_bbox:
|
44 |
+
ball_x_center = (ball_bbox[0] + ball_bbox[2]) / 2
|
45 |
+
# Check if ball hits player in stump zone
|
46 |
+
if (stump_zone[0] <= ball_x_center <= stump_zone[1] and
|
47 |
+
ball_bbox[1] < player_bbox[3] and ball_bbox[3] > player_bbox[1]): # Vertical overlap
|
48 |
+
decision = "Out"
|
49 |
+
final_decision = "Out" # Update final decision if "Out" is detected
|
50 |
+
|
51 |
frame = overlay_decision_text(frame, decision)
|
52 |
out.write(frame)
|
53 |
|
54 |
cap.release()
|
55 |
out.release()
|
56 |
+
return output_path, final_decision # Return both output path and final decision
|