基于CANN昇腾环境下做目标检测（使用YOLO模型）
·
我是提前训练好了YOLO模型，pt文件的。现在要在华为加速卡环境下部署推理。
0. 准备工作：先根据官方文档，把开发环境装好了
1、ONNX模型转OM模型（通过昇腾张量编译器ATC编译出来的）
由于我之前用的是pt模型，现在需要先将pt转为ONNX，再把ONNX转为OM模型。
转换代码如下，直接用YOLO的转换方法
from ultralytics import YOLO
from pathlib import Path
WEIGHTS_DIR = Path(__file__).resolve().parent / "weights"
ONNX_DIR = Path(__file__).resolve().parent / "onnx"
def pt2onnx(pt_path: str, onnx_path: str):
model = YOLO(pt_path)
model.export(format="onnx", imgsz=640, opset=11, simplify=True)
print(f"ONNX model saved to: {onnx_path}")
if __name__ == "__main__":
ONNX_DIR.mkdir(exist_ok=True)
pt_path = WEIGHTS_DIR / "best.pt"
onnx_path = ONNX_DIR / "best.onnx"
if not pt_path.exists():
print(f"Error: {pt_path} not found")
else:
pt2onnx(str(pt_path), str(onnx_path))
ONXX转OM模型zh
去服务器上执行转换命令
# 3. 执行 ATC 转换命令
# --soc_version 需替换为实际芯片型号，
# --input_shape 需根据模型实际输入定义
# --framework=5：代表 ONNX 模型。
atc --framework=5 \
--model=./best.onnx \
--output=./model_bs1 \
--input_format=NCHW \
--input_shape="images:1,3,640,640" \
--soc_version=Ascend310P3 \
--log=error过程运行的时候，会发现可能缺少一些依赖，可以执行下面的命令安装
第一次运行：pip3 install numpy
第二次运行：pip3 install sympy attrs psutil scipy
转换完成后，当前目录就出现了 om 模型文件
2. 推理运行
以下为官方推理demo，放到服务器上。执行。
函数：create_tensor_desc-CANN商用版9.0.0-昇腾社区 这里是api文档
"""
Ascend 昇腾推理脚本 - YOLO11 目标检测
基于 pyACL (Python Ascend Computing Language) 开发，用于在昇腾 NPU 上执行 YOLO 推理。
使用方法：
python3 ascend_infer.py --model ./model_bs1.om --image ./fire11.jpg --output result.jpg
前置条件：
1. 已通过 ATC 工具将 ONNX 模型转换为 .om 格式
2. 已安装 pyACL、OpenCV、NumPy
3. 已设置 Ascend 环境变量 (source set_env.sh)
"""
import os
import sys
import numpy as np
import cv2
# 导入 pyACL 模块（仅在昇腾设备上可用）
try:
import acl
except ImportError:
print("pyACL not found. Please run on Ascend device.")
sys.exit(1)
# ==================== 常量定义 ====================
# 内存分配策略：优先分配大页内存，提高大块内存分配效率
ACL_MEM_MALLOC_HUGE_FIRST = 0
# 内存拷贝方向
ACL_MEMCPY_HOST_TO_DEVICE = 1 # 主机(Host) -> 设备(Device)，用于上传输入数据
ACL_MEMCPY_DEVICE_TO_HOST = 2 # 设备(Device) -> 主机(Host)，用于下载输出结果
class YOLOInference:
"""
YOLO 推理类，封装了昇腾 NPU 上完整的推理流程：
1. 初始化 AscendCL 环境
2. 加载 .om 模型
3. 准备输入输出数据集（Device 内存）
4. 图像预处理（Resize、归一化、格式转换）
5. 执行推理
6. 后处理（坐标还原、NMS、阈值过滤）
7. 资源释放
"""
def __init__(self, model_path, input_size=640, conf_thresh=0.25, iou_thresh=0.45):
"""
初始化推理参数（尚未初始化 AscendCL，需调用 init()）
参数:
model_path: .om 模型文件路径
input_size: 模型输入图像尺寸（正方形），默认 640x640
conf_thresh: 置信度阈值，低于此值的检测框将被过滤，默认 0.25
iou_thresh: NMS 的 IoU 阈值，重叠度高于此值的框将被抑制，默认 0.45
"""
self.model_path = model_path
self.input_size = input_size
self.conf_thresh = conf_thresh
self.iou_thresh = iou_thresh
self.device_id = 0 # 使用第 0 号昇腾设备
# 以下变量在后续步骤中初始化
self.model_id = None # 模型 ID（由 load_from_file 返回）
self.model_desc = None # 模型描述符（包含输入输出信息）
self.input_dataset = None # 输入数据集（包含输入 buffer）
self.output_dataset = None # 输出数据集（包含输出 buffer）
self.input_data = None # 输入数据列表，每个元素包含 buffer 地址、大小等
self.output_data = None # 输出数据列表，每个元素包含 buffer 地址、大小等
self.output_dims = None # 输出 tensor 形状，如 [1, 6, 8400]
# ==================== 第一步：初始化环境与资源 ====================
def init(self):
"""
初始化 AscendCL 运行环境：
1. acl.init() - 初始化 AscendCL 框架
2. acl.rt.set_device() - 绑定指定的昇腾 NPU 设备
"""
# 初始化 AscendCL，必须第一个调用
ret = acl.init()
if ret != 0:
raise RuntimeError(f"acl.init failed: {ret}")
# 绑定设备，后续所有操作都在该设备上执行
ret = acl.rt.set_device(self.device_id)
if ret != 0:
raise RuntimeError(f"acl.rt.set_device failed: {ret}")
# ==================== 第二步：加载 OM 模型 ====================
def load_model(self):
"""
加载 .om 模型并获取模型信息：
1. 从文件加载模型，获得 model_id
2. 创建模型描述符，获取输入输出的维度、大小等信息
3. 在 Device 上申请输入输出内存
4. 创建输入输出数据集（Dataset + DataBuffer）
"""
# 从 .om 文件加载模型到设备
self.model_id, ret = acl.mdl.load_from_file(self.model_path)
if ret != 0:
raise RuntimeError(f"load_from_file failed: {ret}")
# 创建模型描述符并填充模型信息
self.model_desc = acl.mdl.create_desc()
ret = acl.mdl.get_desc(self.model_desc, self.model_id)
if ret != 0:
raise RuntimeError(f"get_desc failed: {ret}")
# 获取输出 tensor 形状，如 [1, 6, 8400]
# get_output_dims 返回格式: ({'name': '...', 'dims': [1,6,8400]}, 0)
out_dims_info = acl.mdl.get_output_dims(self.model_desc, 0)
self.output_dims = list(out_dims_info[0]["dims"])
print(f"Output shape: {self.output_dims}")
# 创建输入和输出数据集，包含 Device 内存分配和 DataBuffer 创建
self.input_dataset, self.input_data = self._prepare_dataset("input")
self.output_dataset, self.output_data = self._prepare_dataset("output")
def _prepare_dataset(self, io_type):
"""
准备输入或输出数据集（核心函数）
流程：
1. 获取输入/输出的数量和各自的 buffer 大小
2. 为每个输入/输出在 Device 上申请内存 (acl.rt.malloc)
3. 将 Device 内存地址包装为 DataBuffer (acl.create_data_buffer)
4. 将 DataBuffer 加入 Dataset (acl.mdl.add_dataset_buffer)
参数:
io_type: "input" 或 "output"
返回:
(dataset, datas):
dataset: acl Dataset 对象，传给 mdl.execute 使用
datas: 列表，每个元素包含 buffer 地址、DataBuffer 句柄、大小
"""
# 根据输入/输出类型，获取对应的数量和大小查询函数
if io_type == "input":
io_num = acl.mdl.get_num_inputs(self.model_desc) # 输入数量（YOLO 通常为 1）
get_size = acl.mdl.get_input_size_by_index # 获取指定输入的字节数
else:
io_num = acl.mdl.get_num_outputs(self.model_desc) # 输出数量（YOLO 通常为 1）
get_size = acl.mdl.get_output_size_by_index # 获取指定输出的字节数
# 创建数据集容器
dataset = acl.mdl.create_dataset()
datas = []
for i in range(io_num):
# 获取当前输入/输出的 buffer 大小（字节）
# 例如输入: 1*3*640*640*4 = 4915200 字节 (float32)
# 例如输出: 1*6*8400*4 = 201600 字节 (float32)
buffer_size = get_size(self.model_desc, i)
# 在 Device（NPU）上申请内存
# 返回: (device_ptr, error_code)
# ACL_MEM_MALLOC_HUGE_FIRST 表示优先分配大页内存
buffer, ret = acl.rt.malloc(buffer_size, ACL_MEM_MALLOC_HUGE_FIRST)
if ret != 0:
raise RuntimeError(f"rt.malloc failed: {ret}")
# 将 Device 内存地址包装为 DataBuffer 对象
# DataBuffer 是 pyACL 用于描述一块内存地址和大小的句柄
data_buffer = acl.create_data_buffer(buffer, buffer_size)
# 将 DataBuffer 加入到 Dataset 中
# Dataset 是 mdl.execute 需要的输入/输出容器
_, ret = acl.mdl.add_dataset_buffer(dataset, data_buffer)
if ret != 0:
raise RuntimeError(f"add_dataset_buffer failed: {ret}")
# 保存信息，后续 memcpy 和释放时需要用到
datas.append({
"buffer": buffer, # Device 内存地址（用于 memcpy）
"data": data_buffer, # DataBuffer 句柄（用于销毁）
"size": buffer_size # buffer 大小（字节）
})
if io_type == "input":
print(f"Input {i}: {buffer_size} bytes")
else:
print(f"Output {i}: {buffer_size} bytes")
return dataset, datas
# ==================== 第三步：图像预处理 ====================
def preprocess(self, image_path):
"""
将原始图像预处理为模型输入格式：
处理步骤：
1. 读取图像（OpenCV 默认 BGR 格式）
2. BGR -> RGB 通道转换
3. Resize 到模型输入尺寸（如 640x640）
4. 像素值归一化到 [0, 1]（除以 255）
5. HWC -> CHW 维度转换（Height,Width,Channel -> Channel,Height,Width）
6. 增加 batch 维度（NCHW 格式）
参数:
image_path: 输入图像文件路径
返回:
(input_data, orig_img, orig_w, orig_h):
input_data: 预处理后的 numpy 数组，shape=(1,3,640,640)，dtype=float32
orig_img: 原始图像（用于绘制结果）
orig_w: 原始图像宽度
orig_h: 原始图像高度
"""
# 读取图像，OpenCV 默认 BGR 格式
img = cv2.imread(image_path)
if img is None:
raise FileNotFoundError(f"Cannot read image: {image_path}")
# 记录原始尺寸，后续后处理时需要将坐标映射回原图
orig_h, orig_w = img.shape[:2]
# BGR -> RGB（YOLO 训练时使用 RGB 格式）
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# Resize 到模型输入尺寸（640x640）
img_resized = cv2.resize(img_rgb, (self.input_size, self.input_size))
# 归一化到 [0, 1]（训练时同样做了此操作）
img_float = img_resized.astype(np.float32) / 255.0
# HWC -> CHW：将通道维移到最前面
# 原始: (640, 640, 3) -> 转置后: (3, 640, 640)
img_nchw = img_float.transpose(2, 0, 1)[np.newaxis, ...]
# 增加 batch 维度：(3, 640, 640) -> (1, 3, 640, 640)
# 返回 float32，因为模型输入是 float32
return img_nchw.astype(np.float32), img, orig_w, orig_h
# ==================== 第四步：执行推理 ====================
def forward(self, input_data):
"""
在昇腾 NPU 上执行模型推理（同步）：
流程：
1. 将预处理后的数据从 Host 拷贝到 Device（输入）
2. 调用 mdl.execute 执行推理
3. 将推理结果从 Device 拷贝回 Host（输出）
参数:
input_data: 预处理后的输入数据，shape=(1,3,640,640)，dtype=float32
返回:
output: 模型输出，shape=(1,6,8400)，dtype=float32
6 = 4(坐标: cx,cy,w,h) + 2(类别置信度: fire,smoke)
8400 = 检测框数量（不同尺度的特征图拼接后）
"""
# ========== 1. Host -> Device：将输入数据上传到 NPU ==========
# 将 numpy 数组转为字节流
bytes_data = input_data.tobytes()
# 将字节流转为 pyACL 可识别的指针地址
bytes_ptr = acl.util.bytes_to_ptr(bytes_data)
# 执行内存拷贝：从 Host（主机内存）拷贝到 Device（NPU 显存）
# 参数：目标地址, 目标大小, 源地址, 源大小, 拷贝方向
ret = acl.rt.memcpy(
self.input_data[0]["buffer"], # Device 上的输入 buffer 地址
self.input_data[0]["size"], # Device buffer 大小
bytes_ptr, # Host 上的输入数据指针
len(bytes_data), # 数据字节数
ACL_MEMCPY_HOST_TO_DEVICE # 拷贝方向：Host -> Device
)
if ret != 0:
raise RuntimeError(f"memcpy H2D failed: {ret}")
# ========== 2. 执行推理 ==========
# 调用 NPU 执行模型推理（同步阻塞，直到推理完成）
# 参数：模型ID, 输入数据集, 输出数据集
ret = acl.mdl.execute(self.model_id, self.input_dataset, self.output_dataset)
if ret != 0:
raise RuntimeError(f"mdl.execute failed: {ret}")
# ========== 3. Device -> Host：将输出数据下载到主机 ==========
# 在 Host 上申请临时内存，用于接收输出数据
buffer_host, ret = acl.rt.malloc_host(self.output_data[0]["size"])
if ret != 0:
raise RuntimeError(f"malloc_host failed: {ret}")
# 执行内存拷贝：从 Device（NPU 显存）拷贝到 Host（主机内存）
ret = acl.rt.memcpy(
buffer_host, # Host 上的目标地址
self.output_data[0]["size"], # 目标大小
self.output_data[0]["buffer"], # Device 上的输出 buffer 地址
self.output_data[0]["size"], # 源大小
ACL_MEMCPY_DEVICE_TO_HOST # 拷贝方向：Device -> Host
)
if ret != 0:
raise RuntimeError(f"memcpy D2H failed: {ret}")
# 将 Host 内存指针转为字节流，再转为 numpy 数组
bytes_out = acl.util.ptr_to_bytes(buffer_host, self.output_data[0]["size"])
# 按模型输出形状 reshape，dtype=float32
output = np.frombuffer(bytes_out, dtype=np.float32).reshape(self.output_dims)
# 释放 Host 上的临时内存
acl.rt.free_host(buffer_host)
return output
# ==================== 第五步：后处理 ====================
def postprocess(self, output, orig_w, orig_h):
"""
对模型原始输出进行后处理，得到最终检测结果：
处理步骤：
1. 解析输出 tensor，分离边界框坐标和类别分数
2. 置信度过滤：去掉低置信度的检测框
3. 坐标还原：从模型输入尺寸映射回原始图像尺寸
4. NMS（非极大值抑制）：去除重叠度高的冗余框
YOLO 输出格式：(1, 6, 8400)
- 第 0~3 维：cx, cy, w, h（中心点坐标和宽高，归一化到 0~640）
- 第 4~5 维：fire 置信度, smoke 置信度
- 8400：检测框总数（不同尺度特征图的锚框数量之和）
参数:
output: 模型输出，shape=(1,6,8400)
orig_w: 原始图像宽度
orig_h: 原始图像高度
返回:
(boxes, scores, class_ids):
boxes: 检测框坐标 [[x1,y1,x2,y2], ...]，已映射到原图尺寸
scores: 置信度分数
class_ids: 类别 ID（0=fire, 1=smoke）
"""
# 去掉 batch 维度：(1,6,8400) -> (6,8400)
output = output[0]
# 如果输出格式是 (6, 8400)，转置为 (8400, 6) 方便按行处理
if output.shape[0] == 6:
output = output.T
# 分离坐标和分数
boxes = output[:, :4] # 前 4 列：cx, cy, w, h
scores = output[:, 4:] # 后 2 列：fire_score, smoke_score
# 取每个检测框的最大类别分数及其类别 ID
class_ids = np.argmax(scores, axis=1) # 每行最大值的索引（0 或 1）
max_scores = np.max(scores, axis=1) # 每行的最大值
# 置信度过滤：去掉分数低于阈值的框
mask = max_scores > self.conf_thresh
boxes = boxes[mask]
max_scores = max_scores[mask]
class_ids = class_ids[mask]
# 没有检测到目标，返回空数组
if len(boxes) == 0:
return np.array([]), np.array([]), np.array([])
# 坐标转换：从 (cx, cy, w, h) 转为 (x1, y1, x2, y2)，并映射回原图尺寸
# 公式：x1 = (cx - w/2) / input_size * orig_w
x1 = (boxes[:, 0] - boxes[:, 2] / 2) / self.input_size * orig_w
y1 = (boxes[:, 1] - boxes[:, 3] / 2) / self.input_size * orig_h
x2 = (boxes[:, 0] + boxes[:, 2] / 2) / self.input_size * orig_w
y2 = (boxes[:, 1] + boxes[:, 3] / 2) / self.input_size * orig_h
# 拼接为 (N, 4) 的 xyxy 格式
boxes_xyxy = np.stack([x1, y1, x2, y2], axis=1)
# NMS（非极大值抑制）：去除高度重叠的冗余检测框
# OpenCV 的 NMSBoxes 要求输入为 list 格式
indices = cv2.dnn.NMSBoxes(
boxes_xyxy.tolist(), # 检测框列表
max_scores.tolist(), # 对应分数列表
self.conf_thresh, # 置信度阈值
self.iou_thresh # IoU 阈值（重叠度高于此值则抑制）
)
# 处理 NMS 返回结果
if len(indices) == 0:
return np.array([]), np.array([]), np.array([])
if isinstance(indices, np.ndarray):
indices = indices.flatten()
elif isinstance(indices, list) and len(indices) > 0:
indices = np.array(indices).flatten()
else:
return np.array([]), np.array([]), np.array([])
# 返回 NMS 后保留的检测结果
return boxes_xyxy[indices], max_scores[indices], class_ids[indices]
# ==================== 绘制检测结果 ====================
def draw_results(self, image, boxes, scores, class_ids, class_names=None):
"""
在图像上绘制检测框和标签
参数:
image: 原始图像
boxes: 检测框坐标 [[x1,y1,x2,y2], ...]
scores: 置信度分数
class_ids: 类别 ID
class_names: 类别名称字典，默认 {0: "fire", 1: "smoke"}
返回:
img: 绘制了检测框的图像
"""
img = image.copy()
# 默认类别名称（火灾检测场景）
if class_names is None:
class_names = {0: "fire", 1: "smoke"}
# 不同类别的颜色（BGR 格式）
colors = [(0, 0, 255), (0, 165, 255)] # 红色=fire, 橙色=smoke
# 没有检测结果，直接返回原图
if len(boxes) == 0:
return img
# 遍历每个检测框，绘制矩形和标签
for box, score, cls_id in zip(boxes, scores, class_ids):
x1, y1, x2, y2 = map(int, box)
color = colors[int(cls_id) % len(colors)]
name = class_names.get(int(cls_id), f"class_{cls_id}")
# 绘制检测框（矩形）
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
# 绘制标签背景和文字
label = f"{name} {score:.2f}"
(tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
cv2.rectangle(img, (x1, y1 - th - 6), (x1 + tw, y1), color, -1) # 背景矩形
cv2.putText(img, label, (x1, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)
return img
# ==================== 第六步：释放资源 ====================
def destroy(self):
"""
按相反顺序释放所有 AscendCL 资源：
释放顺序（与初始化顺序相反）：
1. 销毁 DataBuffer，释放 Device 内存
2. 销毁 Dataset
3. 销毁模型描述符
4. 卸载模型
5. 重置设备
6. AscendCL 去初始化
"""
# 释放输入和输出的 DataBuffer 和 Device 内存
for dataset_data in [self.input_data, self.output_data]:
if dataset_data:
for item in dataset_data:
acl.destroy_data_buffer(item["data"]) # 销毁 DataBuffer 句柄
acl.rt.free(item["buffer"]) # 释放 Device 内存
# 销毁 Dataset
if self.input_dataset:
acl.mdl.destroy_dataset(self.input_dataset)
if self.output_dataset:
acl.mdl.destroy_dataset(self.output_dataset)
# 销毁模型描述符
if self.model_desc:
acl.mdl.destroy_desc(self.model_desc)
# 卸载模型
if self.model_id is not None:
acl.mdl.unload(self.model_id)
# 重置设备并去初始化
acl.rt.reset_device(self.device_id)
acl.finalize()
def main():
"""
主函数：解析命令行参数，执行完整的推理流程
"""
import argparse
parser = argparse.ArgumentParser(description="YOLO11 Fire Detection on Ascend")
parser.add_argument("--model", type=str, required=True, help="Path to .om model file")
parser.add_argument("--image", type=str, required=True, help="Path to input image")
parser.add_argument("--output", type=str, default="result.jpg", help="Output image path")
parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold")
parser.add_argument("--iou", type=float, default=0.45, help="IoU threshold for NMS")
parser.add_argument("--input-size", type=int, default=640, help="Model input size")
args = parser.parse_args()
# 创建推理实例
detector = YOLOInference(
model_path=args.model,
input_size=args.input_size,
conf_thresh=args.conf,
iou_thresh=args.iou
)
try:
# 第一步：初始化 AscendCL 环境
print("Initializing Ascend environment...")
detector.init()
# 第二步：加载 .om 模型
print(f"Loading model: {args.model}")
detector.load_model()
# 第三步：图像预处理
print(f"Preprocessing image: {args.image}")
input_data, orig_img, orig_w, orig_h = detector.preprocess(args.image)
# 第四步：执行推理
print("Running inference...")
output = detector.forward(input_data)
# 第五步：后处理
print("Postprocessing results...")
boxes, scores, class_ids = detector.postprocess(output, orig_w, orig_h)
# 打印检测结果
print(f"Detected {len(boxes)} objects")
for i, (box, score, cls_id) in enumerate(zip(boxes, scores, class_ids)):
print(f" [{i}] class={int(cls_id)}, score={score:.4f}, box={box.astype(int).tolist()}")
# 绘制结果并保存
result_img = detector.draw_results(orig_img, boxes, scores, class_ids)
cv2.imwrite(args.output, result_img)
print(f"Result saved to: {args.output}")
finally:
# 第六步：释放所有资源（无论是否出错都会执行）
detector.destroy()
print("Resources released.")
if __name__ == "__main__":
main()
执行推理
python3 ascend_infer.py --model ./model_bs1.om --image ./fire11.jpg
图片下载下来查看，也是推理成功了，这次的推理用到的就是华为的NPU了
作为“人工智能6S店”的官方数字引擎，为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。
更多推荐
0
0- 0
已为社区贡献1条内容
所有评论(0)