import tempfile

import cv2
import torch
from ultralytics import YOLO

if __name__ == "__main__":
    model_config = {
        'model_path': r'E:\DC\Yolov\encrypted_models\0deb0a16a7c17415.enc',
        'imgs': r'F:\PyModelScope\Yolov\images\lQDPJxJjZJOlMDPNDkDNFWCw5IKSN_bsEtAI7HvAI16LAA_5472_3648.jpg',
    }

    encryption_key = "Mcj4#lpC4Qad$ba%QrM&@mw5OREpsX&G"
    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(temp_path).to('cuda:0')
    with torch.no_grad():
        results = model.predict(
            source=model_config['imgs'],
            stream=False,
            verbose=False,
            conf=0.5,
            iou=0.45,
            imgsz=1920,
        )

    # 提取检测结果
    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
        )
        result.save(r'E:\DC\Yolov\uploads\result5.jpg')
        # 显示图像
        cv2.destroyAllWindows()