Update app.py

app.py

@@ -1,223 +1,120 @@
 import os
-import shutil
-import subprocess
 import tempfile
+from typing import Dict, Tuple
 
 import cv2
-import imageio
-import numpy as np
 import gradio as gr
+import numpy as np
 import supervision as sv
 from ultralytics import YOLO
 
-DEFAULT_MAX_FRAME_SIZE = 480
-DEFAULT_DETECT_EVERY_N_FRAMES = 3
-DEFAULT_ZONE_MARGIN = 0.10
 
-model = YOLO("yolov8n.pt")
-CLASS_NAMES_DICT = model.model.names
+MODEL_PATH = "yolo11n.pt"
+model = YOLO(MODEL_PATH)
 
-SELECTED_CLASS_NAMES = ['car', 'truck', 'bus', 'motorcycle', ]
+CLASS_NAMES_DICT = model.model.names
+SELECTED_CLASS_NAMES = ["car", "bus", "truck", "motorcycle"]
 SELECTED_CLASS_IDS = [
-    {value: key for key, value in CLASS_NAMES_DICT.items()}[name]
-    for name in SELECTED_CLASS_NAMES
+    {value: key for key, value in CLASS_NAMES_DICT.items()}[class_name]
+    for class_name in SELECTED_CLASS_NAMES
 ]
 
-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)
+
 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",
-    zone_margin: float = DEFAULT_ZONE_MARGIN,
-):
-    if video_path is None:
-        return None
-
-    if isinstance(video_path, dict):
-        video_path = video_path.get("path", video_path)
+    video_path: str,
+    frame_stride: int = 2,
+    max_seconds: int = 60,
+    line_orientation: str = "horizontal",
+) -> Tuple[str, Dict[str, int]]:
+    """
+    Xử lý video: phát hiện + tracking + đếm phương tiện.
+    Trả về: đường dẫn video đã annotate và dict số lượng theo class.
+    """
+    video_info = sv.VideoInfo.from_video_path(video_path)
+    fps = video_info.fps
+    max_frames_for_detection = int(fps * max_seconds) if max_seconds is not None else None
 
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
-        return None
+        raise RuntimeError("Không thể mở video đầu vào.")
+
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
 
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    if fps is None or fps <= 0 or np.isnan(fps):
-        fps = 30
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    tmp_dir = tempfile.mkdtemp()
+    out_path = os.path.join(tmp_dir, "result.mp4")
+    writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
 
     byte_tracker = sv.ByteTrack(
         track_activation_threshold=0.25,
         lost_track_buffer=30,
         minimum_matching_threshold=0.8,
         frame_rate=fps,
-        minimum_consecutive_frames=3
+        minimum_consecutive_frames=3,
     )
-    byte_tracker.reset()
-
-    class_counts = {name: 0 for name in SELECTED_CLASS_NAMES}
-    counted_ids = set()
 
-    def callback(frame: np.ndarray, index: int) -> np.ndarray:
-        nonlocal class_counts, counted_ids
-
-        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)
-
-        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
+    if line_orientation == "vertical":
+        line_pos = int(width * 0.5)
+        line_start = sv.Point(line_pos, int(height * 0.05))
+        line_end = sv.Point(line_pos, int(height * 0.95))
+        axis = "x"
+    else:
+        line_pos = int(height * 0.5)
+        line_start = sv.Point(int(width * 0.05), line_pos)
+        line_end = sv.Point(int(width * 0.95), line_pos)
+        axis = "y"
+
+    line_zone = sv.LineZone(start=line_start, end=line_end)
+
+    previous_positions: Dict[int, float] = {}
+    class_counts: Dict[str, int] = {name: 0 for name in SELECTED_CLASS_NAMES}
+    crossed_ids = set()
+
+    box_annotator = sv.BoxAnnotator(thickness=4)
+    label_annotator = sv.LabelAnnotator(
+        text_thickness=2, text_scale=1.0, text_color=sv.Color.BLACK
+    )
+    trace_annotator = sv.TraceAnnotator(thickness=4, trace_length=50)
+    line_zone_annotator = sv.LineZoneAnnotator(
+        thickness=4,
+        color=sv.Color.RED,
+        text_thickness=2,
+        text_scale=2,
+        display_in_count=False,
+        display_out_count=False,
+    )
 
-        fh, fw = frame_infer.shape[:2]
-        if line_orientation == "Dọc":
-            line_pos = int(fw * 0.5)
-            is_horizontal = False
-        else:
-            line_pos = int(fh * 0.5)
-            is_horizontal = True
+    frame_idx = 0
+    last_detections = None
 
