Update drowsiness_detector.py

drowsiness_detector.py

@@ -1,122 +1,122 @@
-import os
-import cv2
-import numpy as np
-from speed_detector import SpeedDetector
-from face_analyzer import FaceAnalyzer
-import pandas as pd
-import time
-
-class DrowsinessDetector:
-    def __init__(self):
-        self.speed_detector = SpeedDetector()
-        self.face_analyzer = FaceAnalyzer()
-        
-    def process_frame(self, frame_path, face_path):
-        """
-        處理單個幀
-        :param frame_path: 場景圖片路徑
-        :param face_path: 人臉圖片路徑
-        :return: (速度, 是否犯困)
-        """
-        try:
-            # 讀取圖片
-            frame = cv2.imread(frame_path)
-            face = cv2.imread(face_path)
-            
-            if frame is None or face is None:
-                print(f"處理 {os.path.basename(frame_path)} 時出錯: 無法讀取圖片")
-                return None, None
-                
-            # 檢測速度
-            speed = self.speed_detector.detect_speed(frame)
-            
-            # 檢測是否犯困
-            is_drowsy = self.face_analyzer.is_drowsy(face)
-            
-            return speed, is_drowsy
-        except Exception as e:
-            print(f"處理 {os.path.basename(frame_path)} 時出錯: {str(e)}")
-            return None, None
-        
-    def process_video_folder(self, folder_path):
-        """
-        處理一個視頻文件夾中的所有幀
-        :param folder_path: 視頻文件夾路徑
-        :return: 處理結果列表
-        """
-        results = []
-        
-        # 獲取所有幀圖片
-        frame_files = [f for f in os.listdir(folder_path) if f.endswith('.jpg') and not f.endswith('_face.jpg')]
-        total_frames = len(frame_files)
-        
-        for i, frame_file in enumerate(frame_files, 1):
-            # 構建完整的文件路徑
-            frame_path = os.path.join(folder_path, frame_file)
-            face_path = os.path.join(folder_path, frame_file.replace('.jpg', '_face.jpg'))
-            
-            # 顯示進度
-            print(f"\r處理進度: {i}/{total_frames} ({i/total_frames*100:.1f}%)", end="")
-            
-            try:
-                speed, is_drowsy = self.process_frame(frame_path, face_path)
-                if speed is not None and is_drowsy is not None:
-                    results.append({
-                        'frame': frame_file,
-                        'speed': speed,
-                        'is_drowsy': is_drowsy
-                    })
-            except KeyboardInterrupt:
-                print("\n檢測到中斷，保存當前結果...")
-                return results
-            except Exception as e:
-                print(f"\n處理 {frame_file} 時出錯: {str(e)}")
-                continue
-                
-        print()  # 換行
-        return results
-        
-def main():
-    # 初始化檢測器
-    detector = DrowsinessDetector()
-    
-    # 獲取所有視頻文件夾
-    dataset_path = os.path.join('dataset', 'driver')
-    video_folders = [f for f in os.listdir(dataset_path) if os.path.isdir(os.path.join(dataset_path, f))]
-    total_folders = len(video_folders)
-    
-    all_results = []
-    batch_size = 100  # 每處理100個文件夾保存一次結果
-    
-    try:
-        # 處理每個視頻文件夾
-        for i, folder in enumerate(video_folders, 1):
-            print(f"\n處理文件夾 {i}/{total_folders}: {folder}")
-            folder_path = os.path.join(dataset_path, folder)
-            results = detector.process_video_folder(folder_path)
-            all_results.extend(results)
-            
-            # 每處理完一批文件夾就保存一次結果
-            if i % batch_size == 0 or i == total_folders:
-                print(f"\n保存第 {i//batch_size + 1} 批結果...")
-                df = pd.DataFrame(all_results)
-                df.to_csv(f'drowsiness_results_batch_{i//batch_size + 1}.csv', index=False)
-                all_results = []  # 清空結果列表
-                
-    except KeyboardInterrupt:
-        print("\n檢測到中斷，保存當前結果...")
-        if all_results:
-            df = pd.DataFrame(all_results)
-            df.to_csv('drowsiness_results_final.csv', index=False)
-            print("結果已保存到 drowsiness_results_final.csv")
-    except Exception as e:
-        print(f"\n發生錯誤: {str(e)}")
-        if all_results:
-            df = pd.DataFrame(all_results)
-            df.to_csv('drowsiness_results_error.csv', index=False)
-            print("結果已保存到 drowsiness_results_error.csv")
-    finally:
-        print("\n處理完成")
-    
-if __name__ == "__main__":
+import os
+import cv2
+import numpy as np
+from speed_detector import SpeedDetector
+from face_analyzer import FaceAnalyzer
+import pandas as pd
+import time
+
+class DrowsinessDetector:
+    def __init__(self):
+        self.speed_detector = SpeedDetector()
+        self.face_analyzer = FaceAnalyzer()
+        
+    def process_frame(self, frame_path, face_path):
+        """
+        Process a single frame
+        :param frame_path: Path to scene image
+        :param face_path: Path to face image
+        :return: (speed, is_drowsy)
+        """
+        try:
+            # Read images
+            frame = cv2.imread(frame_path)
+            face = cv2.imread(face_path)
+            
+            if frame is None or face is None:
+                print(f"Error processing {os.path.basename(frame_path)}: Unable to read image")
+                return None, None
+                
+            # Detect speed
+            speed = self.speed_detector.detect_speed(frame)
+            
+            # Detect drowsiness
+            is_drowsy = self.face_analyzer.is_drowsy(face)
+            
+            return speed, is_drowsy
+        except Exception as e:
+            print(f"Error processing {os.path.basename(frame_path)}: {str(e)}")
+            return None, None
+        
+    def process_video_folder(self, folder_path):
+        """
+        Process all frames in a video folder
+        :param folder_path: Path to video folder
+        :return: Processing results list
+        """
+        results = []
+        
+        # Get all frame images
+        frame_files = [f for f in os.listdir(folder_path) if f.endswith('.jpg') and not f.endswith('_face.jpg')]
+        total_frames = len(frame_files)
+        
+        for i, frame_file in enumerate(frame_files, 1):
+            # Build full file path
+            frame_path = os.path.join(folder_path, frame_file)
+            face_path = os.path.join(folder_path, frame_file.replace('.jpg', '_face.jpg'))
+            
+            # Show progress
+            print(f"\rProcessing progress: {i}/{total_frames} ({i/total_frames*100:.1f}%)", end="")
+            
+            try:
+                speed, is_drowsy = self.process_frame(frame_path, face_path)
+                if speed is not None and is_drowsy is not None:
+                    results.append({
+                        'frame': frame_file,
+                        'speed': speed,
+                        'is_drowsy': is_drowsy
+                    })
+            except KeyboardInterrupt:
+                print("\nInterrupt detected, saving current results...")
+                return results
+            except Exception as e:
+                print(f"\nError processing {frame_file}: {str(e)}")
+                continue
+                
+        print()  # New line
+        return results
+        
+def main():
+    # Initialize detector
+    detector = DrowsinessDetector()
+    
+    # Get all video folders
+    dataset_path = os.path.join('dataset', 'driver')
+    video_folders = [f for f in os.listdir(dataset_path) if os.path.isdir(os.path.join(dataset_path, f))]
+    total_folders = len(video_folders)
+    
+    all_results = []
+    batch_size = 100  # Save results after processing 100 folders
+    
+    try:
+        # Process each video folder
+        for i, folder in enumerate(video_folders, 1):
+            print(f"\nProcessing folder {i}/{total_folders}: {folder}")
+            folder_path = os.path.join(dataset_path, folder)
+            results = detector.process_video_folder(folder_path)
+            all_results.extend(results)
+            
+            # Save results after processing each batch of folders
+            if i % batch_size == 0 or i == total_folders:
+                print(f"\nSaving results for batch {i//batch_size + 1}...")
+                df = pd.DataFrame(all_results)
+                df.to_csv(f'drowsiness_results_batch_{i//batch_size + 1}.csv', index=False)
+                all_results = []  # Clear results list
+                
+    except KeyboardInterrupt:
+        print("\nInterrupt detected, saving current results...")
+        if all_results:
+            df = pd.DataFrame(all_results)
+            df.to_csv('drowsiness_results_final.csv', index=False)
+            print("The result has been saved to drowsiness_results_final.csv")
+    except Exception as e:
+        print(f"\nA error occurred: {str(e)}")
+        if all_results:
+            df = pd.DataFrame(all_results)
+            df.to_csv('drowsiness_results_error.csv', index=False)
+            print("Results saved to drawsiness_results_error.csv")
+    finally:
+        print("\nProcessing completed")
+    
+if __name__ == "__main__":
     main()