Update app.py

app.py

@@ -1,104 +1,190 @@
-import gradio as gr
+import os
 import cv2
+import gradio as gr
 import requests
-import os
-
 from ultralytics import YOLO
 
+# ==== CẤU HÌNH PHÁT HIỆN NGỦ ====
+# Các tên class trong model được coi là "ngủ gật"
+SLEEPY_CLASS_NAMES = {
+    "drowsy",
+    "sleepy",
+    "closed_eyes",
+    "sleep",
+    "ngủ",
+    "buonngu",
+}
+
+# Ngưỡng confidence để kết luận là ngủ
+SLEEP_CONF_THRESHOLD = 0.4
+
+# ==== DEMO FILES (có thể bỏ nếu không cần) ====
 file_urls = [
-    'https://www.dropbox.com/s/b5g97xo901zb3ds/pothole_example.jpg?dl=1',
-    'https://www.dropbox.com/s/86uxlxxlm1iaexa/pothole_screenshot.png?dl=1',
-    'https://www.dropbox.com/s/7sjfwncffg8xej2/video_7.mp4?dl=1'
+    # Bạn có thể thay bằng ảnh/video tài xế của bạn
+    "https://www.dropbox.com/s/b5g97xo901zb3ds/pothole_example.jpg?dl=1",
+    "https://www.dropbox.com/s/86uxlxxlm1iaexa/pothole_screenshot.png?dl=1",
+    "https://www.dropbox.com/s/7sjfwncffg8xej2/video_7.mp4?dl=1",
 ]
 
+
 def download_file(url, save_name):
-    url = url
     if not os.path.exists(save_name):
         file = requests.get(url)
-        open(save_name, 'wb').write(file.content)
+        open(save_name, "wb").write(file.content)
+
 
 for i, url in enumerate(file_urls):
-    if 'mp4' in file_urls[i]:
-        download_file(
-            file_urls[i],
-            f"video.mp4"
-        )
+    if "mp4" in url:
+        download_file(url, "video.mp4")
     else:
-        download_file(
-            file_urls[i],
-            f"image_{i}.jpg"
-        )
+        download_file(url, f"image_{i}.jpg")
 
-model = YOLO('best.pt')
-path  = [['image_0.jpg'], ['image_1.jpg']]
-video_path = [['video.mp4']]
+# ==== LOAD MODEL YOLO (đã train phát hiện buồn ngủ) ====
+model = YOLO("best.pt")
+
+image_examples = [["image_0.jpg"], ["image_1.jpg"]]
+video_examples = [["video.mp4"]]
+
+
+def _normalize_name(name: str) -> str:
+    return name.lower().replace(" ", "_")
+
+
+def draw_and_decide_state(image, results):
+    """
+    Vẽ bounding box + label lên ảnh
+    Đồng thời quyết định xem tài xế đang ngủ hay tỉnh
+    """
+    sleepy_detected = False
+
+    names = results.names
+    boxes = results.boxes.xyxy
+    confs = results.boxes.conf
+    clss = results.boxes.cls
+
+    for i, box in enumerate(boxes):
+        x1, y1, x2, y2 = map(int, box)
+        cls_id = int(clss[i])
+        conf = float(confs[i])
+        cls_name = names[cls_id]
+
+        norm_name = _normalize_name(cls_name)
+        if norm_name in SLEEPY_CLASS_NAMES and conf >= SLEEP_CONF_THRESHOLD:
+            sleepy_detected = True
+
+        label = f"{cls_name} ({conf:.2f})"
+        color = (0, 0, 255) if norm_name in SLEEPY_CLASS_NAMES else (0, 255, 0)
 
-def show_preds_image(image_path):
-    image = cv2.imread(image_path)
-    outputs = model.predict(source=image_path)
-    results = outputs[0].cpu().numpy()
-    for i, det in enumerate(results.boxes.xyxy):
         cv2.rectangle(
             image,
-            (int(det[0]), int(det[1])),
-            (int(det[2]), int(det[3])),
-            color=(0, 0, 255),
+            (x1, y1),
+            (x2, y2),
+            color=color,
             thickness=2,
-            lineType=cv2.LINE_AA
+            lineType=cv2.LINE_AA,
+        )
+        cv2.putText(
+            image,
+            label,
+            (x1, max(y1 - 10, 10)),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.6,
+            color,
+            2,
+            lineType=cv2.LINE_AA,
         )
