yipai-tech/SupervisorAI
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

print改为logger打印

Browse Source
This commit is contained in:
zqc
2026-02-03 09:15:37 +08:00
parent 42d7d24c4d
commit f6ea6082f3
5 changed files with 46 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
algorithm/checkpoint/npu_yolo_pose_onnx.py
View File

@@ -9,7 +9,8 @@ import cv2
import numpy as np
import onnxruntime as ort
from utils.logger import get_logger
logger = get_logger(__name__)
# -------------------------------------------------
@@ -89,14 +90,14 @@ class YOLOv8_Pose_ONNX:
# 获取真实工作 provider
actual_providers = self.session.get_providers()
print("YOLO Providers:", actual_providers)
logger.info("YOLO Providers:", actual_providers)
if "CANNExecutionProvider" in actual_providers:
print("[INFO] YOLO 使用 CANNExecutionProvider（昇腾）")
logger.info("[INFO] YOLO 使用 CANNExecutionProvider（昇腾）")
elif 'CUDAExecutionProvider' in actual_providers:
print("[INFO] YOLO 使用 CUDAExecutionProvider（NVIDIA GPU）")
logger.info("[INFO] YOLO 使用 CUDAExecutionProvider（NVIDIA GPU）")
else:
print("[INFO] YOLO 使用 CPUExecutionProvider（非昇腾环境）")
logger.info("[INFO] YOLO 使用 CPUExecutionProvider（非昇腾环境）")
self.conf_threshold = conf_threshold
self.iou_threshold = iou_threshold
@@ -104,9 +105,9 @@ class YOLOv8_Pose_ONNX:
self.input_name = self.session.get_inputs()[0].name
self.input_size = (input_size, input_size)
print(f"模型输入名称: {self.input_name}")
print(f"模型输入形状: {self.session.get_inputs()[0].shape}")
print(f"模型输出形状: {self.session.get_outputs()[0].shape}")
logger.info(f"模型输入名称: {self.input_name}")
logger.info(f"模型输入形状: {self.session.get_inputs()[0].shape}")
logger.info(f"模型输出形状: {self.session.get_outputs()[0].shape}")
def nms(self, boxes, scores, iou_threshold=0.45):