songhaoyu/opi-yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

npu端基本测试

Browse Source
This commit is contained in:
songhaoyu
2026-05-21 16:53:31 +08:00
parent 0ec79c9781
commit 719d4081e8
10 changed files with 430 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
infer.py
+184 -64
View File
@@ -1,78 +1,198 @@
import os
import cv2
import numpy as np
import argparse
from ais_bench.infer.interface import InferSession
# 导入昇腾核心推理库（板子自带）
from acllite_model import AclLiteModel
from acllite_resource import AclLiteResource
class YOLO26Infer:
def __init__(self, model_path, device_id=0):
"""
初始化推理器
:param model_path: .om 模型的本地路径
:param device_id: NPU 设备ID
"""
self.device_id = device_id
self.model_path = model_path
# 创建 ais_bench 推理会话
self.session = InferSession(device_id, model_path)
# 获取模型输入尺寸，用于后续预处理
self.input_width = 640
self.input_height = 640
def preprocess(img_path, model_w=640, model_h=640):
"""图片预处理：缩放、通道转换、归一化"""
orig_img = cv2.imread(img_path)
h, w, _ = orig_img.shape
# 1. 缩放到模型要求的 640x640
img = cv2.resize(orig_img, (model_w, model_h))
# 2. BGR 转 RGB (YOLO 训练时通常是 RGB)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# 3. 归一化并调整维度为 [1, 3, 640, 640]
img = img.astype(np.float32) / 255.0
img = img.transpose(2, 0, 1) # HWC -> CHW
img = np.expand_dims(img, axis=0) # CHW -> NCHW
# 必须将其转换为连续的内存，方便 NPU 读取
return np.ascontiguousarray(img), orig_img
def letterbox(self, img, new_shape=(640, 640), color=(114, 114, 114)):
shape = img.shape[:2]
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
# 计算缩放后的新尺寸（四舍五入）
new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))
# 计算需要填充的总宽度和高度
dw = new_shape[1] - new_unpad[0]
dh = new_shape[0] - new_unpad[1]
# 分别计算左右/上下填充量（允许奇数不对称）
left = dw // 2
right = dw - left
top = dh // 2
bottom = dh - top
def main():
# 1. 初始化昇腾 NPU 资源
acl_resource = AclLiteResource()
acl_resource.init()
# 缩放图片
if shape[::-1] != new_unpad:
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
# 2. 加载你的 YOLO26 .om 模型
model_path = "./yolo26_8t.om"
model = AclLiteModel(model_path)
# 添加填充（可能左右或上下不对称，但总尺寸正好是 new_shape）
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
return img, r, (left, top) # 返回填充的左上角偏移量
# 3. 图像预处理
input_data, orig_img = preprocess("./test.jpg")
def preprocess(self, image_path):
"""
预处理图片，使其符合模型输入要求
"""
# 使用OpenCV读取图片 (BGR格式)
img_bgr = cv2.imread(image_path)
if img_bgr is None:
raise FileNotFoundError(f"图片未找到: {image_path}")
# 4. 送入 NPU 进行硬件加速推理
# execute 接收一个列表，里面是模型的输入 Tensor
result_list = model.execute([input_data])
# 获取原始图片尺寸，用于后处理时的坐标转换
self.original_height, self.original_width = img_bgr.shape[:2]
# 5. 后处理（针对 YOLO26 的端到端特性）
# YOLO26 输出的直接是最终预测框，result_list[0] 形状通常是 [1, 300, 6] 或类似
# 6个维度的含义一般为：[x_min, y_min, x_max, y_max, score, class_id]
predictions = result_list[0][0]
# 将图片从BGR转换为RGB
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
CONF_THRESHOLD = 0.25 # 设定置信度阈值
orig_h, orig_w = orig_img.shape[:2]
# 调整图片尺寸并保持宽高比，进行Letterbox填充
img_resized, self.ratio, (self.pad_left, self.pad_top) = self.letterbox(
img_rgb, (self.input_width, self.input_height)
)
box_count = 0
for pred in predictions:
# 取出置信度和类别
score = pred[4]
class_id = int(pred[5])
# 过滤掉低置信度的框
if score > CONF_THRESHOLD:
box_count += 1
# 取出缩放后的坐标 [0~640]
x1, y1, x2, y2 = pred[0:4]
# 将坐标还原到原图尺寸
x1 = int(x1 * orig_w / 640)
y1 = int(y1 * orig_h / 640)
x2 = int(x2 * orig_w / 640)
y2 = int(y2 * orig_h / 640)
# 在原图上画框和标签
cv2.rectangle(orig_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(orig_img, f"Class {class_id}: {score:.2f}", (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
# 归一化：将像素值从 [0, 255] 转为 [0, 1]
img_input = (img_resized.astype(np.float32) / 255.0).astype(np.float32)
print(f"推理完成，共检测到 {box_count} 个目标。")
# 6. 保存检测结果图
cv2.imwrite("./result.jpg", orig_img)
# 调整维度：HWC -> NCHW，并转换为连续数组
img_input = np.ascontiguousarray(np.transpose(img_input, (2, 0, 1))[np.newaxis, :, :, :])
return img_input
def infer(self, img_input):
"""
执行推理
"""
# 使用会话的 infer 方法进行推理
# infer 方法接收一个列表，包含输入张量，返回一个列表，包含输出张量
result = self.session.infer([img_input])
print(f"Input tensor shape: {img_input.shape}")
print(f"Input tensor dtype: {img_input.dtype}")
print(f"Input tensor size: {img_input.nbytes}")
return result
def postprocess(self, infer_result, conf_threshold=0.5):
"""
对模型输出进行后处理
:param infer_result: 模型原始输出，一个列表，通常第一个元素是输出张量
:param conf_threshold: 置信度阈值
:return: 包含检测结果的列表，每个元素为 [x1, y1, x2, y2, confidence, class_id]
"""
detections = []
# 获取输出张量，假设输出列表的第一个元素是输出数组
if isinstance(infer_result, (list, tuple)):
output = infer_result[0]
else:
output = infer_result
# ========== 调试信息 ==========
# print(f"[DEBUG] output shape: {output.shape}")
# print(f"[DEBUG] output dtype: {output.dtype}")
# print(f"[DEBUG] output min: {output.min()}, max: {output.max()}")
# # 打印前5个检测框（如果形状是 [N,6] 或 [1,N,6]）
# if output.ndim == 3:
# flat_output = output[0]
# else:
# flat_output = output
# print(f"[DEBUG] first 5 rows:\n{flat_output[:5]}")
# 假设输出维度为 [batch, num_detections, 6]
# 移除batch维度，并只取第一个batch的结果
if output.ndim == 3:
output = output[0]
# 遍历所有检测框
for det in output:
confidence = det[4]
if confidence < conf_threshold:
continue
# 获取原始输出坐标 (x1, y1, x2, y2) 和 类别ID
x1_scale, y1_scale, x2_scale, y2_scale = det[0], det[1], det[2], det[3]
cls_id = int(det[5])
# 将坐标从模型输入尺寸映射回原始图片尺寸
x1 = (x1_scale - self.pad_left) / self.ratio
y1 = (y1_scale - self.pad_top) / self.ratio
x2 = (x2_scale - self.pad_left) / self.ratio
y2 = (y2_scale - self.pad_top) / self.ratio
# 边界限制，防止超出原始图片范围
x1 = max(0, min(x1, self.original_width))
y1 = max(0, min(y1, self.original_height))
x2 = max(0, min(x2, self.original_width))
y2 = max(0, min(y2, self.original_height))
detections.append([x1, y1, x2, y2, confidence, cls_id])
return detections
def draw_boxes(image, detections, class_names=None):
"""
在图片上绘制检测框和标签
:param image: 原始图片数组 (BGR格式)
:param detections: 后处理得到的结果列表
:param class_names: 类别名称列表，默认使用类别ID
"""
img_copy = image.copy()
for det in detections:
x1, y1, x2, y2, conf, cls_id = det
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
# 绘制矩形框
cv2.rectangle(img_copy, (x1, y1), (x2, y2), (0, 255, 0), 2)
# 构建标签文本
label = f"Class: {cls_id}" if class_names is None else class_names[cls_id]
label += f" {conf:.2f}"
# 设置标签背景
(label_w, label_h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
cv2.rectangle(img_copy, (x1, y1 - label_h - 5), (x1 + label_w, y1), (0, 255, 0), -1)
cv2.putText(img_copy, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
return img_copy
if __name__ == "__main__":
main()
parser = argparse.ArgumentParser(description="YOLOv26 目标检测推理")
parser.add_argument("--model", type=str, required=True, help="YOLOv26 OM模型文件路径")
parser.add_argument("--image", type=str, required=True, help="输入图片路径")
parser.add_argument("--conf", type=float, default=0.5, help="置信度阈值")
args = parser.parse_args()
# 初始化推理器
yolo_infer = YOLO26Infer(args.model)
# 图片预处理
print("预处理图片...")
input_tensor = yolo_infer.preprocess(args.image)
# 执行推理
print("执行推理...")
raw_result = yolo_infer.infer(input_tensor)
# 后处理
print("后处理...")
detections = yolo_infer.postprocess(raw_result, conf_threshold=args.conf)
# 打印结果
print(f"检测到 {len(detections)} 个目标:")
for det in detections:
print(f"类别ID: {int(det[5])}, 置信度: {det[4]:.4f}, 边界框: ({det[0]:.2f}, {det[1]:.2f}, {det[2]:.2f}, {det[3]:.2f})")
# 可选：绘制结果并保存
original_img = cv2.imread(args.image)
result_img = draw_boxes(original_img, detections)
output_path = "output_" + os.path.basename(args.image)
cv2.imwrite(output_path, result_img)
print(f"结果图片已保存至: {output_path}")