+
+    # Dòng trạng thái tổng quát
+    if sleepy_detected:
+        state_text = "NGỦ GỤC / DROWSY"
+        state_color = (0, 0, 255)
+    else:
+        state_text = "TỈNH TÁO / ALERT"
+        state_color = (0, 255, 0)
+
+    cv2.putText(
+        image,
+        state_text,
+        (10, 30),
+        cv2.FONT_HERSHEY_SIMPLEX,
+        1.0,
+        state_color,
+        2,
+        lineType=cv2.LINE_AA,
+    )
+
+    return image, sleepy_detected
+
+
+# ==== 1. ẢNH TĨNH ====
+def show_preds_image(image_path):
+    image = cv2.imread(image_path)
+
+    # conf=0.25 cho YOLO, bạn có thể giảm nếu muốn nhạy hơn
+    outputs = model.predict(source=image_path, conf=0.25)
+    results = outputs[0].cpu().numpy()
+
+    image, _ = draw_and_decide_state(image, results)
     return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
+
 inputs_image = [
-    gr.components.Image(type="filepath", label="Input Image"),
+    gr.Image(type="filepath", label="Ảnh đầu vào (driver image)"),
 ]
 outputs_image = [
-    gr.components.Image(type="numpy", label="Output Image"),
+    gr.Image(type="numpy", label="Kết quả nhận diện"),
 ]
+
 interface_image = gr.Interface(
     fn=show_preds_image,
     inputs=inputs_image,
     outputs=outputs_image,
-    title="Pothole detector",
-    examples=path,
+    title="Drowsy Driver Detector - Image",
+    examples=image_examples,
     cache_examples=False,
 )
 
+
+# ==== 2. VIDEO ====
 def show_preds_video(video_path):
     cap = cv2.VideoCapture(video_path)
-    while(cap.isOpened()):
+
+    while cap.isOpened():
         ret, frame = cap.read()
-        if ret:
-            frame_copy = frame.copy()
-            outputs = model.predict(source=frame)
-            results = outputs[0].cpu().numpy()
-            for i, det in enumerate(results.boxes.xyxy):
-                cv2.rectangle(
-                    frame_copy,
-                    (int(det[0]), int(det[1])),
-                    (int(det[2]), int(det[3])),
-                    color=(0, 0, 255),
-                    thickness=2,
-                    lineType=cv2.LINE_AA
-                )
-            yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
+        if not ret:
+            break
 
-inputs_video = [
-    gr.components.Video(type="filepath", label="Input Video"),
+        frame_copy = frame.copy()
+
+        outputs = model.predict(source=frame, conf=0.25, verbose=False)
+        results = outputs[0].cpu().numpy()
+
+        frame_copy, _ = draw_and_decide_state(frame_copy, results)
 
+        yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
+
+    cap.release()
+
+
+inputs_video = [
+    gr.Video(type="filepath", label="Video đầu vào (driver camera)"),
 ]
 outputs_video = [
-    gr.components.Image(type="numpy", label="Output Image"),
+    gr.Image(type="numpy", label="Kết quả từng frame"),
 ]
+
 interface_video = gr.Interface(
     fn=show_preds_video,
     inputs=inputs_video,
     outputs=outputs_video,
-    title="Pothole detector",
-    examples=video_path,
+    title="Drowsy Driver Detector - Video",
+    examples=video_examples,
     cache_examples=False,
 )
 
+
+# ==== Giao diện Tab ====
 gr.TabbedInterface(
     [interface_image, interface_video],
-    tab_names=['Image inference', 'Video inference']
-).queue().launch()
+    tab_names=["Ảnh", "Video"],
+).queue().launch()