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 và in ra để kiểm tra
fps = cap.get(cv2.CAP_PROP_FPS)
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
print(f"Video FPS: {fps}")
print(f"Video size: {frame_width}x{frame_height}")

# 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:
        # Kiểm tra frame có rỗng không
        if frame is None:
            print("Frame rỗng")
            continue

        # Kiểm tra kích thước frame
        print(f"Frame shape: {frame.shape}")

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

        # Chạy mô hình OpenVINO
        results = compiled_model([input_tensor])[compiled_model.output(0).any_name]
        
        # Debug thông tin chi tiết
        print("\n--- Debug Info ---")
        print("Input tensor shape:", input_tensor.shape)
        print("Results shape:", results.shape)
        print("Results min/max:", np.min(results), np.max(results))
        
        # Chuyển đổi kết quả (1, 84, 8400) -> (8400, 84)
        predictions = np.transpose(results[0])
        print("Predictions shape:", predictions.shape)
        
        # Lấy scores và class_ids
        conf_threshold = 0.1  # Giảm ngưỡng xuống để dễ phát hiện
        object_scores = predictions[:, 4]  # Confidence scores
        class_scores = predictions[:, 5:]  # Class scores
        
        # Debug scores
        print("Object confidence scores:", object_scores[:5])  # In 5 giá trị đầu tiên
        print("Class scores shape:", class_scores.shape)
        print("Max object confidence:", np.max(object_scores))
        print("Max class score:", np.max(class_scores))
        
        # Lọc các dự đoán có confidence > threshold
        mask = object_scores > conf_threshold
        boxes_detected = predictions[mask, :4]
        scores_detected = object_scores[mask]
        class_scores_detected = class_scores[mask]
        class_ids_detected = np.argmax(class_scores_detected, axis=1)
        
        print("\n--- Detection Results ---")
        print(f"Number of detections: {len(boxes_detected)}")
        if len(boxes_detected) > 0:
            print("First detection:")
            print(f"- Box coordinates: {boxes_detected[0]}")
            print(f"- Confidence: {scores_detected[0]:.4f}")
            print(f"- Class: {class_names[class_ids_detected[0]]}")

        # Chuyển đổi xywh -> xyxy
        boxes_detected[:, 0] = boxes_detected[:, 0] - boxes_detected[:, 2] / 2  # x1
        boxes_detected[:, 1] = boxes_detected[:, 1] - boxes_detected[:, 3] / 2  # y1
        boxes_detected[:, 2] = boxes_detected[:, 0] + boxes_detected[:, 2]      # x2
        boxes_detected[:, 3] = boxes_detected[:, 1] + boxes_detected[:, 3]      # y2
        
        # Scale boxes về kích thước ảnh gốc
        boxes_detected[:, [0, 2]] *= frame_width
        boxes_detected[:, [1, 3]] *= frame_height
        
        # Chuyển về int
        boxes_detected = boxes_detected.astype(np.int32)
        
        # Áp dụng NMS
        indices = cv2.dnn.NMSBoxes(
            boxes_detected.tolist(), 
            scores_detected.tolist(),
            conf_threshold, 
            0.45
        )
        
        # Vẽ các box
        for i in indices:
            box = boxes_detected[i]
            x1, y1, x2, y2 = box
            
            # Vẽ box
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            
            # Vẽ label
            label = f"{class_names[class_ids_detected[i]]}: {scores_detected[i]:.2f}"
            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(1) & 0xFF  # Giảm độ trễ xuống 1ms
        if key == ord('q'):
            break

    except Exception as e:
        print(f"\nLỗi xử lý frame: {str(e)}")
        import traceback
        traceback.print_exc()
        break

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