Update app.py

app.py

@@ -6,8 +6,6 @@ import cv2
 from PIL import Image
 import io
 import os
-import sys
-import time
 
 class DrowsinessDetector:
     def __init__(self):
@@ -18,11 +16,12 @@ class DrowsinessDetector:
         self.id2label = {0: "notdrowsy", 1: "drowsy"}
         self.label2id = {"notdrowsy": 0, "drowsy": 1}
 
-    def load_model(self, model_path):
-        """Load the ViT model and processor from the specified path or directory"""
+    def load_model(self):
+        """Load the ViT model and processor from Hugging Face Hub"""
         try:
+            model_id = "ckcl/driver-drowsiness-detector"  # 使用你的模型ID
             self.model = ViTForImageClassification.from_pretrained(
-                model_path,  # 直接給資料夾路徑
+                model_id,
                 num_labels=2,
                 id2label=self.id2label,
                 label2id=self.label2id,
@@ -30,7 +29,7 @@ class DrowsinessDetector:
             )
             self.model.eval()
             self.processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
-            print(f"ViT model loaded successfully from {model_path}")
+            print(f"ViT model loaded successfully from {model_id}")
         except Exception as e:
             print(f"Error loading ViT model: {str(e)}")
             raise
@@ -80,88 +79,54 @@ class DrowsinessDetector:
 # Initialize detector
 detector = DrowsinessDetector()
 
-def find_model_file():
-    """Find the model directory or file in common locations"""
-    possible_paths = [
-        "huggingface_model",  # 優先資料夾
-        "pytorch_model.bin",
-        "model_weights.h5",
-        "drowsiness_model.h5",
-        "model/drowsiness_model.h5",
-        "models/drowsiness_model.h5",
-        "huggingface_model/model_weights.h5",
-        "huggingface_model/drowsiness_model.h5",
-        "../model_weights.h5",
-        "../drowsiness_model.h5"
-    ]
-    for path in possible_paths:
-        if os.path.exists(path):
-            return path
-    return None
-
-def load_model():
-    """Load the model"""
-    model_path = find_model_file()
-    
-    if model_path is None:
-        print("\nError: Model file not found!")
-        print("\nPlease ensure one of the following files exists:")
-        print("1. model_weights.h5")
-        print("2. drowsiness_model.h5")
-        print("3. model/drowsiness_model.h5")
-        print("4. models/drowsiness_model.h5")
-        print("\nYou can download the model from Hugging Face Hub or train it using train_model.py")
-        sys.exit(1)
+def process_image(image):
+    """Process a single image"""
+    if image is None:
+        return None, "No image provided"
         
     try:
-        detector.load_model(model_path)
-    except Exception as e:
-        print(f"\nError loading model: {str(e)}")
-        sys.exit(1)
-
-def process_frame(frame):
-    """Process a single frame"""
-    if frame is None:
-        return None
+        # Convert image to numpy array if it's a PIL Image
+        if isinstance(image, Image.Image):
+            image = np.array(image)
         
-    try:
         # Convert frame to RGB if needed
-        if len(frame.shape) == 2:
-            frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
-        elif frame.shape[2] == 4:
-            frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2RGB)
+        if len(image.shape) == 2:
+            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+        elif image.shape[2] == 4:
+            image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
         
         # Make prediction
-        drowsy_prob, face_coords, error = detector.predict(frame)
+        drowsy_prob, face_coords, error = detector.predict(image)
         
         if error:
-            return frame
+            return image, error
             
         if face_coords is not None:
             x, y, w, h = face_coords
             # Draw rectangle around face
             color = (0, 0, 255) if drowsy_prob > 0.7 else (0, 255, 0)
-            cv2.rectangle(frame, (x, y), (x+w, y+h), color, 2)
+            cv2.rectangle(image, (x, y), (x+w, y+h), color, 2)
             
             # Add text
             status = "DROWSY" if drowsy_prob > 0.7 else "ALERT"
