Update app.py

app.py

@@ -1,6 +1,3 @@
-"""
-Đếm xe qua line - Deploy lên Hugging Face Spaces với Gradio
-"""
 import os
 import tempfile
 import cv2
@@ -9,21 +6,18 @@ import gradio as gr
 import supervision as sv
 from ultralytics import YOLO
 
-# Cấu hình mặc định cho tối ưu
-DEFAULT_MAX_FRAME_SIZE = 640        # Giới hạn kích thước khung hình đưa vào YOLO
-DEFAULT_DETECT_EVERY_N_FRAMES = 2   # Chỉ detect mỗi N frame để giảm tải CPU
+DEFAULT_MAX_FRAME_SIZE = 640
+DEFAULT_DETECT_EVERY_N_FRAMES = 2
 
-# Load model (yolov8n nhẹ cho CPU miễn phí)
 model = YOLO("yolov8n.pt")
 CLASS_NAMES_DICT = model.model.names
 
-SELECTED_CLASS_NAMES = ['person', 'bus', 'motorcycle', 'car', 'truck']
+SELECTED_CLASS_NAMES = ['car', 'truck', 'bus', 'motorcycle', ]
 SELECTED_CLASS_IDS = [
     {value: key for key, value in CLASS_NAMES_DICT.items()}[name]
     for name in SELECTED_CLASS_NAMES
 ]
 
-# Annotators (tạo 1 lần, dùng lại)
 box_annotator = sv.BoxAnnotator(thickness=4)
 label_annotator = sv.LabelAnnotator(text_thickness=2, text_scale=1.5, text_color=sv.Color.BLACK)
 trace_annotator = sv.TraceAnnotator(thickness=4, trace_length=50)
@@ -32,24 +26,14 @@ def process_video(
     use_resize: bool = True,
     max_frame_size: int = DEFAULT_MAX_FRAME_SIZE,
     detect_every_n: int = DEFAULT_DETECT_EVERY_N_FRAMES,
-    line_orientation: str = "Ngang",  # "Ngang" | "Dọc"
+    line_orientation: str = "Ngang",
 ):
-    """Xử lý video: đếm xe qua line và trả về video đã annotate.
-
-    Các tham số tối ưu và cấu hình được truyền từ UI:
-    - use_resize: có resize khung hình trước khi detect hay không
-    - max_frame_size: cạnh dài tối đa sau resize
-    - detect_every_n: chỉ detect mỗi N frame (1 = detect mọi frame)
-    - line_orientation: "Ngang" hoặc "Dọc"
-    """
     if video_path is None:
         return None
 
-    # Gradio Video có thể trả về dict với key "path"
     if isinstance(video_path, dict):
         video_path = video_path.get("path", video_path)
 
