Update app.py

app.py

@@ -1,3 +1,6 @@
+"""
+Đếm xe qua line - Deploy lên Hugging Face Spaces với Gradio
+"""
 import os
 import tempfile
 import cv2
@@ -6,39 +9,49 @@ 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
+
+# Load model
 model = YOLO("yolo11n.pt")
 CLASS_NAMES_DICT = model.model.names
 
-SELECTED_CLASS_NAMES = ['car', 'truck', 'bus', 'motorcycle']
+SELECTED_CLASS_NAMES = ['person', 'bus', 'motorcycle', 'car', 'truck']
 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)
-line_zone_annotator = sv.LineZoneAnnotator(
-    thickness=4, text_thickness=2, text_scale=2,
-    display_in_count=False, display_out_count=False
-)
-
+def process_video(
+    video_path,
+    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"
+):
+    """Xử lý video: đếm xe qua line và trả về video đã annotate.
 
-def process_video(video_path):
+    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)
-    w, h = video_info.width, video_info.height
-
-    line_y = int(h * 0.5)
-    line_start = sv.Point(int(w * 0.005), line_y)
-    line_end = sv.Point(int(w * 0.995), line_y)
 
-    line_zone = sv.LineZone(start=line_start, end=line_end)
     byte_tracker = sv.ByteTrack(
         track_activation_threshold=0.25,
         lost_track_buffer=30,
@@ -55,11 +68,100 @@ def process_video(video_path):
     def callback(frame: np.ndarray, index: int) -> np.ndarray:
         nonlocal previous_positions, class_counts, crossed_ids
 
-        results = model(frame, verbose=False)[0]
+        # Đả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:
+            max_size = int(max_frame_size)
+
+        if detect_every_n is None or detect_every_n < 1:
+            detect_every = 1
+        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))
+            if scale < 1.0:
+                frame_infer = cv2.resize(
+                    frame, (int(fw_orig * scale), int(fh_orig * scale))
+                )
+            else:
+                frame_infer = frame
+        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
+        else:
+            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,
+                    (int(fw * 0.05), line_pos),
+                    (int(fw * 0.95), line_pos),
+                    (0, 255, 255),
+                    2,
+                )
+            else:
+                cv2.line(
+                    annotator_frame,
+                    (line_pos, int(fh * 0.05)),
+                    (line_pos, int(fh * 0.95)),
+                    (0, 255, 255),
+                    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
+
+            overlay = annotator_frame.copy()
+            cv2.rectangle(overlay, (x0, y0), (x0 + box_w, y0 + box_h), (0, 0, 0), -1)
+            annotator_frame = cv2.addWeighted(overlay, 0.6, annotator_frame, 0.4, 0)
+
+            total = sum(class_counts.values())
+            cv2.putText(
+                annotator_frame,
+                f'Total: {total}',
+                (x0 + 10, y0 + 30),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.8,
+                (0, 255, 0),
+                2,
+            )
+            for i, cls_name in enumerate(SELECTED_CLASS_NAMES):
+                cnt = class_counts.get(cls_name, 0)
+                cv2.putText(
+                    annotator_frame,
+                    f'{cls_name.capitalize()}: {cnt}',
+                    (x0 + 10, y0 + 60 + i * 28),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.7,
+                    (255, 255, 255),
+                    2,
+                )
+
+            return annotator_frame
+
+        results = model(frame_infer, verbose=False)[0]
         detections = sv.Detections.from_ultralytics(results)
         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)):
@@ -69,15 +171,26 @@ def process_video(video_path):
                 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:
-                    py = previous_positions[tid]
-                    if py < line_y and cy > line_y and (tid, 'out') not in crossed_ids:
-                        crossed_ids.add((tid, 'out'))
+                    prev_coord = previous_positions[tid]
+                    if (
+                        prev_coord < line_pos
+                        and curr_coord > line_pos
+                        and (tid, "out") not in crossed_ids
+                    ):
+                        crossed_ids.add((tid, "out"))
                         class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
-                    elif py > line_y and cy < line_y and (tid, 'in') not in crossed_ids:
-                        crossed_ids.add((tid, 'in'))
+                    elif (
+                        prev_coord > line_pos
+                        and curr_coord < line_pos
+                        and (tid, "in") not in crossed_ids
+                    ):
+                        crossed_ids.add((tid, "in"))
                         class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
-                previous_positions[tid] = cy
+                previous_positions[tid] = curr_coord
 
         labels = [
             f"#{tid} {CLASS_NAMES_DICT[cid]} {conf:0.2f}"
@@ -86,15 +199,32 @@ def process_video(video_path):
             )
         ]
 
-        annotator_frame = frame.copy()
+        annotator_frame = frame_infer.copy()
         annotator_frame = trace_annotator.annotate(scene=annotator_frame, detections=detections)
         annotator_frame = box_annotator.annotate(scene=annotator_frame, detections=detections)
-        annotator_frame = label_annotator.annotate(scene=annotator_frame, detections=detections, labels=labels)
+        annotator_frame = label_annotator.annotate(
+            scene=annotator_frame, detections=detections, labels=labels
+        )
 
-        line_zone.trigger(detections)
-        annotator_frame = line_zone_annotator.annotate(annotator_frame, line_counter=line_zone)
+        # Vẽ line đếm theo cấu hình
+        if is_horizontal:
+            cv2.line(
+                annotator_frame,
+                (int(fw * 0.05), line_pos),
+                (int(fw * 0.95), line_pos),
+                (0, 255, 255),
+                2,
+            )
+        else:
+            cv2.line(
+                annotator_frame,
+                (line_pos, int(fh * 0.05)),
+                (line_pos, int(fh * 0.95)),
+                (0, 255, 255),
+                2,
+            )
 
-        fh, fw, _ = annotator_frame.shape
+        # Khung tổng đếm
         box_w, box_h = 280, 50 + len(SELECTED_CLASS_NAMES) * 28
         x0, y0 = fw - box_w - 20, 20
 
@@ -103,13 +233,26 @@ def process_video(video_path):
         annotator_frame = cv2.addWeighted(overlay, 0.6, annotator_frame, 0.4, 0)
 
         total = sum(class_counts.values())
-        cv2.putText(annotator_frame, f'Total: {total}', (x0 + 10, y0 + 30),
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
+        cv2.putText(
+            annotator_frame,
+            f'Total: {total}',
+            (x0 + 10, y0 + 30),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.8,
+            (0, 255, 0),
+            2,
+        )
         for i, cls_name in enumerate(SELECTED_CLASS_NAMES):
             cnt = class_counts.get(cls_name, 0)
-            cv2.putText(annotator_frame, f'{cls_name.capitalize()}: {cnt}',
-                        (x0 + 10, y0 + 60 + i * 28),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
+            cv2.putText(
+                annotator_frame,
+                f'{cls_name.capitalize()}: {cnt}',
+                (x0 + 10, y0 + 60 + i * 28),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.7,
+                (255, 255, 255),
+                2,
+            )
 
         return annotator_frame
 
@@ -122,22 +265,58 @@ def process_video(video_path):
     return output_path
 
 
-with gr.Blocks(title="Đếm xe", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚗 Đếm xe (YOLOv8 + ByteTrack)")
-    gr.Markdown("Upload video.")
+# Gradio UI
+with gr.Blocks(title="Đếm xe qua line", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🚗 Đếm xe qua line (YOLOv8 + ByteTrack)")
+    gr.Markdown(
+        "Upload video, hệ thống sẽ đếm person, car, bus, truck, motorcycle khi qua đường line giữa khung hình."
+    )
 
     with gr.Row():
         video_input = gr.Video(label="Video đầu vào")
         video_output = gr.Video(label="Video đã xử lý")
 
+    # Tùy chọn cấu hình
+    with gr.Accordion("⚙️ Tùy chọn nâng cao", open=False):
+        use_resize = gr.Checkbox(
+            value=True, label="Giảm kích thước khung hình trước khi detect"
+        )
+        max_frame_size = gr.Slider(
+            minimum=320,
+            maximum=1280,
+            value=DEFAULT_MAX_FRAME_SIZE,
+            step=64,
+            label="Kích thước tối đa (px)",
+        )
+        detect_every_n = gr.Slider(
+            minimum=1,
+            maximum=5,
+            value=DEFAULT_DETECT_EVERY_N_FRAMES,
+            step=1,
+            label="Detect mỗi N frame (1 = mọi frame)",
+        )
+        line_orientation = gr.Radio(
+            choices=["Ngang", "Dọc"],
+            value="Ngang",
+            label="Hướng line đếm",
+        )
+
     btn = gr.Button("▶️ Xử lý video")
-    btn.click(fn=process_video, inputs=video_input, outputs=video_output)
+    btn.click(
+        fn=process_video,
+        inputs=[video_input, use_resize, max_frame_size, detect_every_n, line_orientation],
+        outputs=video_output,
+    )
 
-    gr.Markdown("""
+    gr.Markdown(
+        """
     ### Lưu ý
+    - Mặc định line đếm nằm **ngang** ở giữa khung hình (50% chiều cao)
+    - Có thể chọn lại line **dọc** trong phần "Tùy chọn nâng cao"
     - Chạy trên CPU miễn phí → nên dùng video **ngắn** (< 30 giây) để tránh chờ lâu
-    - Model: YOLOv8n
-    """)
+    - Model: YOLOv8n (nhẹ cho CPU)
+    """
+    )
 
 if __name__ == "__main__":
     demo.launch()