旭日X5部署YOLO11目标检测保姆级教程
·
1.环境准备
D-Robotics OpenExplore(RDK X5, Bayes-e BPU) Version: >= 1.2.8
Ultralytics YOLO Version: >= 8.3.0
Ubuntu 22.04
Python 3.10
2.模型训练
2.1 下载ultralytics/ultralytics仓库
git clone https://github.com/ultralytics/ultralytics.git2.2 创建conda环境
conda create -n rdkx5 python=3.10
conda activate rdkx5
pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple2.3训练模型
model替换为你的预训练模型，data替换为你的yaml文件，训练时无需修改任何程序, 无需修改forward方法.
Ultralytics YOLO 官方文档:
yolo detect train model=yolo11s.pt data=detection.yaml epochs=100 imgsz=640 batch=32 device=0 name=yolo11_run amp=True cache=True3.模型导出（pt转onnx，windows环境也适用）
官方推荐使用Ubuntu 22.04, Python 3.10的环境。我在windows环境Python3.9环境下也可以正常使用。默认你已经拥有了模型文件(.pt)，将--pt中替换成自己训练好的pt模型。
在Ultralytics YOLO的训练环境中, 运行RDK Model Zoo 提供的一键YOLO导出脚本https://github.com/D-Robotics/rdk_model_zoo/blob/main/demos/Vision/ultralytics_YOLO/x86/export_monkey_patch.py, 对模型进行导出. 这个脚本会使用ultralytics.YOLO类对YOLO的pt模型进行加载, 使用猴子补丁(Monkey Patch)的方法对模型在PyTorch层面进行替换, 进行并调用ultralytics.YOLO.export方法对模型进行导出. 导出的ONNX模型会保存在pt模型同级目录下.
python3 export_monkey_patch.py --pt yolo11n.pt
#根据自己指令选择
python export_monkey_patch.py --pt yolo11n.ptexport_monkey_patch.py代码如下：
import argparse
from ultralytics import YOLO
from ultralytics.nn.modules.head import Detect, v10Detect, Segment, OBB, Pose, Classify
from ultralytics.nn.modules.block import Attention, AAttn
import torch
import types
import os
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--pt', type=str, default='./yolo11n.pt', help='path to *.pt model.')
parser.add_argument('--optse', type=int, default=11, help='opset version.')
opt = parser.parse_args()
# Init Ultralytics YOLO Model
m = YOLO(opt.pt)
# Replace some efficient modules
modelZooOptimizer(m.model.model)
# Export to ONNX
m.export(format='onnx', simplify=False, opset=11)
def modelZooOptimizer(model): # Monkey Patch
for name, child in model.named_children():
# print(name)
if type(child) == Classify:
child.forward = types.MethodType(Classify_forward, child)
print("\033[1;31m" + f"[Cauchy] Replaced Classify_forward in {name}" + "\033[0m")
elif type(child) == Detect:
child.forward = types.MethodType(Detect_forward, child)
print("\033[1;31m" + f"[Cauchy] Replaced Detect_forward in {name}" + "\033[0m")
elif type(child) == v10Detect:
child.forward = types.MethodType(v10Detect_forward, child)
print("\033[1;31m" + f"[Cauchy] Replaced v10Detect_forward in {name}" + "\033[0m")
elif type(child) == Segment:
child.forward = types.MethodType(Segment_forward, child)
print("\033[1;31m" + f"[Cauchy] Replaced Segment_forward in {name}" + "\033[0m")
elif type(child) == Pose:
child.forward = types.MethodType(Pose_forward, child)
print("\033[1;31m" + f"[Cauchy] Replaced Pose_forward in {name}" + "\033[0m")
elif type(child) == OBB:
child.forward = types.MethodType(OBB_forward, child)
print("\033[1;31m" + f"[Cauchy] Replaced OBB_forward in {name}" + "\033[0m")
elif type(child) == AAttn:
child.forward = types.MethodType(AAttn_forward, child)
print("\033[1;31m" + f"[Cauchy] Replaced AAttn_forward in {name}" + "\033[0m")
elif type(child) == Attention:
child.forward = types.MethodType(Attention_forward, child)
print("\033[1;31m" + f"[Cauchy] Replaced Attention_forward in {name}" + "\033[0m")
modelZooOptimizer(child)
def Attention_forward(self, x):
# Effieicient for Bayes-e BPU
B, C, H, W = x.shape
N = H * W
qkv = self.qkv(x)
q, k, v = qkv.view(B, self.num_heads, self.key_dim * 2 + self.head_dim, N).split([self.key_dim, self.key_dim, self.head_dim], dim=2)
attn = (q.transpose(-2, -1) @ k) * self.scale
attn = attn.permute(0, 3, 1, 2).contiguous() # CHW2HWC like
max_attn = attn.max(dim=1, keepdim=True).values
exp_attn = torch.exp(attn - max_attn)
sum_attn = exp_attn.sum(dim=1, keepdim=True)
attn = exp_attn / sum_attn
attn = attn.permute(0, 2, 3, 1).contiguous() # HWC2CHW like
x = (v @ attn.transpose(-2, -1)).view(B, C, H, W) + self.pe(v.reshape(B, C, H, W))
x = self.proj(x)
return x
def AAttn_forward(self, x):
# Effieicient for Bayes-e BPU
B, C, H, W = x.shape
N = H * W
qkv = self.qkv(x).flatten(2).transpose(1, 2)
if self.area > 1:
qkv = qkv.reshape(B * self.area, N // self.area, C * 3)
B, N, _ = qkv.shape
q, k, v = (qkv.view(B, N, self.num_heads, self.head_dim * 3).permute(0, 2, 3, 1).split([self.head_dim, self.head_dim, self.head_dim], dim=2))
attn = (q.transpose(-2, -1) @ k) * (self.head_dim**-0.5)
attn = attn.permute(0, 3, 1, 2).contiguous() # CHW2HWC like
max_attn = attn.max(dim=1, keepdim=True).values
exp_attn = torch.exp(attn - max_attn)
sum_attn = exp_attn.sum(dim=1, keepdim=True)
attn = exp_attn / sum_attn
attn = attn.permute(0, 2, 3, 1).contiguous() # HWC2CHW like
x = v @ attn.transpose(-2, -1)
x = x.permute(0, 3, 1, 2)
v = v.permute(0, 3, 1, 2)
if self.area > 1:
x = x.reshape(B // self.area, N * self.area, C)
v = v.reshape(B // self.area, N * self.area, C)
B, N, _ = x.shape
x = x.reshape(B, H, W, C).permute(0, 3, 1, 2).contiguous()
v = v.reshape(B, H, W, C).permute(0, 3, 1, 2).contiguous()
x = x + self.pe(v)
return self.proj(x)
def Classify_forward(self, x):
# Effieicient for Bernoulli2, Bayes, Bayes-e, Nash-{e/m/p} BPU
x = torch.cat(x, 1) if isinstance(x, list) else x
return self.linear(self.drop(self.pool(self.conv(x)).flatten(1)))
def Detect_forward(self, x):
# Effieicient for Bernoulli2, Bayes, Bayes-e, Nash-{e/m/p} BPU
result = []
for i in range(self.nl):
result.append(self.cv3[i](x[i]).permute(0, 2, 3, 1).contiguous()) # cls
result.append(self.cv2[i](x[i]).permute(0, 2, 3, 1).contiguous()) # bbox
return result
def v10Detect_forward(self, x):
# Effieicient for Bernoulli2, Bayes, Bayes-e, Nash-{e/m/p} BPU
result = []
for i in range(self.nl):
result.append(self.one2one_cv3[i](x[i]).permute(0, 2, 3, 1).contiguous()) # cls
result.append(self.one2one_cv2[i](x[i]).permute(0, 2, 3, 1).contiguous()) # bbox
return result
def Segment_forward(self, x):
# Effieicient for Bernoulli2, Bayes, Bayes-e, Nash-{e/m/p} BPU
result = []
for i in range(self.nl):
result.append(self.cv3[i](x[i]).permute(0, 2, 3, 1).contiguous()) # cls
result.append(self.cv2[i](x[i]).permute(0, 2, 3, 1).contiguous()) # bbox
result.append(self.cv4[i](x[i]).permute(0, 2, 3, 1).contiguous()) # proto weights
result.append(self.proto(x[0]).permute(0, 2, 3, 1).contiguous()) # proto mask
return result
def Pose_forward(self, x):
# Effieicient for Bernoulli2, Bayes, Bayes-e, Nash-{e/m/p} BPU
result = []
for i in range(self.nl):
result.append(self.cv3[i](x[i]).permute(0, 2, 3, 1).contiguous()) # cls
result.append(self.cv2[i](x[i]).permute(0, 2, 3, 1).contiguous()) # bbox
result.append(self.cv4[i](x[i]).permute(0, 2, 3, 1).contiguous()) # kpts
return result
def OBB_forward(self, x):
# Effieicient for Bernoulli2, Bayes, Bayes-e, Nash-{e/m/p} BPU
# TODO: Test and PostProcess Code in Model Zoo.
result = []
for i in range(self.nl):
result.append(self.cv3[i](x[i]).permute(0, 2, 3, 1).contiguous()) # cls
result.append(self.cv2[i](x[i]).permute(0, 2, 3, 1).contiguous()) # bbox
result.append(self.cv4[i](x[i]).permute(0, 2, 3, 1).contiguous()) # theta logits
return result
if __name__ == '__main__':
main()4.模型转换（onnx转bin）
4.1 下载oe包和docker（在linux系统中执行这一步）
4.1.1 下载OE包：
wget -c ftp://x5ftp@vrftp.horizon.ai/OpenExplorer/v1.2.8_release/horizon_x5_open_explorer_v1.2.8-py310_20240926.tar.gz --ftp-password=x5ftp@123$%4.1.2 下载docker镜像：
这里安装的是CPU版本，GPU版本将cpu换成gpu即可。
wget -c ftp://x5ftp@vrftp.horizon.ai/OpenExplorer/v1.2.8_release/docker_openexplorer_ubuntu_20_x5_cpu_v1.2.8.tar.gz --ftp-password=x5ftp@123$%4.2 配置OE环境
4.2.1 加载docker镜像(压缩文件无需解压)
docker load < docker_openexplorer_ubuntu_20_x5_cpu_v1.2.8.tar.gz加载后输入docker images查看镜像，并且记住TAG，等下要用
4.2.2 解压OE开发包并挂载镜像
解压OE包：
tar -xvf horizon_x5_open_explorer_v1.2.8-py310_20240926.tar.gz挂载镜像：
sudo docker run -it --rm -v {OE包路径}:/open_explorer -v {data存放路径，后续存放onnx模型以及校准数据}:/data openexplorer/ai_toolchain_ubuntu_20_x5_cpu:v1.2.8-py310示例：
docker run -it --rm -v /horizon_x5_open_explorer_v1.2.8-py310_20240926/:/open_explorer -v /data/:/data openexplorer/ai_toolchain_ubuntu_20_x5_cpu:v1.2.8-py310
文件目录示例：
/horizon_x5_open_explorer_v1.2.8-py310_20240926/为OE包解压后的路径
/data/为模型及校准数据存放路径
参考图片：
挂载成功后键入hb_mapper命令验证安装成功
4.3 转换bin模型
- 进入前面准备好的存放onnx模型及校准数据的data目录中
cd /data/
- 将转换好的onnx模型、校准图片（原图片，不做任何处理）(我的命名为calibration_images）一键YOLO转化脚本
https://github.com/D-Robotics/rdk_model_zoo/blob/main/demos/Vision/ultralytics_YOLO/x86/mapper.py放到data目录下
mapper.py如下：
import os
import argparse
import logging
import subprocess
import shutil
# 获取脚本所在目录和当前工作目录
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
WORK_DIR = os.getcwd()
try:
import cv2
except ImportError:
os.system('pip install opencv-python')
import cv2
try:
import numpy as np
except ImportError:
os.system('pip install numpy')
import numpy as np
try:
import onnxruntime as ort
except ImportError:
os.system('pip install onnxruntime')
import onnxruntime as ort
# 日志模块配置
logging.basicConfig(
level = logging.DEBUG,
format = '[%(name)s] [%(asctime)s.%(msecs)03d] [%(levelname)s] %(message)s',
datefmt='%H:%M:%S')
logger = logging.getLogger("MZOO")
def resolve_path(path, base_dir=None):
"""解析路径，支持相对路径和绝对路径"""
if os.path.isabs(path):
return path
if base_dir is None:
base_dir = WORK_DIR
return os.path.abspath(os.path.join(base_dir, path))
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--cal-images', type=str, default='./cal_images', help='*.jpg, *.png calibration images path, 20 ~ 50 pictures is OK.')
parser.add_argument('--onnx', type=str, default='./yolo11n.onnx', help='origin float onnx model path.')
parser.add_argument('--output-dir', type=str, default='.', help='output directory for converted model.')
# default below
parser.add_argument('--quantized', type=str, default="int8", help='int8 first / int16 first')
parser.add_argument('--jobs', type=int, default=16, help='model combine jobs.')
parser.add_argument('--optimize-level', type=str, default='O3', help='O0, O1, O2, O3')
parser.add_argument('--cal-sample', type=bool, default=True, help='sample calibration data or not.')
parser.add_argument('--cal-sample-num', type=int, default=20, help='num of sample calibration data.')
parser.add_argument('--save-cache', type=bool, default=False, help='remove bpu output files or not.')
# private settings
parser.add_argument('--cal', type=str, default='.calibration_data_temporary_folder', help='calibration_data_temporary_folder')
parser.add_argument('--ws', type=str, default='.temporary_workspace', help='temporary workspace')
opt = parser.parse_args()
# 首先打印原始参数
logger.info(opt)
# 解析所有路径为绝对路径
opt.onnx = resolve_path(opt.onnx)
opt.cal_images = resolve_path(opt.cal_images)
# 如果输出目录是默认值（当前目录），则设置为ONNX文件同级目录
if opt.output_dir == '.':
opt.output_dir = os.path.dirname(opt.onnx)
logger.info(f"Output directory set to ONNX file directory: {opt.output_dir}")
else:
opt.output_dir = resolve_path(opt.output_dir)
opt.ws = resolve_path(opt.ws)
logger.info(f"Resolved paths:")
logger.info(f" ONNX model: {opt.onnx}")
logger.info(f" Calibration images: {opt.cal_images}")
logger.info(f" Output directory: {opt.output_dir}")
logger.info(f" Workspace: {opt.ws}")
# check hb_mapper
try:
subprocess.run(['hb_mapper', '--version'], capture_output=True, text=True, check=True)
logger.info("hb_mapper is available and working.")
except (subprocess.CalledProcessError, FileNotFoundError):
logger.error("hb_mapper is not available.")
exit(1)
# check onnx file
session = None
width = 640
height = 640
try:
logger.info(f"Loading ONNX model from: {opt.onnx}")
if not os.path.exists(opt.onnx):
logger.error(f"ONNX file not found: {opt.onnx}")
exit(1)
session = ort.InferenceSession(opt.onnx, providers=['CPUExecutionProvider'])
inputs = session.get_inputs()
# single input check
if len(inputs) != 1:
logger.error(f"Error: Model has {len(inputs)} inputs, expected exactly 1.")
exit(1)
logger.debug("Model has a single input.")
input_tensor = inputs[0]
input_shape = input_tensor.shape
input_type = input_tensor.type
# float32 input check
if input_type != 'tensor(float)':
logger.error(f"Error: Input type is {input_type}, expected 'tensor(float)' (fp32).")
exit(1)
logger.debug("Input data type is float32 (tensor(float)).")
# NCHW input check
if len(input_shape) != 4:
logger.error(f"Error: Input shape has {len(input_shape)} dimensions, expected 4 (NCHW).")
exit(1)
logger.debug("NCHW check success.")
# get input_h, input_w
height = input_shape[2]
width = input_shape[3]
assert isinstance(height, int), "input height dtype error."
assert isinstance(width, int), "input width dtype error"
except FileNotFoundError:
logger.error(f"Error: Model file not found at '{opt.onnx}'.")
exit(1)
except Exception as e:
logger.error(f"Error analyzing ONNX model: {e}")
exit(1)
finally:
if session is not None:
del session
logger.debug("ONNX Runtime session released.")
# check cal-images folder
if not os.path.exists(opt.cal_images):
logger.error(f"cal-images folder: '{opt.cal_images}' does not exist, please check!")
exit(1)
if len(os.listdir(opt.cal_images)) == 0:
logger.error(f"cal-images folder: '{opt.cal_images}' is empty, please check!")
exit(1)
# check cal-images file
img_cnt = 0
img_names = []
for name in os.listdir(opt.cal_images):
if name.lower().endswith(('.jpg', '.png', '.jpeg')):
img_cnt += 1
img_names.append(name)
else:
logger.warning(f"cal-images folder: '{opt.cal_images}' contains non-image files, skipping: {name}")
if img_cnt == 0:
logger.error(f"cal-images folder: '{opt.cal_images}' contains no valid images, please check!")
exit(1)
if img_cnt > opt.cal_sample_num and opt.cal_sample:
sampled_indices = np.random.choice(len(img_names), size=opt.cal_sample_num, replace=False)
img_names = [img_names[i] for i in sampled_indices]
logger.info(f"Sampling enabled. Sampled {opt.cal_sample_num} images from {img_cnt} total images.")
img_cnt = len(img_names)
if img_cnt < 20:
logger.warning(f"There are {img_cnt} ( < 20 ) images in the calibration dataset, which may cause the calibration to fail.")
if img_cnt > 50:
logger.warning(f"There are {img_cnt} ( > 50 ) images in the calibration dataset, which may cost a long time to calibrate.")
# workspace folder check and setup
if os.path.exists(opt.ws) and os.path.isdir(opt.ws):
logger.info(f"Folder '{opt.ws}' exists, removing...")
try:
shutil.rmtree(opt.ws)
logger.info(f"Folder '{opt.ws}' removed successfully.")
except Exception as e:
logger.error(f"Remove folder '{opt.ws}' error: {e}")
exit(1)
try:
cal_data_dir = os.path.join(opt.ws, opt.cal)
os.makedirs(cal_data_dir, exist_ok=True)
logger.info(f"Workspace '{opt.ws}' created successfully.")
except Exception as e:
logger.error(f"Create folder '{opt.ws}' error: {e}")
exit(1)
# ensure output directory exists
os.makedirs(opt.output_dir, exist_ok=True)
# 获取模型文件名用于生成输出文件名
model_name = os.path.splitext(os.path.basename(opt.onnx))[0]
output_model_prefix = f"{model_name}_bayese_{width}x{height}_nv12"
# int16
int16_config_str = ",set_all_nodes_int16" if opt.quantized == "int16" else ""
# 使用绝对路径生成配置文件
cal_data_dir = os.path.join(opt.ws, opt.cal)
bpu_output_dir = os.path.join(opt.ws, 'bpu_model_output')
yaml_content = f'''model_parameters:
onnx_model: '{opt.onnx}'
march: "bayes-e"
layer_out_dump: False
working_dir: '{bpu_output_dir}'
output_model_file_prefix: '{output_model_prefix}'
input_parameters:
input_name: ""
input_type_rt: 'nv12'
input_type_train: 'rgb'
input_layout_train: 'NCHW'
norm_type: 'data_scale'
scale_value: 0.003921568627451
calibration_parameters:
cal_data_dir: '{cal_data_dir}'
cal_data_type: 'float32'
calibration_type: 'default'
optimization: set_Softmax_input_int8,set_Softmax_output_int8{int16_config_str}
compiler_parameters:
jobs: {opt.jobs}
compile_mode: 'latency'
debug: true
optimize_level: '{opt.optimize_level}'
'''
# 在workspace中创建配置文件
config_path = os.path.join(opt.ws, "config.yaml")
with open(config_path, "w", encoding="utf-8") as file:
file.write(yaml_content)
logger.info(f"Configuration file created: {config_path}")
# prepare calibration data
logger.info("Preparing calibration data...")
for img_name in img_names:
img_path = os.path.join(opt.cal_images, img_name)
img = cv2.imread(img_path)
if img is None:
logger.warning(f"Failed to load image: {img_path}")
continue
# 此处的前处理以ONNX的前处理为基础，总的来说是和训练时的前处理保持一致
# 如果yaml中有配置mean和scale, 则此处无须计算mean和scale.
input_tensor = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # BGR2RGB
input_tensor = cv2.resize(input_tensor, (width, height)) # resize
input_tensor = np.transpose(input_tensor, (2, 0, 1)) # HWC2CHW
input_tensor = np.expand_dims(input_tensor, axis=0).astype(np.float32) # CHW -> NCHW
dst_path = os.path.join(cal_data_dir, img_name + '.rgbchw')
input_tensor.tofile(dst_path)
logger.info("Calibration data has been successfully generated.")
# 切换到workspace目录执行转换
original_cwd = os.getcwd()
try:
os.chdir(opt.ws)
logger.info(f"Changed working directory to: {opt.ws}")
# mapper conversion
cmd = f"hb_mapper makertbin --config config.yaml --model-type onnx"
logger.info(f"Executing model conversion...")
logger.info(f"Command: {cmd}")
result = os.system(cmd)
if result != 0:
logger.error("Model conversion failed!")
exit(1)
logger.info("Model conversion completed successfully!")
# 移动输出文件到指定目录
output_bin_path = os.path.join(bpu_output_dir, f"{output_model_prefix}.bin")
final_output_path = os.path.join(opt.output_dir, f"{output_model_prefix}.bin")
logger.info(f"Looking for output file: {output_bin_path}")
if os.path.exists(output_bin_path):
shutil.move(output_bin_path, final_output_path)
logger.info(f"Output file moved to: {final_output_path}")
else:
logger.error(f"Output file not found: {output_bin_path}")
# 列出实际生成的文件
if os.path.exists(bpu_output_dir):
actual_files = os.listdir(bpu_output_dir)
logger.error(f"Files found in output directory: {actual_files}")
exit(1)
# 移动hb_mapper日志文件到输出目录
mapper_log_source = os.path.join(opt.ws, "hb_mapper_makertbin.log")
mapper_log_dest = os.path.join(opt.output_dir, "hb_mapper_makertbin.log")
if os.path.exists(mapper_log_source):
shutil.move(mapper_log_source, mapper_log_dest)
logger.info(f"hb_mapper log moved to: {mapper_log_dest}")
else:
logger.warning(f"hb_mapper log not found at: {mapper_log_source}")
finally:
# 恢复原始工作目录
os.chdir(original_cwd)
logger.info(f"Restored working directory to: {original_cwd}")
# clean the work space
logger.info("Cleaning up...")
if not opt.save_cache:
if os.path.exists(opt.ws):
shutil.rmtree(opt.ws)
logger.info("Temporary files cleaned up.")
else:
logger.info(f"Cache files preserved in: {opt.ws}")
logger.info(f"Note: hb_mapper log is available at: {mapper_log_dest}")
logger.info(f"Conversion completed successfully!")
logger.info(f"Output file: {final_output_path}")
logger.info(f"Mapper log: {mapper_log_dest}")
if __name__ == "__main__":
main()- 运行mapper.py即可获得到转化后的bin模型
python3 mapper.py --onnx [*.onnx] --cal-images [cal images path] #示例 python3 mapper.py --onnx sign-best1.onnx --cal-images /data/calibration_ima ges/
按照本文方法最后转换的bin模型结构为：检测框bboxes为[1],[3],[5]，类别clses为[0],[2],[4]
至此，bin模型已经转换成功
5.模型部署测试
测试单张图片（自行修改相关参数）：
#!/user/bin/env python
# Copyright (c) 2024，WuChao D-Robotics.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# 注意: 此程序在RDK板端端运行
# Attention: This program runs on RDK board.
import cv2
import numpy as np
from scipy.special import softmax
# from scipy.special import expit as sigmoid
from hobot_dnn import pyeasy_dnn as dnn # BSP Python API
from time import time
import argparse
import logging
# 日志模块配置
# logging configs
logging.basicConfig(
level = logging.DEBUG,
format = '[%(name)s] [%(asctime)s.%(msecs)03d] [%(levelname)s] %(message)s',
datefmt='%H:%M:%S')
logger = logging.getLogger("RDK_YOLO")
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--model-path', type=str, default='models/yolo11n_detect_bayese_640x640_nv12.bin',
help="""Path to BPU Quantized *.bin Model.
RDK X3(Module): Bernoulli2.
RDK Ultra: Bayes.
RDK X5(Module): Bayes-e.
RDK S100: Nash-e.
RDK S100P: Nash-m.""")
parser.add_argument('--test-img', type=str, default='../../../resource/assets/bus.jpg', help='Path to Load Test Image.')
parser.add_argument('--img-save-path', type=str, default='jupyter_result.jpg', help='Path to Load Test Image.')
parser.add_argument('--classes-num', type=int, default=80, help='Classes Num to Detect.')
parser.add_argument('--reg', type=int, default=16, help='DFL reg layer.')
parser.add_argument('--iou-thres', type=float, default=0.45, help='IoU threshold.')
parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold.')
opt = parser.parse_args()
logger.info(opt)
# 实例化
model = YOLO11_Detect(opt.model_path, opt.conf_thres, opt.iou_thres)
# 读图
img = cv2.imread(opt.test_img)
# 准备输入数据
input_tensor = model.bgr2nv12(img)
# 推理
outputs = model.c2numpy(model.forward(input_tensor))
# 后处理
ids, scores, bboxes = model.postProcess(outputs)
# 渲染
logger.info("\033[1;32m" + "Draw Results: " + "\033[0m")
for class_id, score, bbox in zip(ids, scores, bboxes):
x1, y1, x2, y2 = bbox
logger.info("(%d, %d, %d, %d) -> %s: %.2f"%(x1,y1,x2,y2, coco_names[class_id], score))
draw_detection(img, (x1, y1, x2, y2), score, class_id)
# 保存结果
cv2.imwrite(opt.img_save_path, img)
logger.info("\033[1;32m" + f"saved in path: \"./{opt.img_save_path}\"" + "\033[0m")
class BaseModel:
def __init__(
self,
model_file: str
) -> None:
# 加载BPU的bin模型, 打印相关参数
# Load the quantized *.bin model and print its parameters
try:
begin_time = time()
self.quantize_model = dnn.load(model_file)
logger.debug("\033[1;31m" + "Load D-Robotics Quantize model time = %.2f ms"%(1000*(time() - begin_time)) + "\033[0m")
except Exception as e:
logger.error("❌ Failed to load model file: %s"%(model_file))
logger.error("You can download the model file from the following docs: ./models/download.md")
logger.error(e)
exit(1)
logger.info("\033[1;32m" + "-> input tensors" + "\033[0m")
for i, quantize_input in enumerate(self.quantize_model[0].inputs):
logger.info(f"intput[{i}], name={quantize_input.name}, type={quantize_input.properties.dtype}, shape={quantize_input.properties.shape}")
logger.info("\033[1;32m" + "-> output tensors" + "\033[0m")
for i, quantize_input in enumerate(self.quantize_model[0].outputs):
logger.info(f"output[{i}], name={quantize_input.name}, type={quantize_input.properties.dtype}, shape={quantize_input.properties.shape}")
self.model_input_height, self.model_input_weight = self.quantize_model[0].inputs[0].properties.shape[2:4]
def resizer(self, img: np.ndarray)->np.ndarray:
img_h, img_w = img.shape[0:2]
self.y_scale, self.x_scale = img_h/self.model_input_height, img_w/self.model_input_weight
return cv2.resize(img, (self.model_input_height, self.model_input_weight), interpolation=cv2.INTER_NEAREST) # 利用resize重新开辟内存
def preprocess(self, img: np.ndarray)->np.array:
"""
Preprocesses an input image to prepare it for model inference.
Args:
img (np.ndarray): The input image in BGR format as a NumPy array.
Returns:
np.array: The preprocessed image tensor in NCHW format ready for model input.
Procedure:
1. Resizes the image to a specified dimension (`input_image_size`) using nearest neighbor interpolation.
2. Converts the image color space from BGR to RGB.
3. Transposes the dimensions of the image tensor to channel-first order (CHW).
4. Adds a batch dimension, thus conforming to the NCHW format expected by many models.
Note: Normalization to [0, 1] is assumed to be handled elsewhere based on configuration.
"""
begin_time = time()
input_tensor = self.resizer(img)
input_tensor = cv2.cvtColor(input_tensor, cv2.COLOR_BGR2RGB)
# input_tensor = np.array(input_tensor) / 255.0 # yaml文件中已经配置前处理
input_tensor = np.transpose(input_tensor, (2, 0, 1))
input_tensor = np.expand_dims(input_tensor, axis=0).astype(np.uint8) # NCHW
logger.debug("\033[1;31m" + f"pre process time = {1000*(time() - begin_time):.2f} ms" + "\033[0m")
return input_tensor
def bgr2nv12(self, bgr_img: np.ndarray) -> np.ndarray:
"""
Convert a BGR image to the NV12 format.
NV12 is a common video encoding format where the Y component (luminance) is full resolution,
and the UV components (chrominance) are half-resolution and interleaved. This function first
converts the BGR image to YUV 4:2:0 planar format, then rearranges the UV components to fit
the NV12 format.
Parameters:
bgr_img (np.ndarray): The input BGR image array.
Returns:
np.ndarray: The converted NV12 format image array.
"""
begin_time = time()
bgr_img = self.resizer(bgr_img)
height, width = bgr_img.shape[0], bgr_img.shape[1]
area = height * width
yuv420p = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2YUV_I420).reshape((area * 3 // 2,))
y = yuv420p[:area]
uv_planar = yuv420p[area:].reshape((2, area // 4))
uv_packed = uv_planar.transpose((1, 0)).reshape((area // 2,))
nv12 = np.zeros_like(yuv420p)
nv12[:height * width] = y
nv12[height * width:] = uv_packed
logger.debug("\033[1;31m" + f"bgr8 to nv12 time = {1000*(time() - begin_time):.2f} ms" + "\033[0m")
return nv12
def forward(self, input_tensor: np.array) -> list[dnn.pyDNNTensor]:
begin_time = time()
quantize_outputs = self.quantize_model[0].forward(input_tensor)
logger.debug("\033[1;31m" + f"forward time = {1000*(time() - begin_time):.2f} ms" + "\033[0m")
return quantize_outputs
def c2numpy(self, outputs) -> list[np.array]:
begin_time = time()
outputs = [dnnTensor.buffer for dnnTensor in outputs]
logger.debug("\033[1;31m" + f"c to numpy time = {1000*(time() - begin_time):.2f} ms" + "\033[0m")
return outputs
class YOLO11_Detect(BaseModel):
def __init__(self,
model_file: str,
conf: float,
iou: float
):
super().__init__(model_file)
# 将反量化系数准备好, 只需要准备一次
# prepare the quantize scale, just need to generate once
self.s_bboxes_scale = self.quantize_model[0].outputs[0].properties.scale_data[np.newaxis, :]
self.m_bboxes_scale = self.quantize_model[0].outputs[1].properties.scale_data[np.newaxis, :]
self.l_bboxes_scale = self.quantize_model[0].outputs[2].properties.scale_data[np.newaxis, :]
logger.info(f"{self.s_bboxes_scale.shape=}, {self.m_bboxes_scale.shape=}, {self.l_bboxes_scale.shape=}")
# DFL求期望的系数, 只需要生成一次
# DFL calculates the expected coefficients, which only needs to be generated once.
self.weights_static = np.array([i for i in range(16)]).astype(np.float32)[np.newaxis, np.newaxis, :]
logger.info(f"{self.weights_static.shape = }")
# anchors, 只需要生成一次
self.s_anchor = np.stack([np.tile(np.linspace(0.5, 79.5, 80), reps=80),
np.repeat(np.arange(0.5, 80.5, 1), 80)], axis=0).transpose(1,0)
self.m_anchor = np.stack([np.tile(np.linspace(0.5, 39.5, 40), reps=40),
np.repeat(np.arange(0.5, 40.5, 1), 40)], axis=0).transpose(1,0)
self.l_anchor = np.stack([np.tile(np.linspace(0.5, 19.5, 20), reps=20),
np.repeat(np.arange(0.5, 20.5, 1), 20)], axis=0).transpose(1,0)
logger.info(f"{self.s_anchor.shape = }, {self.m_anchor.shape = }, {self.l_anchor.shape = }")
# 输入图像大小, 一些阈值, 提前计算好
self.input_image_size = 640
self.conf = conf
self.iou = iou
self.conf_inverse = -np.log(1/conf - 1)
logger.info("iou threshol = %.2f, conf threshol = %.2f"%(iou, conf))
logger.info("sigmoid_inverse threshol = %.2f"%self.conf_inverse)
def postProcess(self, outputs: list[np.ndarray]) -> tuple[list]:
begin_time = time()
# reshape
s_bboxes = outputs[1].reshape(-1, 64)
m_bboxes = outputs[3].reshape(-1, 64)
l_bboxes = outputs[5].reshape(-1, 64)
s_clses = outputs[0].reshape(-1, 80)
m_clses = outputs[2].reshape(-1, 80)
l_clses = outputs[4].reshape(-1, 80)
# classify: 利用numpy向量化操作完成阈值筛选(优化版 2.0)
s_max_scores = np.max(s_clses, axis=1)
s_valid_indices = np.flatnonzero(s_max_scores >= self.conf_inverse) # 得到大于阈值分数的索引，此时为小数字
s_ids = np.argmax(s_clses[s_valid_indices, : ], axis=1)
s_scores = s_max_scores[s_valid_indices]
m_max_scores = np.max(m_clses, axis=1)
m_valid_indices = np.flatnonzero(m_max_scores >= self.conf_inverse) # 得到大于阈值分数的索引，此时为小数字
m_ids = np.argmax(m_clses[m_valid_indices, : ], axis=1)
m_scores = m_max_scores[m_valid_indices]
l_max_scores = np.max(l_clses, axis=1)
l_valid_indices = np.flatnonzero(l_max_scores >= self.conf_inverse) # 得到大于阈值分数的索引，此时为小数字
l_ids = np.argmax(l_clses[l_valid_indices, : ], axis=1)
l_scores = l_max_scores[l_valid_indices]
# 3个Classify分类分支：Sigmoid计算
s_scores = 1 / (1 + np.exp(-s_scores))
m_scores = 1 / (1 + np.exp(-m_scores))
l_scores = 1 / (1 + np.exp(-l_scores))
# 3个Bounding Box分支：筛选
s_bboxes_float32 = s_bboxes[s_valid_indices,:]#.astype(np.float32) * self.s_bboxes_scale
m_bboxes_float32 = m_bboxes[m_valid_indices,:]#.astype(np.float32) * self.m_bboxes_scale
l_bboxes_float32 = l_bboxes[l_valid_indices,:]#.astype(np.float32) * self.l_bboxes_scale
# 3个Bounding Box分支：dist2bbox (ltrb2xyxy)
s_ltrb_indices = np.sum(softmax(s_bboxes_float32.reshape(-1, 4, 16), axis=2) * self.weights_static, axis=2)
s_anchor_indices = self.s_anchor[s_valid_indices, :]
s_x1y1 = s_anchor_indices - s_ltrb_indices[:, 0:2]
s_x2y2 = s_anchor_indices + s_ltrb_indices[:, 2:4]
s_dbboxes = np.hstack([s_x1y1, s_x2y2])*8
m_ltrb_indices = np.sum(softmax(m_bboxes_float32.reshape(-1, 4, 16), axis=2) * self.weights_static, axis=2)
m_anchor_indices = self.m_anchor[m_valid_indices, :]
m_x1y1 = m_anchor_indices - m_ltrb_indices[:, 0:2]
m_x2y2 = m_anchor_indices + m_ltrb_indices[:, 2:4]
m_dbboxes = np.hstack([m_x1y1, m_x2y2])*16
l_ltrb_indices = np.sum(softmax(l_bboxes_float32.reshape(-1, 4, 16), axis=2) * self.weights_static, axis=2)
l_anchor_indices = self.l_anchor[l_valid_indices,:]
l_x1y1 = l_anchor_indices - l_ltrb_indices[:, 0:2]
l_x2y2 = l_anchor_indices + l_ltrb_indices[:, 2:4]
l_dbboxes = np.hstack([l_x1y1, l_x2y2])*32
# 大中小特征层阈值筛选结果拼接
dbboxes = np.concatenate((s_dbboxes, m_dbboxes, l_dbboxes), axis=0)
scores = np.concatenate((s_scores, m_scores, l_scores), axis=0)
ids = np.concatenate((s_ids, m_ids, l_ids), axis=0)
# nms
indices = cv2.dnn.NMSBoxes(dbboxes, scores, self.conf, self.iou)
# 还原到原始的img尺度
bboxes = dbboxes[indices] * np.array([self.x_scale, self.y_scale, self.x_scale, self.y_scale])
bboxes = bboxes.astype(np.int32)
logger.debug("\033[1;31m" + f"Post Process time = {1000*(time() - begin_time):.2f} ms" + "\033[0m")
return ids[indices], scores[indices], bboxes
coco_names = [
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle",
"wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
"broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed",
"dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven",
"toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
]
rdk_colors = [
(56, 56, 255), (151, 157, 255), (31, 112, 255), (29, 178, 255),(49, 210, 207), (10, 249, 72), (23, 204, 146), (134, 219, 61),
(52, 147, 26), (187, 212, 0), (168, 153, 44), (255, 194, 0),(147, 69, 52), (255, 115, 100), (236, 24, 0), (255, 56, 132),
(133, 0, 82), (255, 56, 203), (200, 149, 255), (199, 55, 255)]
def draw_detection(img: np.array,
bbox: tuple[int, int, int, int],
score: float,
class_id: int) -> None:
"""
Draws a detection bounding box and label on the image.
Parameters:
img (np.array): The input image.
bbox (tuple[int, int, int, int]): A tuple containing the bounding box coordinates (x1, y1, x2, y2).
score (float): The detection score of the object.
class_id (int): The class ID of the detected object.
"""
x1, y1, x2, y2 = bbox
color = rdk_colors[class_id%20]
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
label = f"{coco_names[class_id]}: {score:.2f}"
(label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
label_x, label_y = x1, y1 - 10 if y1 - 10 > label_height else y1 + 10
cv2.rectangle(
img, (label_x, label_y - label_height), (label_x + label_width, label_y + label_height), color, cv2.FILLED
)
cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
if __name__ == "__main__":
main()修改：
-
parser.add_argument('--classes-num', type=int, default=你的模型检测的数量, help='Classes Num to Detect.') -
coco_names = [改成你的检测信息] -
s_bboxes = outputs[1].reshape(-1, 64) m_bboxes = outputs[3].reshape(-1, 64) l_bboxes = outputs[5].reshape(-1, 64) s_clses = outputs[0].reshape(-1, 你的识别项数量) m_clses = outputs[2].reshape(-1, 你的识别项数量) l_clses = outputs[4].reshape(-1, 你的识别项数量)
如果有遗漏的地方需要补充的还请指出!
参考的文章
更多推荐
0
0- 0
已为社区贡献1条内容
所有评论(0)