58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
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()
|