-        if zone_margin is None or zone_margin <= 0:
-            zm_ratio = DEFAULT_ZONE_MARGIN
-        else:
-            zm_ratio = max(0.01, min(0.5, float(zone_margin)))
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
 
-        if is_horizontal:
-            z_half = int(fh * zm_ratio)
-            z_top = max(0, line_pos - z_half)
-            z_bot = min(fh - 1, line_pos + z_half)
+        annotated_frame = frame.copy()
+
+        can_run_detection = True
+        if max_frames_for_detection is not None and frame_idx >= max_frames_for_detection:
+            can_run_detection = False
+
+        if can_run_detection and frame_idx % frame_stride == 0:
+            results = model(
+                frame,
+                imgsz=640,
+                device="cpu",
+                verbose=False,
+                classes=SELECTED_CLASS_IDS,
+            )[0]
+            detections = sv.Detections.from_ultralytics(results)
+            detections = byte_tracker.update_with_detections(detections)
+            last_detections = detections
         else:
-            z_half = int(fw * zm_ratio)
-            z_left = max(0, line_pos - z_half)
-            z_right = min(fw - 1, line_pos + z_half)
-
-        if detect_every > 1 and index % detect_every != 0:
-            annotator_frame = frame_infer.copy()
-            overlay_zone = annotator_frame.copy()
-            if is_horizontal:
-                cv2.rectangle(
-                    overlay_zone,
-                    (0, z_top),
-                    (fw, z_bot),
-                    (0, 0, 200),
-                    -1,
-                )
-            else:
-                cv2.rectangle(
-                    overlay_zone,
-                    (z_left, 0),
-                    (z_right, fh),
-                    (0, 0, 200),
-                    -1,
-                )
-            annotator_frame = cv2.addWeighted(
-                overlay_zone, 0.18, annotator_frame, 0.82, 0
-            )
-
-            thickness_base = max(2, int(2 * (max(fw, fh) / 1920)))
-            if is_horizontal:
-                cv2.line(
-                    annotator_frame,
-                    (0, z_top),
-                    (fw, z_top),
-                    (0, 100, 255),
-                    thickness_base,
-                )
-                cv2.line(
-                    annotator_frame,
-                    (0, z_bot),
-                    (fw, z_bot),
-                    (0, 100, 255),
-                    thickness_base,
-                )
-                cv2.line(
-                    annotator_frame,
-                    (0, line_pos),
-                    (fw, line_pos),
-                    (0, 0, 255),
-                    max(3, thickness_base + 1),
-                )
-            else:
-                cv2.line(
-                    annotator_frame,
-                    (z_left, 0),
-                    (z_left, fh),
-                    (0, 100, 255),
-                    thickness_base,
-                )
-                cv2.line(
-                    annotator_frame,
-                    (z_right, 0),
-                    (z_right, fh),
-                    (0, 100, 255),
-                    thickness_base,
-                )
-                cv2.line(
-                    annotator_frame,
-                    (line_pos, 0),
-                    (line_pos, fh),
-                    (0, 0, 255),
-                    max(3, thickness_base + 1),
-                )
-
-            # Kích thước khung thống kê tỉ lệ theo kích thước khung hình
-            scale_ui = max(fw, fh) / 1280.0
-            base_box_w = 260
-            base_box_h = 60 + len(SELECTED_CLASS_NAMES) * 26
-            box_w = int(base_box_w * scale_ui)
-            box_h = int(base_box_h * scale_ui)
-            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
+            detections = last_detections
 
-        results = model(frame_infer, verbose=False, imgsz=480)[0]
-        detections = sv.Detections.from_ultralytics(results)
-        detections = detections[np.isin(detections.class_id, SELECTED_CLASS_IDS)]
-        detections = byte_tracker.update_with_detections(detections)
-
-        if detections.tracker_id is not None:
+        if detections is not None and detections.tracker_id is not None:
             xyxy = detections.xyxy
             for i in range(len(detections)):
                 tid = int(detections.tracker_id[i])
