import time

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



def preprocess_image(image, input_size=(640, 640)):

    img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转换为 RGB
    img_resized = cv2.resize(img, input_size)  # 调整大小
    img_normalized = img_resized / 255.0  # 归一化
    img_transposed = np.transpose(img_normalized, (2, 0, 1))  # HWC -> CHW
    img_batch = np.expand_dims(img_transposed, axis=0).astype(np.float32)  # 增加 batch 维度
    return img_batch, img_resized

def postprocess(output, conf_threshold=0.5, iou_threshold=0.4):
    # 解析输出
    boxes = []
    scores = []
    class_ids = []

    for detection in output[0]:
        confidence = detection[4]
        if confidence > conf_threshold:
            class_id = np.argmax(detection[5:])
            cx, cy, w, h = detection[:4]  # 中心点和宽高
            x_min = int((cx - w / 2) * original_image.shape[1])  # 转换回原始图像尺寸
            y_min = int((cy - h / 2) * original_image.shape[0])
            width = int(w * original_image.shape[1])
            height = int(h * original_image.shape[0])

            boxes.append([x_min, y_min, width, height])
            scores.append(confidence)
            class_ids.append(class_id)

    # 非极大值抑制 (NMS)
    indices = cv2.dnn.NMSBoxes(boxes, scores, conf_threshold, iou_threshold)

    final_boxes = []
    for idx in indices:
        box = boxes[idx]
        final_boxes.append({
            "box": box,
            "score": scores[idx],
            "class_id": class_ids[idx]
        })

    return final_boxes


def draw_detections(image, detections):
    for det in detections:
        x_min, y_min, width, height = det["box"]
        class_id = det["class_id"]
        score = det["score"]

        # 确保坐标和尺寸是整数
        x_min = int(x_min)
        y_min = int(y_min)
        width = int(width)
        height = int(height)

        # 绘制边界框
        cv2.rectangle(image, (x_min, y_min), (x_min + width, y_min + height), (0, 255, 0), 2)

        # 显示类别和置信度
        label = f"Class {class_id}: {score:.2f}"
        cv2.putText(image, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

    return image


if __name__ == "__main__":
    # 初始化 OpenVINO Runtime
    core = Core()

    # 加载模型
    model_path = "model/safe_det/best.xml"  # 替换为你的模型路径
    model = core.read_model(model=model_path)
    compiled_model = core.compile_model(model=model, device_name="CPU")  # 设备可以是 "CPU", "GPU", "MYRIAD" 等

    # 获取输入和输出层
    input_layer = compiled_model.input(0)
    output_layer = compiled_model.output(0)

    # 打印输入输出信息
    print(f"Input shape: {input_layer.shape}")
    print(f"Output shape: {output_layer.shape}")

    # 预处理图像，改为从摄像头获取
    cap = cv2.VideoCapture(0)

    # 初始化变量
    start_time = time.time()  # 记录开始时间
    frame_count = 0  # 帧计数器

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        input_data, original_image = preprocess_image(frame)

        # 推理
        results = compiled_model([input_data])[output_layer]
        # 输出结果
        print(results)
        detections = postprocess(results)
        # 打印检测结果
        for det in detections:
            print(det)
        draw_img = draw_detections(frame, detections)

        # 计算FPS
        frame_count += 1
        elapsed_time = time.time() - start_time

        fps = frame_count / elapsed_time
        # start_time = time.time()  # 重置开始时间
        # frame_count = 0  # 重置帧计数器

        # 显示FPS
        cv2.putText(draw_img, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        cv2.imshow("Detections", draw_img)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break