File size: 2,986 Bytes
41fadd8 530a851 41fadd8 530a851 41fadd8 530a851 41fadd8 530a851 41fadd8 530a851 41fadd8 530a851 41fadd8 530a851 41fadd8 530a851 41fadd8 530a851 41fadd8 530a851 41fadd8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
# video_processor.py
import cv2 as cv
# We import the necessary functions and state from drowsiness_detection
from drowsiness_detection import process_frame, reset_counters, DROWSY_COUNTER, YAWN_COUNTER, HEAD_DOWN_COUNTER
def process_video_with_progress(input_path, output_path, progress_callback=None):
"""
Processes a video file to detect drowsiness, providing progress updates.
Args:
input_path (str): Path to the input video file.
output_path (str): Path to save the processed video file.
progress_callback (function, optional): A function to call with progress updates.
It receives (current_frame, total_frames).
Returns:
dict: A dictionary containing the final detection statistics.
"""
reset_counters() # Ensure all counters are zeroed before starting
video_stream = cv.VideoCapture(input_path)
if not video_stream.isOpened():
raise ValueError(f"Could not open video file {input_path}")
# Get video properties for the writer
fps = video_stream.get(cv.CAP_PROP_FPS)
original_width = int(video_stream.get(cv.CAP_PROP_FRAME_WIDTH))
original_height = int(video_stream.get(cv.CAP_PROP_FRAME_HEIGHT))
total_frames = int(video_stream.get(cv.CAP_PROP_FRAME_COUNT))
# --- FIX: Calculate correct output dimensions to prevent distortion ---
# The process_frame function resizes frames to a fixed width of 640.
output_width = 640
aspect_ratio = original_height / original_width
output_height = int(output_width * aspect_ratio)
output_dims = (output_width, output_height)
# Setup video writer
fourcc = cv.VideoWriter_fourcc(*'mp4v')
video_writer = cv.VideoWriter(output_path, fourcc, fps, output_dims)
frame_count = 0
try:
while True:
ret, frame = video_stream.read()
if not ret:
break
frame_count += 1
processed_frame = process_frame(frame)
# The processed frame is already resized to 640px width.
# We must ensure it fits the calculated output_dims.
# If aspect ratios differ slightly, resizing is a safe fallback.
if processed_frame.shape[1] != output_dims[0] or processed_frame.shape[0] != output_dims[1]:
processed_frame = cv.resize(processed_frame, output_dims)
video_writer.write(processed_frame)
if progress_callback:
progress_callback(frame_count, total_frames)
stats = {
'total_frames_processed': frame_count,
'drowsy_events': DROWSY_COUNTER,
'yawn_events': YAWN_COUNTER,
'head_down_events': HEAD_DOWN_COUNTER
}
return stats
finally:
video_stream.release()
if video_writer:
video_writer.release() |