63 lines
2.0 KiB
Python
63 lines
2.0 KiB
Python
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()
|