-    # Lấy thông tin video (dùng cho tracker)
     video_info = sv.VideoInfo.from_video_path(video_path)
 
     byte_tracker = sv.ByteTrack(
@@ -68,7 +52,6 @@ def process_video(
     def callback(frame: np.ndarray, index: int) -> np.ndarray:
         nonlocal previous_positions, class_counts, crossed_ids
 
-        # Đảm bảo giá trị hợp lệ
         if max_frame_size is None or max_frame_size <= 0:
             max_size = DEFAULT_MAX_FRAME_SIZE
         else:
@@ -79,7 +62,6 @@ def process_video(
         else:
             detect_every = int(detect_every_n)
 
-        # Resize khung hình trước khi đưa vào YOLO để giảm tải (nếu bật)
         fh_orig, fw_orig = frame.shape[:2]
         if use_resize:
             scale = min(1.0, max_size / max(fh_orig, fw_orig))
@@ -92,9 +74,7 @@ def process_video(
         else:
             frame_infer = frame
 
-        # Tính lại kích thước sau khi resize
         fh, fw = frame_infer.shape[:2]
-        # Vị trí line theo hướng người dùng chọn
         if line_orientation == "Dọc":
             line_pos = int(fw * 0.5)
             is_horizontal = False
@@ -102,11 +82,9 @@ def process_video(
             line_pos = int(fh * 0.5)
             is_horizontal = True
 
-        # Bỏ qua một số frame để giảm số lần detect
         if detect_every > 1 and index % detect_every != 0:
             annotator_frame = frame_infer.copy()
 
-            # Vẽ line đếm theo cấu hình
             if is_horizontal:
                 cv2.line(
                     annotator_frame,
@@ -124,7 +102,6 @@ def process_video(
                     2,
                 )
 
-            # Khung tổng đếm (dùng class_counts hiện tại)
             box_w, box_h = 280, 50 + len(SELECTED_CLASS_NAMES) * 28
             x0, y0 = fw - box_w - 20, 20
 
@@ -161,7 +138,6 @@ def process_video(
         detections = detections[np.isin(detections.class_id, SELECTED_CLASS_IDS)]
         detections = byte_tracker.update_with_detections(detections)
 
-        # Đếm theo loại khi qua line
         if detections.tracker_id is not None:
             xyxy = detections.xyxy
             for i in range(len(detections)):
@@ -171,7 +147,6 @@ def process_video(
                 cx = (xyxy[i, 0] + xyxy[i, 2]) / 2
                 cy = (xyxy[i, 1] + xyxy[i, 3]) / 2
 
-                # Toạ độ 1D dùng để kiểm tra qua line (y nếu line ngang, x nếu line dọc)
                 curr_coord = cy if is_horizontal else cx
 
                 if tid in previous_positions:
@@ -206,7 +181,6 @@ def process_video(
             scene=annotator_frame, detections=detections, labels=labels
         )
 
-        # Vẽ line đếm theo cấu hình
         if is_horizontal:
             cv2.line(
                 annotator_frame,
@@ -224,7 +198,6 @@ def process_video(
                 2,
             )
 
-        # Khung tổng đếm
         box_w, box_h = 280, 50 + len(SELECTED_CLASS_NAMES) * 28
         x0, y0 = fw - box_w - 20, 20
 
@@ -265,17 +238,12 @@ def process_video(
     return output_path
 
 
-# Gradio UI
-with gr.Blocks(title="Đếm xe qua line", theme=gr.themes.Soft(primary_hue="blue", secondary_hue="gray")) as demo:
-    # Header đẹp hơn + thông tin sinh viên
+with gr.Blocks(title="Nhận dạng phương tiện giao thông", theme=gr.themes.Soft(primary_hue="blue", secondary_hue="gray")) as demo:
     with gr.Row():
         gr.Markdown(
             """
             <div style="display:flex;flex-direction:column;gap:4px;">
-              <h1 style="margin-bottom:4px;">🚗 Đếm xe qua line (YOLOv8 + ByteTrack)</h1>
-              <p style="margin:0;font-size:16px;color:#4b5563;">
-                Dự án: Đếm đối tượng qua line (Horizonal / Vertical) dùng YOLOv8 + ByteTrack.
-              </p>
+              <h1 style="margin-bottom:4px;">🚗 Nhận diện phương tiện giao thông (YOLOv8 + ByteTrack)</h1>
               <p style="margin:0;font-size:14px;color:#6b7280;">
                 Thực hiện: <strong>Trần Hải Nam - 223332840</strong>
               </p>
@@ -295,7 +263,7 @@ with gr.Blocks(title="Đếm xe qua line", theme=gr.themes.Soft(primary_hue="blu
         with gr.Column(scale=1):
             gr.Markdown(
                 "### ✅ Kết quả đã xử lý\n"
-                "Hiển thị tracking, line đếm và thống kê số lượng theo lớp."
+                "Hiển thị và thống kê số lượng theo lớp."
             )
             video_output = gr.Video(label="Video đã xử lý")
 
@@ -303,12 +271,12 @@ with gr.Blocks(title="Đếm xe qua line", theme=gr.themes.Soft(primary_hue="blu
     with gr.Accordion("⚙️ Tùy chọn nâng cao", open=False):
         with gr.Row():
             use_resize = gr.Checkbox(
-                value=True, label="Giảm kích thước khung hình trước khi detect"
+                value=True, label="Giảm kích thước khung hình trước khi nhận dạng"
             )
             line_orientation = gr.Radio(
                 choices=["Ngang", "Dọc"],
                 value="Ngang",
-                label="Hướng line đếm",
+                label="Hướng phương tiện di chuyển",
             )
         with gr.Row():
             max_frame_size = gr.Slider(
@@ -337,12 +305,12 @@ with gr.Blocks(title="Đếm xe qua line", theme=gr.themes.Soft(primary_hue="blu
         """
     ---
     ### ℹ️ Gợi ý sử dụng
-    - Mặc định line đếm nằm **ngang** ở giữa khung hình (50% chiều cao).
-    - Có thể chuyển sang line **dọc** trong phần _"Tùy chọn nâng cao"_.
-    - Với CPU miễn phí (Hugging Face Spaces), nên:
+    - Mặc định hướng phương tiện di chuyển để nhận dạng là nằm **ngang** ở giữa khung hình (50% chiều cao).
+    - Có thể chuyển sang hướng **dọc** trong phần _"Tùy chọn nâng cao"_.
+    - Vì sử dụng CPU, nên:
       - Dùng video **ngắn** (&lt; 30 giây).
       - Tăng `Detect mỗi N frame` nếu muốn xử lý nhanh hơn.
-    - Model sử dụng: **YOLOv8n** (phiên bản nhẹ, phù hợp CPU).
+    - Model sử dụng: **YOLOv8n**.
     """
     )