simran0608 commited on
Commit
41fadd8
·
verified ·
1 Parent(s): 3e098b3

Update video_processor.py

Browse files
Files changed (1) hide show
  1. 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 os
8
- import json
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
- Process video with progress callback
55
- progress_callback: function that takes (current_frame, total_frames)
 
 
 
 
 
 
 
 
56
  """
57
- # Import the drowsiness detection functions
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 = int(video_stream.get(cv.CAP_PROP_FPS))
71
- width = int(video_stream.get(cv.CAP_PROP_FRAME_WIDTH))
72
- height = int(video_stream.get(cv.CAP_PROP_FRAME_HEIGHT))
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, (640, 480))
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
- # Write frame to output video
 
 
 
 
 
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
- 'total_frames': frame_count,
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.release()
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()