import tempfile
import time

import cv2
import torch
from ultralytics import YOLO

from log import logger

if __name__ == "__main__":
    model_config = {
        'model_path': r"F:\PyModelScope\Yolov\models\best-fbrt.pt",
        'imgs': r"E:\DC\Yolov\uploads\费县8-13-19-9-14-20-14.JPG"
    }

    # encryption_key = "O3^doTYEpyT%VCYdI6u#YKKi7YFWjGZQ"
    local_path = model_config.get('model_path')
    # from mandatory_model_crypto import MandatoryModelValidator
    #
    # validator = MandatoryModelValidator()
    # decrypt_result = validator.decrypt_and_verify(local_path, encryption_key)
    # decrypted_data = decrypt_result['decrypted_data']
    # # 保存到临时文件并加载
    # with tempfile.NamedTemporaryFile(delete=False, suffix='.pt') as tmp:
    #     tmp.write(decrypted_data)
    #     temp_path = tmp.name

    # 加载YOLO模型
    model = YOLO(local_path).to('cuda:0')
    start = time.time()

    with torch.no_grad():
        results = model.predict(
            source=model_config['imgs'],
            stream=False,
            verbose=False,
            conf=0.5,
            iou=0.45,
            imgsz=1920,
    )
    logger.info(f'推理时间:{time.time() - start}')
    # 提取检测结果
    for result in results:
        boxes = result.boxes.xyxy  # 边界框坐标
        scores = result.boxes.conf  # 置信度分数
        classes = result.boxes.cls  # 类别索引

        # 如果有类别名称，可以通过类别索引获取
        class_names = [model.names[int(cls)] for cls in classes]

        # 打印检测结果
        for box, score, class_name in zip(boxes, scores, class_names):
            print(f"Class: {class_name}, Score: {score:.2f}, Box: {box}")
        # # 可视化检测结果，图像数据格式是BGR
        annotated_img = result.plot(
            line_width=1,
            font_size=9
        )
        logger.info(f'startTime:{start},endTime:{time.time()},时间差:{time.time() - start}')
        result.save(r'E:\DC\Yolov\uploads\result5.jpg')
        # 显示图像
        cv2.destroyAllWindows()