File size: 4,363 Bytes
ce2a8be | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | import gradio as gr
import cv2
from ultralytics import YOLO
import subprocess
import os
# Tải mô hình YOLOv8
model = YOLO('yolov8n.pt')
def process_video(video_path):
if not video_path:
return None
# Cấu hình đếm
vehicle_classes = ['car', 'motorcycle', 'bus', 'truck']
vehicle_class_ids = [2, 3, 5, 7]
vehicle_count = {cls: 0 for cls in vehicle_classes}
track_history = {}
counted_ids = set()
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# Vạch đếm ở 60% màn hình
line_y = int(height * 0.6)
# Xuất ra file tạm bằng OpenCV (mp4v)
temp_output = 'temp_output.mp4'
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(temp_output, fourcc, fps, (width, height))
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
results = model.track(frame, classes=vehicle_class_ids, persist=True, verbose=False)
cv2.line(frame, (0, line_y), (width, line_y), (0, 0, 255), 3)
cv2.putText(frame, "Counting Line", (10, line_y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
if results[0].boxes.id is not None:
boxes = results[0].boxes
ids = boxes.id.cpu().numpy().astype(int)
classes = boxes.cls.cpu().numpy().astype(int)
confs = boxes.conf.cpu().numpy()
for (x1, y1, x2, y2, obj_id, cls_id, conf) in zip(
boxes.xyxy[:,0], boxes.xyxy[:,1], boxes.xyxy[:,2], boxes.xyxy[:,3],
ids, classes, confs
):
label = model.names[cls_id]
cx = int((x1 + x2) / 2)
cy = int((y1 + y2) / 2)
cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (255, 255, 0), 2)
cv2.circle(frame, (cx, cy), 5, (0, 255, 0), -1)
cv2.putText(frame, f"ID:{obj_id} {label}", (int(x1), int(y1) - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)
# Logic đếm qua vạch
if obj_id in track_history and obj_id not in counted_ids:
prev_cy = track_history[obj_id]
if (prev_cy < line_y and cy >= line_y) or (prev_cy > line_y and cy <= line_y):
vehicle_count[label] += 1
counted_ids.add(obj_id)
cv2.line(frame, (0, line_y), (width, line_y), (0, 255, 0), 5)
track_history[obj_id] = cy
# Hiển thị kết quả
total = sum(vehicle_count.values())
overlay = frame.copy()
box_w, box_h = 250, 160
x0, y0 = width - box_w - 20, 20
cv2.rectangle(overlay, (x0, y0), (x0 + box_w, y0 + box_h), (0, 0, 0), -1)
frame = cv2.addWeighted(overlay, 0.6, frame, 0.4, 0)
cv2.putText(frame, f'Total Objects: {total}', (x0 + 10, y0 + 30),
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
for i, (cls, cnt) in enumerate(vehicle_count.items()):
cv2.putText(frame, f'{cls.capitalize()}: {cnt}',
(x0 + 10, y0 + 60 + i * 25),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
out.write(frame)
cap.release()
out.release()
# CHUYỂN ĐỔI VIDEO SANG CHUẨN H.264 ĐỂ HIỂN THỊ TRÊN WEB (Rất quan trọng)
final_output = 'result_video.mp4'
if os.path.exists(final_output):
os.remove(final_output)
subprocess.run(
['ffmpeg', '-i', temp_output, '-vcodec', 'libx264', '-crf', '28', final_output],
check=True
)
return final_output
# Xây dựng giao diện Web với Gradio
interface = gr.Interface(
fn=process_video,
inputs=gr.Video(label="Tải video của bạn lên đây"),
outputs=gr.Video(label="Video kết quả đếm xe"),
title="Hệ thống đếm xe tự động (YOLOv8)",
description="Tải lên một video giao thông để hệ thống tự động nhận diện và đếm số lượng xe cộ qua vạch kẻ ngang.",
allow_flagging="never"
)
# Chạy ứng dụng
if __name__ == "__main__":
interface.launch()
|