lbw_drs_app / video_processor.py
dschandra's picture
Update video_processor.py
18d35d6 verified
import cv2
import os
import tempfile
from detector import LBWDetector
from utils import draw_boxes, overlay_decision_text
def process_video(video_path, output_path="output.mp4"):
"""Process a video, detect objects, make LBW decisions, and return output path and decision."""
detector = LBWDetector()
cap = cv2.VideoCapture(video_path)
width = int(cap.get(3))
height = int(cap.get(4))
fps = cap.get(cv2.CAP_PROP_FPS)
out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
stump_zone = (width // 3, 2 * width // 3)
final_decision = "Not Out"
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
detections, class_names = detector.detect_objects(frame)
frame = draw_boxes(frame, detections, class_names)
decision = "Not Out"
ball_bbox = None
player_bbox = None
for x1, y1, x2, y2, conf, cls_id in detections:
label = class_names[int(cls_id)]
if label == "ball":
ball_bbox = (x1, y1, x2, y2)
elif label in ["player", "pads"]:
player_bbox = (x1, y1, x2, y2)
if ball_bbox and player_bbox:
ball_x_center = (ball_bbox[0] + ball_bbox[2]) / 2
if (stump_zone[0] <= ball_x_center <= stump_zone[1] and
ball_bbox[1] < player_bbox[3] and ball_bbox[3] > player_bbox[1]):
decision = "Out"
final_decision = "Out"
frame = overlay_decision_text(frame, decision)
out.write(frame)
cap.release()
out.release()
return output_path, final_decision