Untitled

import cv2
import numpy as np
from openvino.runtime import Core

# Load mô hình OpenVINO
ie = Core()
model_path = "openvino_model.xml"  # Đảm bảo đúng tên file XML đã lưu
compiled_model = ie.compile_model(model_path, "CPU")

# Đọc video đầu vào
video_path = r"C:\Users\Huy\Pictures\video1.mp4"  # Đường dẫn video của bạn
cap = cv2.VideoCapture(video_path)

# Kiểm tra video có mở được không
if not cap.isOpened():
    print(f"Không thể mở video: {video_path}")
    exit()

# Lấy thông tin video
fps = cap.get(cv2.CAP_PROP_FPS)
delay = int(1000/fps)  # Tính toán độ trễ giữa các frame

# Lấy thông tin kích thước đầu vào của mô hình
input_blob = compiled_model.input(0).any_name
input_shape = compiled_model.input(0).shape  # (1, 3, 640, 640)
input_size = (input_shape[2], input_shape[3])  # (640, 640)

# Danh sách các class của YOLOv8 (80 class COCO)
class_names = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
               "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
               "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
               "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
               "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
               "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
               "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake",
               "chair", "couch", "potted plant", "bed", "dining table", "toilet", "TV", "laptop",
               "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster",
               "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
               "hair drier", "toothbrush"]

# Đọc từng frame của video
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        print("Không thể đọc tiếp frame. Thoát...")
        break

    try:
        # Resize & Chuẩn hóa ảnh về [0,1]
        resized_frame = cv2.resize(frame, input_size)
        normalized_frame = resized_frame / 255.0  # Chuẩn hóa pixel
        input_tensor = np.expand_dims(normalized_frame.transpose(2, 0, 1), axis=0).astype(np.float32)

        # Chạy mô hình OpenVINO
        results = compiled_model([input_tensor])[compiled_model.output(0).any_name]

        # Xử lý đầu ra YOLOv8
        boxes, scores, class_ids = [], [], []
        conf_threshold = 0.4
        img_h, img_w, _ = frame.shape

        for i in range(results.shape[2]):  # Duyệt qua N box dự đoán
            score = results[0, 4, i]  # Confidence score
            if score > conf_threshold:
                x_center, y_center, w, h = results[0, 0:4, i] * np.array([img_w, img_h, img_w, img_h])
                x1, y1, x2, y2 = int(x_center - w / 2), int(y_center - h / 2), int(x_center + w / 2), int(y_center + h / 2)
                class_id = np.argmax(results[0, 5:, i])  # Lấy class có xác suất cao nhất
                class_conf = results[0, class_id + 5, i]  # Xác suất của class đó

                boxes.append([x1, y1, x2, y2])
                scores.append(class_conf)
                class_ids.append(class_id)

        # Vẽ bounding box lên frame
        for box, score, class_id in zip(boxes, scores, class_ids):
            x1, y1, x2, y2 = box
            label = f"{class_names[class_id]}: {score:.2f}"
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        # Hiển thị kết quả
        cv2.imshow("YOLOv8 OpenVINO Video Detection", frame)

        # Thêm độ trễ và kiểm tra phím thoát
        key = cv2.waitKey(delay) & 0xFF
        if key == ord('q'):
            break

    except Exception as e:
        print(f"Lỗi xử lý frame: {str(e)}")
        break

# Giải phóng tài nguyên
cap.release()
cv2.destroyAllWindows()