HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.3.65 Rockchip RKNN Integration for Ultralytics YOLO models (#16308)

Browse source 
Signed-off-by: Francesco Mattioli <Francesco.mttl@gmail.com>
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Burhan <62214284+Burhan-Q@users.noreply.github.com>
Co-authored-by: Lakshantha Dissanayake <lakshantha@ultralytics.com>
Co-authored-by: Burhan <Burhan-Q@users.noreply.github.com>
Co-authored-by: Laughing-q <1185102784@qq.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Ultralytics Assistant <135830346+UltralyticsAssistant@users.noreply.github.com>
Co-authored-by: Lakshantha Dissanayake <lakshanthad@yahoo.com>
Co-authored-by: Francesco Mattioli <Francesco.mttl@gmail.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ivor Zhu 2025-01-20 20:25:54 -05:00 • committed by GitHub
parent 617dea8e25
commit b5e0cee943
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
41 changed files with 390 additions and 118 deletions
Download patch file Download diff file Expand all files Collapse all files

31
ultralytics/nn/autobackend.py
View file

@ -14,7 +14,7 @@ import torch.nn as nn
from PIL import Image
from ultralytics.utils import ARM64, IS_JETSON, IS_RASPBERRYPI, LINUX, LOGGER, PYTHON_VERSION, ROOT, yaml_load
from ultralytics.utils.checks import check_requirements, check_suffix, check_version, check_yaml
from ultralytics.utils.checks import check_requirements, check_suffix, check_version, check_yaml, is_rockchip
from ultralytics.utils.downloads import attempt_download_asset, is_url
@ -60,7 +60,7 @@ class AutoBackend(nn.Module):
Supported Formats and Naming Conventions:
| Format | File Suffix |
|-----------------------|-------------------|
| --------------------- | ----------------- |
| PyTorch | *.pt |
| TorchScript | *.torchscript |
| ONNX Runtime | *.onnx |
@ -75,6 +75,8 @@ class AutoBackend(nn.Module):
| PaddlePaddle | *_paddle_model/ |
| MNN | *.mnn |
| NCNN | *_ncnn_model/ |
| IMX | *_imx_model/ |
| RKNN | *_rknn_model/ |
This class offers dynamic backend switching capabilities based on the input model format, making it easier to deploy
models across various platforms.
@ -124,10 +126,11 @@ class AutoBackend(nn.Module):
mnn,
ncnn,
imx,
rknn,
triton,
) = self._model_type(w)
fp16 &= pt or jit or onnx or xml or engine or nn_module or triton # FP16
nhwc = coreml or saved_model or pb or tflite or edgetpu # BHWC formats (vs torch BCWH)
nhwc = coreml or saved_model or pb or tflite or edgetpu or rknn # BHWC formats (vs torch BCWH)
stride = 32 # default stride
model, metadata, task = None, None, None
@ -466,6 +469,22 @@ class AutoBackend(nn.Module):
model = TritonRemoteModel(w)
metadata = model.metadata
# RKNN
elif rknn:
if not is_rockchip():
raise OSError("RKNN inference is only supported on Rockchip devices.")
LOGGER.info(f"Loading {w} for RKNN inference...")
check_requirements("rknn-toolkit-lite2")
from rknnlite.api import RKNNLite
w = Path(w)
if not w.is_file(): # if not *.rknn
w = next(w.rglob("*.rknn")) # get *.rknn file from *_rknn_model dir
rknn_model = RKNNLite()
rknn_model.load_rknn(w)
ret = rknn_model.init_runtime()
metadata = Path(w).parent / "metadata.yaml"
# Any other format (unsupported)
else:
from ultralytics.engine.exporter import export_formats
@ -652,6 +671,12 @@ class AutoBackend(nn.Module):
im = im.cpu().numpy() # torch to numpy
y = self.model(im)
# RKNN
elif self.rknn:
im = (im.cpu().numpy() * 255).astype("uint8")
im = im if isinstance(im, (list, tuple)) else [im]
y = self.rknn_model.inference(inputs=im)
# TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
else:
im = im.cpu().numpy()