app.py

import os
import tempfile
import cv2
import numpy as np
import gradio as gr
import supervision as sv
from ultralytics import YOLO

DEFAULT_MAX_FRAME_SIZE = 640
DEFAULT_DETECT_EVERY_N_FRAMES = 2
DEFAULT_ZONE_MARGIN = 0.10

model = YOLO("yolov8n.pt")
CLASS_NAMES_DICT = model.model.names

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
]

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_info = sv.VideoInfo.from_video_path(video_path)

    byte_tracker = sv.ByteTrack(
        track_activation_threshold=0.25,
        lost_track_buffer=30,
        minimum_matching_threshold=0.8,
        frame_rate=video_info.fps or 30,
        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

        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

        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)))

        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)
        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),
                )

            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)

        if detections.tracker_id is not None:
            xyxy = detections.xyxy
            for i in range(len(detections)):
                tid = int(detections.tracker_id[i])
                cls_id = int(detections.class_id[i])
                cls_name = CLASS_NAMES_DICT[cls_id]
                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:
                        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:
                        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
        )

        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),
            )

        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

    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    sv.process_video(
        source_path=video_path,
        target_path=output_path,
        callback=callback
    )
    return output_path


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.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ý")

    # 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"
            )
            line_orientation = gr.Radio(
                choices=["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)",
            )
            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",
            )

    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,
    )

    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()