-            cv2.putText(frame, f"{status} ({drowsy_prob:.2%})", 
+            cv2.putText(image, f"{status} ({drowsy_prob:.2%})", 
                        (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
-        
-        return frame
+            
+            return image, f"Status: {status} (Confidence: {drowsy_prob:.2%})"
+        else:
+            return image, "No face detected"
         
     except Exception as e:
-        print(f"Error processing frame: {str(e)}")
-        return frame
+        return image, f"Error processing image: {str(e)}"
 
-def process_video(video_input):
+def process_video(video):
     """Process video input"""
-    if video_input is None:
-        return None
+    if video is None:
+        return None, "No video provided"
     
     try:
         # Get input video properties
-        cap = cv2.VideoCapture(video_input)
+        cap = cv2.VideoCapture(video)
         fps = cap.get(cv2.CAP_PROP_FPS)
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -176,7 +141,7 @@ def process_video(video_input):
             if not ret:
                 break
                 
-            processed_frame = process_frame(frame)
+            processed_frame = process_image(frame)[0]
             if processed_frame is not None:
                 out.write(processed_frame)
                 
@@ -186,14 +151,12 @@ def process_video(video_input):
         
         # Check if video was created
         if os.path.exists(temp_output) and os.path.getsize(temp_output) > 0:
-            return temp_output
+            return temp_output, "Video processed successfully"
         else:
-            print("Error: Failed to create output video")
-            return None
+            return None, "Error: Failed to create output video"
         
     except Exception as e:
-        print(f"Error processing video: {str(e)}")
-        return None
+        return None, f"Error processing video: {str(e)}"
     finally:
         # Clean up temporary file
         if 'out' in locals():
@@ -201,27 +164,8 @@ def process_video(video_input):
         if 'cap' in locals():
             cap.release()
 
-def webcam_feed():
-    """Process webcam feed"""
-    try:
-        cap = cv2.VideoCapture(0)
-        while True:
-            ret, frame = cap.read()
-            if not ret:
-                break
-                
-            processed_frame = process_frame(frame)
-            if processed_frame is not None:
-                yield processed_frame
-                
-    except Exception as e:
-        print(f"Error processing webcam feed: {str(e)}")
-        yield None
-    finally:
-        cap.release()
-
 # Load the model at startup
-load_model()
+detector.load_model()
 
 # Create interface
 with gr.Blocks(title="Driver Drowsiness Detection") as demo:
@@ -231,34 +175,38 @@ with gr.Blocks(title="Driver Drowsiness Detection") as demo:
     This system detects driver drowsiness using computer vision and deep learning.
     
     ## Features:
-    - Real-time webcam monitoring
-    - Video file processing
-    - Single image analysis
+    - Image analysis
+    - Video processing
     - Face detection and drowsiness prediction
     """)
     
     with gr.Tabs():
-        with gr.Tab("Webcam"):
-            gr.Markdown("Real-time drowsiness detection using your webcam")
-            webcam_output = gr.Image(label="Live Detection")
-            webcam_button = gr.Button("Start Webcam")
-            webcam_button.click(fn=webcam_feed, inputs=None, outputs=webcam_output)
+        with gr.Tab("Image"):
+            gr.Markdown("Upload an image for drowsiness detection")
+            with gr.Row():
+                image_input = gr.Image(label="Input Image", type="numpy")
+                image_output = gr.Image(label="Processed Image")
+            with gr.Row():
+                status_output = gr.Textbox(label="Status")
+            image_input.change(
+                fn=process_image,
+                inputs=[image_input],
+                outputs=[image_output, status_output]
+            )
             
         with gr.Tab("Video"):
             gr.Markdown("Upload a video file for drowsiness detection")
             with gr.Row():
                 video_input = gr.Video(label="Input Video")
-                video_output = gr.Video(label="Detection Result")
-            video_button = gr.Button("Process Video")
-            video_button.click(fn=process_video, inputs=video_input, outputs=video_output)
-            
-        with gr.Tab("Image"):
-            gr.Markdown("Upload an image for drowsiness detection")
+                video_output = gr.Video(label="Processed Video")
             with gr.Row():
-                image_input = gr.Image(type="numpy", label="Input Image")
-                image_output = gr.Image(label="Detection Result")
-            image_button = gr.Button("Process Image")
-            image_button.click(fn=process_frame, inputs=image_input, outputs=image_output)
+                video_status = gr.Textbox(label="Status")
+            video_input.change(
+                fn=process_video,
+                inputs=[video_input],
+                outputs=[video_output, video_status]
+            )
 
+# Launch the app
 if __name__ == "__main__":
     demo.launch()