Update video_processor.py
Browse files- video_processor.py +36 -102
video_processor.py
CHANGED
@@ -1,80 +1,44 @@
|
|
1 |
-
|
2 |
-
Video Processing Utility for Drowsiness Detection
|
3 |
-
This script provides a more robust video processing interface
|
4 |
-
"""
|
5 |
|
6 |
import cv2 as cv
|
7 |
-
import
|
8 |
-
import
|
9 |
-
from datetime import datetime
|
10 |
-
import argparse
|
11 |
-
|
12 |
-
def get_video_info(video_path):
|
13 |
-
"""Get detailed video information"""
|
14 |
-
cap = cv.VideoCapture(video_path)
|
15 |
-
|
16 |
-
if not cap.isOpened():
|
17 |
-
return None
|
18 |
-
|
19 |
-
info = {
|
20 |
-
'fps': cap.get(cv.CAP_PROP_FPS),
|
21 |
-
'width': int(cap.get(cv.CAP_PROP_FRAME_WIDTH)),
|
22 |
-
'height': int(cap.get(cv.CAP_PROP_FRAME_HEIGHT)),
|
23 |
-
'total_frames': int(cap.get(cv.CAP_PROP_FRAME_COUNT)),
|
24 |
-
'duration': cap.get(cv.CAP_PROP_FRAME_COUNT) / cap.get(cv.CAP_PROP_FPS) if cap.get(cv.CAP_PROP_FPS) > 0 else 0,
|
25 |
-
'codec': int(cap.get(cv.CAP_PROP_FOURCC)),
|
26 |
-
'file_size': os.path.getsize(video_path)
|
27 |
-
}
|
28 |
-
|
29 |
-
cap.release()
|
30 |
-
return info
|
31 |
-
|
32 |
-
def create_processing_report(input_path, output_path, stats):
|
33 |
-
"""Create a JSON report of the processing results"""
|
34 |
-
report = {
|
35 |
-
'timestamp': datetime.now().isoformat(),
|
36 |
-
'input_file': input_path,
|
37 |
-
'output_file': output_path,
|
38 |
-
'video_info': get_video_info(input_path),
|
39 |
-
'detection_stats': stats,
|
40 |
-
'processing_info': {
|
41 |
-
'software': 'Drowsiness Detection System',
|
42 |
-
'version': '1.0'
|
43 |
-
}
|
44 |
-
}
|
45 |
-
|
46 |
-
report_path = output_path.replace('.mp4', '_report.json')
|
47 |
-
with open(report_path, 'w') as f:
|
48 |
-
json.dump(report, f, indent=2)
|
49 |
-
|
50 |
-
return report_path
|
51 |
|
52 |
def process_video_with_progress(input_path, output_path, progress_callback=None):
|
53 |
"""
|
54 |
-
|
55 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
56 |
"""
|
57 |
-
#
|
58 |
-
from drowsiness_detection import process_frame, reset_counters
|
59 |
-
from drowsiness_detection import DROWSY_COUNTER, YAWN_COUNTER, HEAD_DOWN_COUNTER
|
60 |
-
|
61 |
-
reset_counters()
|
62 |
|
63 |
-
# Open video file
|
64 |
video_stream = cv.VideoCapture(input_path)
|
65 |
-
|
66 |
if not video_stream.isOpened():
|
67 |
raise ValueError(f"Could not open video file {input_path}")
|
68 |
|
69 |
-
# Get video properties
|
70 |
-
fps =
|
71 |
-
|
72 |
-
|
73 |
total_frames = int(video_stream.get(cv.CAP_PROP_FRAME_COUNT))
|
74 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
75 |
# Setup video writer
|
76 |
fourcc = cv.VideoWriter_fourcc(*'mp4v')
|
77 |
-
video_writer = cv.VideoWriter(output_path, fourcc, fps,
|
78 |
|
79 |
frame_count = 0
|
80 |
|
@@ -85,58 +49,28 @@ def process_video_with_progress(input_path, output_path, progress_callback=None)
|
|
85 |
break
|
86 |
|
87 |
frame_count += 1
|
88 |
-
|
89 |
-
# Process frame
|
90 |
processed_frame = process_frame(frame)
|
91 |
|
92 |
-
#
|
|
|
|
|
|
|
|
|
|
|
93 |
video_writer.write(processed_frame)
|
94 |
|
95 |
-
# Call progress callback if provided
|
96 |
if progress_callback:
|
97 |
progress_callback(frame_count, total_frames)
|
98 |
-
|
99 |
-
# Get final stats
|
100 |
stats = {
|
101 |
-
'
|
102 |
'drowsy_events': DROWSY_COUNTER,
|
103 |
'yawn_events': YAWN_COUNTER,
|
104 |
'head_down_events': HEAD_DOWN_COUNTER
|
105 |
}
|
106 |
-
|
107 |
return stats
|
108 |
|
109 |
finally:
|
110 |
video_stream.release()
|
111 |
-
video_writer
|
112 |
-
|
113 |
-
def main():
|
114 |
-
parser = argparse.ArgumentParser(description='Video Processing Utility for Drowsiness Detection')
|
115 |
-
parser.add_argument('--input', '-i', required=True, help='Input video file path')
|
116 |
-
parser.add_argument('--output', '-o', help='Output video file path (optional)')
|
117 |
-
parser.add_argument('--report', '-r', action='store_true', help='Generate processing report')
|
118 |
-
parser.add_argument('--info', action='store_true', help='Show video information only')
|
119 |
-
|
120 |
-
args = parser.parse_args()
|
121 |
-
|
122 |
-
if not os.path.exists(args.input):
|
123 |
-
print(f"Error: Input file {args.input} does not exist")
|
124 |
-
return
|
125 |
-
|
126 |
-
# Show video info
|
127 |
-
if args.info:
|
128 |
-
info = get_video_info(args.input)
|
129 |
-
if info:
|
130 |
-
print(f"Video Information for: {args.input}")
|
131 |
-
print(f"Resolution: {info['width']}x{info['height']}")
|
132 |
-
print(f"FPS: {info['fps']:.2f}")
|
133 |
-
print(f"Duration: {info['duration']:.2f} seconds")
|
134 |
-
print(f"Total Frames: {info['total_frames']}")
|
135 |
-
print(f"File Size: {info['file_size'] / (1024*1024):.2f} MB")
|
136 |
-
else:
|
137 |
-
print("Error: Could not read video file")
|
138 |
-
return
|
139 |
-
|
140 |
-
# Generate output path if not provided
|
141 |
-
if not args.output:
|
142 |
-
base_name
|
|
|
1 |
+
# video_processor.py
|
|
|
|
|
|
|
2 |
|
3 |
import cv2 as cv
|
4 |
+
# We import the necessary functions and state from drowsiness_detection
|
5 |
+
from drowsiness_detection import process_frame, reset_counters, DROWSY_COUNTER, YAWN_COUNTER, HEAD_DOWN_COUNTER
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
6 |
|
7 |
def process_video_with_progress(input_path, output_path, progress_callback=None):
|
8 |
"""
|
9 |
+
Processes a video file to detect drowsiness, providing progress updates.
|
10 |
+
|
11 |
+
Args:
|
12 |
+
input_path (str): Path to the input video file.
|
13 |
+
output_path (str): Path to save the processed video file.
|
14 |
+
progress_callback (function, optional): A function to call with progress updates.
|
15 |
+
It receives (current_frame, total_frames).
|
16 |
+
|
17 |
+
Returns:
|
18 |
+
dict: A dictionary containing the final detection statistics.
|
19 |
"""
|
20 |
+
reset_counters() # Ensure all counters are zeroed before starting
|
|
|
|
|
|
|
|
|
21 |
|
|
|
22 |
video_stream = cv.VideoCapture(input_path)
|
|
|
23 |
if not video_stream.isOpened():
|
24 |
raise ValueError(f"Could not open video file {input_path}")
|
25 |
|
26 |
+
# Get video properties for the writer
|
27 |
+
fps = video_stream.get(cv.CAP_PROP_FPS)
|
28 |
+
original_width = int(video_stream.get(cv.CAP_PROP_FRAME_WIDTH))
|
29 |
+
original_height = int(video_stream.get(cv.CAP_PROP_FRAME_HEIGHT))
|
30 |
total_frames = int(video_stream.get(cv.CAP_PROP_FRAME_COUNT))
|
31 |
|
32 |
+
# --- FIX: Calculate correct output dimensions to prevent distortion ---
|
33 |
+
# The process_frame function resizes frames to a fixed width of 640.
|
34 |
+
output_width = 640
|
35 |
+
aspect_ratio = original_height / original_width
|
36 |
+
output_height = int(output_width * aspect_ratio)
|
37 |
+
output_dims = (output_width, output_height)
|
38 |
+
|
39 |
# Setup video writer
|
40 |
fourcc = cv.VideoWriter_fourcc(*'mp4v')
|
41 |
+
video_writer = cv.VideoWriter(output_path, fourcc, fps, output_dims)
|
42 |
|
43 |
frame_count = 0
|
44 |
|
|
|
49 |
break
|
50 |
|
51 |
frame_count += 1
|
|
|
|
|
52 |
processed_frame = process_frame(frame)
|
53 |
|
54 |
+
# The processed frame is already resized to 640px width.
|
55 |
+
# We must ensure it fits the calculated output_dims.
|
56 |
+
# If aspect ratios differ slightly, resizing is a safe fallback.
|
57 |
+
if processed_frame.shape[1] != output_dims[0] or processed_frame.shape[0] != output_dims[1]:
|
58 |
+
processed_frame = cv.resize(processed_frame, output_dims)
|
59 |
+
|
60 |
video_writer.write(processed_frame)
|
61 |
|
|
|
62 |
if progress_callback:
|
63 |
progress_callback(frame_count, total_frames)
|
64 |
+
|
|
|
65 |
stats = {
|
66 |
+
'total_frames_processed': frame_count,
|
67 |
'drowsy_events': DROWSY_COUNTER,
|
68 |
'yawn_events': YAWN_COUNTER,
|
69 |
'head_down_events': HEAD_DOWN_COUNTER
|
70 |
}
|
|
|
71 |
return stats
|
72 |
|
73 |
finally:
|
74 |
video_stream.release()
|
75 |
+
if video_writer:
|
76 |
+
video_writer.release()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|