Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import cv2 | |
| from ultralytics import YOLO | |
| import os | |
| from moviepy.editor import VideoFileClip | |
| # 1. Khởi tạo mô hình | |
| model = YOLO('yolov8n.pt') | |
| def process_video(video_path): | |
| if not video_path: | |
| return None | |
| # --- CẤU HÌNH BIẾN ĐẾM (Dựa trên code gốc của bạn) --- | |
| 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() | |
| # Đọc video đầu vào | |
| 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% chiều cao màn hình | |
| line_y = int(height * 0.6) | |
| # File tạm trung gian (định dạng mp4v) | |
| temp_output = 'temp_raw.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 | |
| # Thực hiện Tracking | |
| results = model.track(frame, classes=vehicle_class_ids, persist=True, verbose=False) | |
| # Vẽ vạch đếm màu đỏ | |
| 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) | |
| for (x1, y1, x2, y2, obj_id, cls_id) in zip( | |
| boxes.xyxy[:,0], boxes.xyxy[:,1], boxes.xyxy[:,2], boxes.xyxy[:,3], ids, classes | |
| ): | |
| label = model.names[cls_id] | |
| cx = int((x1 + x2) / 2) | |
| cy = int((y1 + y2) / 2) | |
| # Vẽ khung và tâm xe | |
| 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) | |
| # 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 | |
| # --- BẢNG THỐNG KÊ (Thiết kế gốc của bạn) --- | |
| 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: {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 SANG CHUẨN H.264 ĐỂ XEM TRỰC TIẾP --- | |
| final_output = 'result_display.mp4' | |
| try: | |
| # Chuyển đổi định dạng để trình duyệt web có thể Play trực tiếp | |
| clip = VideoFileClip(temp_output) | |
| clip.write_videofile(final_output, codec="libx264", audio=False) | |
| clip.close() | |
| if os.path.exists(temp_output): | |
| os.remove(temp_output) | |
| except Exception as e: | |
| print(f"Lỗi encode video: {e}") | |
| return temp_output | |
| return final_output | |
| # --- GIAO DIỆN GRADIO --- | |
| interface = gr.Interface( | |
| fn=process_video, | |
| inputs=gr.Video(label="Tải video của bạn lên"), | |
| outputs=gr.Video(label="Kết quả đếm xe (Xem trực tiếp)"), | |
| title="🚗 Hệ thống Đếm Xe Thông Minh (YOLOv8)", | |
| description="Nhận diện và đếm số lượng xe qua vạch. Kết quả xem được trực tiếp trên web.", | |
| ) | |
| if __name__ == "__main__": | |
| interface.launch() | |