@@ -226,111 +123,60 @@ def process_video(
                 cx = (xyxy[i, 0] + xyxy[i, 2]) / 2
                 cy = (xyxy[i, 1] + xyxy[i, 3]) / 2
 
-                if cls_name in SELECTED_CLASS_NAMES and tid not in counted_ids:
-                    if is_horizontal and z_top <= cy <= z_bot:
+                cur_pos = cy if axis == "y" else cx
+
+                if tid in previous_positions:
+                    prev_pos = previous_positions[tid]
+                    if (
+                        prev_pos < line_pos
+                        and cur_pos > line_pos
+                        and (tid, "pos") not in crossed_ids
+                    ):
+                        crossed_ids.add((tid, "pos"))
                         class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
-                        counted_ids.add(tid)
-                    elif (not is_horizontal) and z_left <= cx <= z_right:
+                    elif (
+                        prev_pos > line_pos
+                        and cur_pos < line_pos
+                        and (tid, "neg") not in crossed_ids
+                    ):
+                        crossed_ids.add((tid, "neg"))
                         class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
-                        counted_ids.add(tid)
-
-        labels = [
-            f"#{tid} {CLASS_NAMES_DICT[cid]} {conf:0.2f}"
-            for conf, cid, tid in zip(
-                detections.confidence, detections.class_id, detections.tracker_id
-            )
-        ]
-
-        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
-        )
+                previous_positions[tid] = cur_pos
 
-        overlay_zone = annotator_frame.copy()
-        if is_horizontal:
-            cv2.rectangle(
-                overlay_zone,
-                (0, z_top),
-                (fw, z_bot),
-                (0, 0, 200),
-                -1,
-            )
-        else:
-            cv2.rectangle(
-                overlay_zone,
-                (z_left, 0),
-                (z_right, fh),
-                (0, 0, 200),
-                -1,
-            )
-        annotator_frame = cv2.addWeighted(
-            overlay_zone, 0.18, annotator_frame, 0.82, 0
-        )
+            labels = [
+                f"#{tracker_id} {CLASS_NAMES_DICT[class_id]} {confidence:0.2f}"
+                for confidence, class_id, tracker_id in zip(
+                    detections.confidence, detections.class_id, detections.tracker_id
+                )
+            ]
 
