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

🚗 Nhận diện phương tiện giao thông (YOLOv8 + ByteTrack)

Thực hiện: Trần Hải Nam - 223332840

""", 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 < 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** (< 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()