-        thickness_base = max(2, int(2 * (max(fw, fh) / 1920)))
-        if is_horizontal:
-            cv2.line(
-                annotator_frame,
-                (0, z_top),
-                (fw, z_top),
-                (0, 100, 255),
-                thickness_base,
+            annotated_frame = trace_annotator.annotate(
+                scene=annotated_frame, detections=detections
             )
-            cv2.line(
-                annotator_frame,
-                (0, z_bot),
-                (fw, z_bot),
-                (0, 100, 255),
-                thickness_base,
+            annotated_frame = box_annotator.annotate(
+                scene=annotated_frame, detections=detections
             )
-            cv2.line(
-                annotator_frame,
-                (0, line_pos),
-                (fw, line_pos),
-                (0, 0, 255),
-                max(3, thickness_base + 1),
+            annotated_frame = label_annotator.annotate(
+                scene=annotated_frame, detections=detections, labels=labels
             )
-        else:
-            cv2.line(
-                annotator_frame,
-                (z_left, 0),
-                (z_left, fh),
-                (0, 100, 255),
-                thickness_base,
-            )
-            cv2.line(
-                annotator_frame,
-                (z_right, 0),
-                (z_right, fh),
-                (0, 100, 255),
-                thickness_base,
-            )
-            cv2.line(
-                annotator_frame,
-                (line_pos, 0),
-                (line_pos, fh),
-                (0, 0, 255),
-                max(3, thickness_base + 1),
+
+            line_zone.trigger(detections)
+            annotated_frame = line_zone_annotator.annotate(
+                annotated_frame, line_counter=line_zone
             )
 
-        # Kích thước khung thống kê tỉ lệ theo kích thước khung hình
-        scale_ui = max(fw, fh) / 1280.0
-        base_box_w = 260
-        base_box_h = 60 + len(SELECTED_CLASS_NAMES) * 26
-        box_w = int(base_box_w * scale_ui)
-        box_h = int(base_box_h * scale_ui)
-        x0, y0 = fw - box_w - 20, 20
+        h, w, _ = annotated_frame.shape
+        box_w, box_h = 280, 50 + len(SELECTED_CLASS_NAMES) * 28
+        x0, y0 = w - box_w - 20, 20
 
-        overlay = annotator_frame.copy()
+        overlay = annotated_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)
+        annotated_frame = cv2.addWeighted(overlay, 0.6, annotated_frame, 0.4, 0)
 
         total = sum(class_counts.values())
         cv2.putText(
-            annotator_frame,
-            f'Total: {total}',
+            annotated_frame,
+            f"Total: {total}",
             (x0 + 10, y0 + 30),
             cv2.FONT_HERSHEY_SIMPLEX,
             0.8,
@@ -340,8 +186,8 @@ def process_video(
         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}',
+                annotated_frame,
+                f"{cls_name.capitalize()}: {cnt}",
                 (x0 + 10, y0 + 60 + i * 28),
                 cv2.FONT_HERSHEY_SIMPLEX,
                 0.7,
@@ -349,112 +195,97 @@ def process_video(
                 2,
             )
 
-        return annotator_frame
-
-    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-    writer = imageio.get_writer(output_path, fps=fps)
-
-    index = 0
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-
-        annotated_frame = callback(frame, index)
-        index += 1
-
-        # imageio expects RGB frames
-        annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
-        writer.append_data(annotated_frame_rgb)
+        writer.write(annotated_frame)
+        frame_idx += 1
 
-    writer.close()
     cap.release()
+    writer.release()
+
+    return out_path, class_counts
+
+
+def gradio_infer(video, mode: str = "Cân bằng", line_type: str = "Ngang"):
+    if video is None:
+        return None, "Vui lòng upload 1 video."
+
+    if isinstance(video, dict):
+        video_path = video.get("name") or video.get("data")
+    else:
+        video_path = video
+
+    if not video_path or not os.path.exists(video_path):
+        return None, "Không tìm thấy file video."
+
+    if mode == "Nhanh":
+        frame_stride = 4
+        max_seconds = 30
+    elif mode == "Chính xác":
+        frame_stride = 1
+        max_seconds = 90
+    else:  # Cân bằng
+        frame_stride = 2
+        max_seconds = 60
+
+    if line_type == "Dọc":
+        line_orientation = "vertical"
+    else:
+        line_orientation = "horizontal"
+
+    out_path, counts = process_video(
+        video_path=video_path,
+        frame_stride=frame_stride,
+        max_seconds=max_seconds,
+        line_orientation=line_orientation,
+    )
 
-    return output_path
+    total = sum(counts.values())
+    lines = [f"Tổng số phương tiện: {total}"]
+    for cls_name in SELECTED_CLASS_NAMES:
+        lines.append(f"- {cls_name}: {counts.get(cls_name, 0)}")
+    summary = "\n".join(lines)
 
+    return out_path, summary
 
-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;">🚗 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>
-            </div>
-            """,
-            elem_id="header",
-        )
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown(
-                "### 🎥 Video đầu vào\n"
-                "Upload video ngắn (ưu tiên &lt; 30s để xử lý nhanh hơn)."
-            )
-            video_input = gr.Video(label="Video đầu vào")
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # Nhận diện phương tiện trong video (YOLO + ByteTrack)
 
-        with gr.Column(scale=1):
-            gr.Markdown(
-                "### ✅ Kết quả đã xử lý\n"
-                "Hiển thị và thống kê số lượng theo lớp."
-            )
-            video_output = gr.Video(label="Video đã xử lý", format="mp4")
+        Upload 1 video.
+        """
+    )
 
-    # Tùy chọn cấu hình
-    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 nhận dạng"
+    with gr.Row():
+        with gr.Column():
+            video_input = gr.Video(label="Video đầu vào", sources=["upload"])
+            mode_input = gr.Radio(
+                ["Nhanh", "Cân bằng", "Chính xác"],
+                value="Cân bằng",
+                label="Chế độ xử lý",
+                info="Nhanh: nhanh hơn, ít chính xác hơn. Chính xác: chậm hơn.",
             )
-            line_orientation = gr.Radio(
-                choices=["Ngang", "Dọc"],
+            line_type_input = gr.Radio(
+                ["Ngang", "Dọc"],
                 value="Ngang",
-                label="Hướng phương tiện di chuyển",
-            )
-        with gr.Row():
-            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)",
+                label="Hướng đường đếm",
+                info="Ngang: đường ngang ở giữa khung hình. Dọc: đường dọc ở giữa khung hình.",
             )
-            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)",
-            )
-            zone_margin = gr.Slider(
-                minimum=0.02,
-                maximum=0.30,
-                value=DEFAULT_ZONE_MARGIN,
-                step=0.01,
-                label="Độ dày vùng đếm quanh line",
+            run_btn = gr.Button("Bắt đầu đếm")
+
+        with gr.Column():
+            video_output = gr.Video(label="Video đã xử lý")
+            text_output = gr.Textbox(
+                label="Kết quả đếm", lines=6, interactive=False
             )
 
-    btn = gr.Button("▶️ Xử lý video", variant="primary")
-    btn.click(
-        fn=process_video,
-        inputs=[video_input, use_resize, max_frame_size, detect_every_n, line_orientation, zone_margin],
-        outputs=video_output,
+    run_btn.click(
+        fn=gradio_infer,
+        inputs=[video_input, mode_input, line_type_input],
+        outputs=[video_output, text_output],
     )
 
-    gr.Markdown(
-        """
-    ---
-    ### ℹ️ Gợi ý sử dụng
-    - 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**.
-    """
-    )
 
 if __name__ == "__main__":
     demo.launch()
+