Team2890/PhotonVision
4
0
Fork 0
mirror of https://github.com/PhotonVision/photonvision synced 2026-06-19 00:41:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge rknn conversion scripts into notebook (#2157)

Browse Source
This commit is contained in:
Rikhil Chilka
2025-10-24 00:41:49 -04:00
committed by GitHub
parent 8e88a9a780
commit 3cbac8117e
6 changed files with 618 additions and 721 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/source/docs/objectDetection/opi.md
View File

@@ -14,6 +14,6 @@ PhotonVision currently ONLY supports 640x640 Ultralytics YOLOv5, YOLOv8, and YOL
Only quantized models are supported, so take care when exporting to select the option for quantization.
:::
PhotonVision now ships with a {{ '[Python Notebook](https://github.com/PhotonVision/photonvision/blob/{}/scripts/rknn-convert-tool/rknn_conversion.ipynb)'.format(git_tag_ref) }} that you can use in [Google Colab](https://colab.research.google.com) or in a local environment. In Google Colab, you can simply paste the PhotonVision GitHub URL into the "GitHub" tab and select the `rknn_conversion.ipynb` notebook without needing to manually download anything.
PhotonVision now ships with a {{ '[Python Notebook](https://github.com/PhotonVision/photonvision/blob/{}/scripts/rknn_conversion.ipynb)'.format(git_tag_ref) }} that you can use in [Google Colab](https://colab.research.google.com) or in a local environment. In Google Colab, you can simply paste the PhotonVision GitHub URL into the "GitHub" tab and select the `rknn_conversion.ipynb` notebook without needing to manually download anything.
Please ensure that the model you are attempting to convert is among the {ref}`supported models <docs/objectDetection/opi:Supported Models>` and using the PyTorch format.

7
scripts/rknn-convert-tool/MAINTAINERS.md
View File

@@ -1,7 +0,0 @@
# Notebook
In the first cell of the RKNN conversion notebook, the installation script uses a structured list of dictionaries to define the download URLs and filenames for required scripts. Each dictionary includes a `url` (a permalink to a specific commit) and the corresponding `filename`.
Please ensure that all URLs in this array use permalinks—that is, links pointing to a specific commit hash rather than a branch name (e.g., main). This guarantees that the correct version of each script is always fetched, and prevents unexpected changes if the repository is updated in the future.
You typically won’t need to update these permalinks unless one of the referenced scripts is modified. In that case, update the commit hash in the URLs accordingly.

182
scripts/rknn-convert-tool/create_onnx.py
View File

@@ -1,182 +0,0 @@
import argparse
import os.path
import subprocess
import sys
# This will work for all models that don't use anchors (e.g. all YOLO models except YOLOv5/v7)
# This includes YOLOv5u
yolo_non_anchor_repo = "https://github.com/airockchip/ultralytics_yolo11"
# For original YOLOv5 models
yolov5_repo = "https://github.com/airockchip/yolov5"
valid_yolo_versions = ["yolov5", "yolov8", "yolov11"]
comma_sep_yolo_versions = ", ".join(valid_yolo_versions)
ultralytics_folder_name_yolov5 = "airockchip_yolo_pkg_yolov5"
ultralytics_default_folder_name = "airockchip_yolo_pkg"
bad_model_msg = """
This is usually due to passing in the wrong model version.
Please make sure you have the right model version and try again.
"""
def print_bad_model_msg(cause):
print(f"{cause}{bad_model_msg}")
def run_and_exit_with_error(cmd, error_msg, enable_error_output=True):
try:
if enable_error_output:
subprocess.run(
cmd,
stderr=subprocess.STDOUT,
stdout=subprocess.PIPE,
universal_newlines=True,
).check_returncode()
else:
subprocess.run(cmd).check_returncode()
except subprocess.CalledProcessError as e:
print(error_msg)
if enable_error_output:
print(e.stdout)
sys.exit(1)
def check_git_installed():
run_and_exit_with_error(
["git", "--version"],
"""Git is not installed or not found in your PATH.
Please install Git from https://git-scm.com/downloads and try again.""",
)
def check_or_clone_rockchip_repo(repo_url, repo_name=ultralytics_default_folder_name):
if os.path.exists(repo_name):
print(
f'Existing Rockchip repo "{repo_name}" detected, skipping installation...'
)
else:
print(f'Cloning Rockchip repo to "{repo_name}"')
run_and_exit_with_error(
["git", "clone", repo_url, repo_name],
"Failed to clone Rockchip repo, please see error output",
)
def run_pip_install_or_else_exit(args):
print("Running pip install...")
run_and_exit_with_error(
["pip", "install"] + args,
"Pip install rockchip repo failed, please see error output",
)
def run_onnx_conversion_yolov5(model_path):
check_or_clone_rockchip_repo(yolov5_repo, ultralytics_folder_name_yolov5)
run_pip_install_or_else_exit(
[
"-r",
os.path.join(ultralytics_folder_name_yolov5, "requirements.txt"),
"torch<2.6.0",
"onnx==1.18.0",
"onnxscript",
]
)
model_abspath = os.path.abspath(model_path)
try:
subprocess.run(
[
"python",
f"{ultralytics_folder_name_yolov5}/export.py",
"--weights",
model_abspath,
"--rknpu",
"--include",
"onnx",
],
stderr=subprocess.STDOUT,
stdout=subprocess.PIPE,
universal_newlines=True,
).check_returncode()
except subprocess.CalledProcessError as e:
print("Failed to run YOLOv5 export, please see error output")
if "ModuleNotFoundError" in e.stdout and "ultralytics" in e.stdout:
print_bad_model_msg(
"It seems the YOLOv5 repo could not find an ultralytics installation."
)
elif "AttributeError" in e.stdout and "_register_detect_seperate" in e.stdout:
print_bad_model_msg("It seems that you received a model attribute error.")
else:
print("Unknown Error when converting:")
print(e.stdout)
sys.exit(1)
def run_onnx_conversion_no_anchor(model_path):
check_or_clone_rockchip_repo(yolo_non_anchor_repo)
run_pip_install_or_else_exit(
["-e", ultralytics_default_folder_name, "onnx==1.18.0", "onnxscript"]
)
sys.path.insert(0, os.path.abspath(ultralytics_default_folder_name))
model_abs_path = os.path.abspath(model_path)
from ultralytics import YOLO
try:
model = YOLO(model_abs_path)
model.export(format="rknn")
except TypeError as e:
if "originally trained" in str(e):
print_bad_model_msg(
"Ultralytics has detected that this model is a YOLOv5 model."
)
else:
raise e
sys.exit(1)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Generate valid ONNX file for yolo model"
)
parser.add_argument(
"-m",
"--model_path",
required=True,
help=(f"Path to YOLO model"),
)
parser.add_argument(
"-v",
"--version",
required=True,
choices=valid_yolo_versions,
help=(f"Model version, must be one of: {comma_sep_yolo_versions}"),
)
args = parser.parse_args()
check_git_installed()
try:
if args.version.lower() == "yolov5":
run_onnx_conversion_yolov5(args.model_path)
else:
run_onnx_conversion_no_anchor(args.model_path)
print(
"Model export finished. Please use the generated ONNX file to convert to RKNN."
)
except SystemExit:
print("Model export failed. Please see output above.")

221
scripts/rknn-convert-tool/create_rknn.py
View File

@@ -1,221 +0,0 @@
import argparse
import os
import random
import sys
from rknn.api import RKNN
image_extensions = (".jpg", ".jpeg", ".png", ".bmp", ".gif", ".tiff", ".webp")
DEFAULT_PLATFORM = "rk3588"
def list_img_dir(img_dir):
return [
os.path.abspath(os.path.join(img_dir, f))
for f in os.listdir(img_dir)
if f.lower().endswith(image_extensions)
]
def sample_imgs(num, img_list):
if len(img_list) < num:
return img_list
else:
return random.sample(img_list, num)
def get_image_list_from_dataset(num_imgs, yaml_dir):
print(f"Dataset detected with {yaml_dir} file")
img_raw_paths = []
with open(yaml_dir, "r") as yaml_file:
for line in yaml_file:
line = line.strip()
if (
line.startswith("train:")
or line.startswith("val:")
or line.startswith("test:")
):
img_raw_paths.append(line.split(":", 1)[1].strip())
no_yaml_dir = yaml_dir.replace(
"data.yaml", "dummy_dir"
) # data.yaml sets dirs one level up
img_set_paths = []
for img_raw_path in img_raw_paths:
p = (
img_raw_path
if os.path.isabs(img_raw_path)
else os.path.realpath(os.path.join(no_yaml_dir, img_raw_path))
)
if os.path.exists(p):
img_set_paths.append(p)
if len(img_set_paths) < 1:
return None
all_imgs = [list_img_dir(path) for path in img_set_paths]
for imgs in all_imgs:
print(len(imgs))
total_imgs = sum(len(group) for group in all_imgs)
sampled_imgs = [
sample_imgs(round((len(group) / total_imgs) * num_imgs), group)
for group in all_imgs
]
return [img for group in sampled_imgs for img in group]
def get_image_list_from_img_dir(num_imgs, img_dir):
return sample_imgs(num_imgs, list_img_dir(img_dir))
def get_image_list(num_imgs, image_dir):
yaml_path = os.path.join(image_dir, "data.yaml")
if os.path.exists(yaml_path):
return get_image_list_from_dataset(num_imgs, yaml_path)
else:
return get_image_list_from_img_dir(num_imgs, image_dir)
def run_rknn_conversion(
img_list_txt, disable_quant, model_path, rknn_output, verbose_logging
):
rknn = RKNN(
verbose=verbose_logging,
verbose_file=("rknn_convert.log" if verbose_logging else None),
)
rknn.config(
mean_values=[[0, 0, 0]],
std_values=[[255, 255, 255]],
target_platform=DEFAULT_PLATFORM,
)
print("Attempted RKNN load")
ret = rknn.load_onnx(model=model_path)
if ret != 0:
print("Loading model failed!")
exit(ret)
print("Attempting RKNN build")
ret = rknn.build(do_quantization=(not disable_quant), dataset=img_list_txt)
if ret != 0:
print("Building model failed!")
exit(ret)
print("Build succeeded! Starting export...")
ret = rknn.export_rknn(rknn_output)
if ret != 0:
print("Exporting model failed!")
exit(ret)
print("Finished export!")
# Release
rknn.release()
print(f'Your model is in "{rknn_output}" and ready to use!')
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Generate valid ONNX file for yolo model"
)
parser.add_argument(
"-ni",
"--num_imgs",
type=int,
default=300,
help="Number of images to use for calibration (default: 300)",
)
parser.add_argument(
"-d",
"--img_dir",
help="Directory where your dataset is located (must have data.yaml), or images are located",
)
parser.add_argument(
"-m",
"--model_path",
required=True,
help=(f"Path to generated ONNX model"),
)
parser.add_argument(
"-dq",
"--disable_quantize",
action="store_true",
help="Whether to skip quantization (default: False)",
)
parser.add_argument(
"-o",
"--rknn_output",
default="out.rknn",
help="Where the rknn model should be outputted (default: ./out.rknn)",
)
parser.add_argument(
"-ds",
"--img_dataset_txt",
default="imgs.txt",
help="Where the list of images used for quantization should be outputted (default: ./imgs.txt)",
)
parser.add_argument(
"-vb",
"--verbose",
action="store_true",
help="Whether to enable verbose logging",
)
args = parser.parse_args()
if not args.rknn_output.endswith(".rknn"):
print("RKNN output path must end in .rknn!")
sys.exit(1)
if not args.disable_quantize:
if args.img_dir == None or len(args.img_dir) < 1:
print(f"Must specify list of images to use with --img_dir")
sys.exit(1)
img_dir_abs = os.path.abspath(args.img_dir)
img_list = get_image_list(args.num_imgs, img_dir_abs)
img_list_len = 0 if img_list is None else len(img_list)
if img_list_len == 0:
print(f"No images found in {img_dir_abs}")
sys.exit(1)
elif img_list_len < args.num_imgs:
print(
f"Not enough images in your dataset/directory, you have {img_list_len} images, but need {args.num_imgs}"
)
sys.exit(1)
if not args.img_dataset_txt.endswith(".txt"):
print(f"Image dataset text file path must end in .txt")
sys.exit(1)
with open(args.img_dataset_txt, "w") as set_file:
set_file.writelines(f"{img}\n" for img in img_list)
try:
run_rknn_conversion(
args.img_dataset_txt,
args.disable_quantize,
args.model_path,
args.rknn_output,
args.verbose,
)
except SystemExit:
print("RKNN Conversion failed, see output above")

310
scripts/rknn-convert-tool/rknn_conversion.ipynb
View File

@@ -1,310 +0,0 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "bb5367ce",
"metadata": {},
"source": [
"# RKNN Conversion Guide\n",
"\n",
"----------------------------\n",
"\n",
"### Before you start\n",
"\n",
"If you are not using Google Colab, it is recommended to create a separate [Python virtual environment](https://docs.python.org/3/library/venv.html) before you run the scripts or the Python notebook from this project. This ensures that packages installed for the conversion process do not conflict with other packages you may already have setup."
]
},
{
"cell_type": "markdown",
"id": "5f42d0a144caceb6",
"metadata": {},
"source": [
"### Preinstallation\n",
"\n",
"This notebook requires the use of external Python scripts. Please run the installation script below to import these external scripts if you do not have them already."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7903189e",
"metadata": {},
"outputs": [],
"source": [
"import subprocess\n",
"\n",
"# Define scripts with URLs and inferred filenames\n",
"# DO NOT modify the filenames\n",
"scripts = [\n",
" {\n",
" \"url\": \"https://raw.githubusercontent.com/PhotonVision/photonvision/2bf166bc3f377fa8af9d9d38ee46e0db978a4036/scripts/rknn-convert-tool/create_onnx.py\",\n",
" \"filename\": \"create_onnx.py\" # CREATE_ONNX_SCRIPT\n",
" },\n",
" {\n",
" \"url\": \"https://raw.githubusercontent.com/PhotonVision/photonvision/2bf166bc3f377fa8af9d9d38ee46e0db978a4036/scripts/rknn-convert-tool/create_rknn.py\",\n",
" \"filename\": \"create_rknn.py\" # CREATE_RKNN_SCRIPT\n",
" }\n",
"]\n",
"\n",
"# Download each script\n",
"for script in scripts:\n",
" try:\n",
" subprocess.run([\"wget\", script[\"url\"], \"-O\", script[\"filename\"]]).check_returncode()\n",
" print(f\"Successfully downloaded: {script['filename']}\")\n",
" except subprocess.CalledProcessError as e:\n",
" print(f\"Failed to download script from URL: {script['url']}\")\n",
" print(e)\n"
]
},
{
"cell_type": "markdown",
"id": "d68be4aba4d3022b",
"metadata": {},
"source": [
"#### *Numpy Fix* - Important for Google Colab Users\n",
"\n",
"Google Colab comes with an incompatible version of Numpy installed. To fix this, please run the following cells below and **restart your session** when prompted."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "de0310a3e4401233",
"metadata": {},
"outputs": [],
"source": [
"%pip uninstall numpy -y\n",
"%pip install \"numpy>=1.23.0,<2.0.0\""
]
},
{
"cell_type": "markdown",
"id": "d498ed79",
"metadata": {},
"source": [
"### Step 1: Convert to ONNX\n",
"\n",
"To convert to ONNX, simply run the `create_onnx.py` script, providing the path to your model weights and specifying the model version, as shown below."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0659e15f",
"metadata": {},
"outputs": [],
"source": [
"# where version is either yolov5, yolov8, or yolov11, and model_path is the path to your weights file (.pt)\n",
"%run create_onnx.py --version yolov8 --model_path myyolov8model.pt"
]
},
{
"cell_type": "markdown",
"id": "86ff07e6",
"metadata": {},
"source": [
"### Step 2: Download RKNN API\n",
"\n",
"You can either use `pip` below to automatically detect and install the correct RKNN API Python library for you, or install it manually.\n",
"\n",
"#### Automatic installation\n",
"\n",
"Please run `pip` below. If it does not work, refer to the instructions for manual installation. You may need to restart your session after running the command below.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7ec11f96",
"metadata": {},
"outputs": [],
"source": [
"%pip install rknn-toolkit2"
]
},
{
"cell_type": "markdown",
"id": "8b57fe4d",
"metadata": {},
"source": [
"#### Manual Installation (If Automatic Installation Fails)\n",
"Visit the [RKNN Toolkit 2](https://github.com/airockchip/rknn-toolkit2) Github repository, then click on rknn-toolkit2, followed by packages.\n",
"If you are running an x86_64 CPU (e.g., most Intel and AMD processors), select that option; otherwise, choose arm64 for ARM-based computers (e.g., M-series Macs or Snapdragon processors). If you're unsure which CPU you're using, check your system settings for processor architecture information.\n",
"\n",
"Once you've selected the correct CPU architecture, you'll see multiple packages. The file names will look something like:\n",
"`rknn_toolkit2-2.3.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl`.\n",
"The numbers after cp correspond to your Python version. For example, if you're using Python 3.10, look for a package with cp310 in the name. For Python 3.8, look for cp38; for Python 3.7, cp37, and so on.\n",
"\n",
"Once you've found the correct package, click the \"Raw\" button to download the .whl file. Then, run the following command in your terminal, replacing rknn_toolkit2.whl with the actual path to the file you downloaded:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7414b120",
"metadata": {},
"outputs": [],
"source": [
"%pip install rknn_toolkit2.whl"
]
},
{
"cell_type": "markdown",
"id": "c1db5ef0",
"metadata": {},
"source": [
"### Step 3: Convert to RKNN\n",
"\n",
"Please review the notes about quantization before running the RKNN conversion."
]
},
{
"cell_type": "markdown",
"id": "5e56b2f64bf6e85f",
"metadata": {},
"source": [
"### Quantization\n",
"\n",
"When performing quantization, it is critical to provide representative images of the objects or scenes you are trying to detect. These images are used to calibrate the model’s internal activations and greatly influence the final performance.\n",
"\n",
"It is recommended to use 300–500 representative images that reflect the real-world input your model will encounter. As the old saying goes, it’s quality over quantity — having a diverse, relevant set matters more than simply having many images.\n",
"\n",
"Quantization will cause some loss in model accuracy. However, if your calibration images are chosen wisely, this accuracy drop should be minimal and acceptable. If the sampled images are too uniform or unrelated, your quantized model's performance may worsen significantly.\n",
"\n",
"The script will automatically sample representative images randomly from the provided dataset. While this usually works well, please verify that the dataset contains diverse and relevant examples of your target objects. As a reminder, the images used to quantize the model are stored in the text file specified by `--img_dataset_txt`.\n"
]
},
{
"cell_type": "markdown",
"id": "93e0d0622df170e",
"metadata": {},
"source": [
"### Optional: Download a dataset from Roboflow for quantization\n",
"\n",
"If you do not already have a dataset or set of images containing the objects you want to detect, follow the steps below to download one from Roboflow Universe.\n",
"\n",
"#### **Step 1: Search for a Dataset**\n",
"\n",
"Go to [Roboflow Universe](https://universe.roboflow.com) and use the search bar to locate a dataset relevant to what you want to detect.\n",
"**Note:** The dataset must include the classes or object types you intend to detect.\n",
"\n",
"#### **Step 2: Access the Dataset Tab**\n",
"\n",
"After selecting a suitable project, navigate to the **Dataset** tab. Click the **\"Download Dataset\"** button. A prompt will appear with several options, including:\n",
"\n",
"- Train a model with this dataset\n",
"- Train from a portion of this dataset\n",
"- Download dataset\n",
"\n",
"Select **Download dataset**.\n",
"\n",
"#### **Step 3: Choose Format and View Download Code**\n",
"\n",
"- Under **Image and Annotation Format**, choose the version of YOLO you are using:\n",
" - For **YOLOv5**, choose `YOLOv5 PyTorch`\n",
" - For **YOLOv8**, choose `YOLOv8`\n",
" - For **YOLOv11**, choose `YOLOv11`\n",
"- If multiple annotation formats are listed for your model, always select the one ending in **\"PyTorch\"**.\n",
"\n",
"Then, under **Download Options**, click **\"Show Download Code\"** and continue.\n",
"\n",
"In the resulting screen, you will see three tabs:\n",
"- **Jupyter**\n",
"- **Terminal**\n",
"- **Raw URL**\n",
"\n",
"Select the **Terminal** tab and copy the provided command.\n",
"\n",
"#### **Step 4: Paste and Run**\n",
"\n",
"Paste the copied command into the notebook cell below and run it. This will download and extract the dataset into your environment, making it ready for use in the quantization process.\n",
"\n",
"Make sure to prefix the command with \"`!`\" so it executes properly in this Jupyter Notebook environment."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8bf75c9dcb328c84",
"metadata": {},
"outputs": [],
"source": [
"!curl -L \"https://universe.roboflow.com/ds/FaF3HbDmF7?key=iMoJR25O9H\" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip"
]
},
{
"cell_type": "markdown",
"id": "72bad9cac670f1ab",
"metadata": {},
"source": [
"### RKNN Conversion Script\n",
"To get started, run the `create_rknn.py` script below, replacing the arguments with your own values. Refer to the table below for detailed information on each argument’s purpose and usage. The `--model_path` argument should point to your exported ONNX model from Step 1, and `--img_dir` must reference a valid directory containing either a dataset or a set of images to be used for quantization.\n",
"\n",
"#### Overview of the `create_rknn.py` script\n",
"\n",
"This script converts a YOLO ONNX model to RKNN format using a set of calibration images. It's designed to work with either:\n",
"\n",
"- A flat directory of images (e.g. `train/images`), **or**\n",
"- A dataset directory containing a `data.yaml` file that defines `train`, `val`, and/or `test` folders.\n",
"\n",
"##### Arguments\n",
"\n",
"| Argument | Type | Description |\n",
"|----------|------|-----------------------------------------------------------------------------------------------------------------|\n",
"| `--img_dir` (`-d`) | `str` (required) | Path to your image directory. This can either be a folder of images **or** a dataset folder with a `data.yaml`. |\n",
"| `--model_path` (`-m`) | `str` (required) | Path to your YOLO ONNX model, created in Step 1. |\n",
"| `--num_imgs` (`-ni`) | `int` (default: `300`) | Number of images to use for quantization calibration. |\n",
"| `--disable_quantize` (`-dq`) | `bool` (default: `False`) | Set to `True` to skip quantization entirely. Not recommended for performance, and should not be used for deployment on PhotonVision, which requires quantization. |\n",
"| `--rknn_output` (`-o`) | `str` (default: `out.rknn`) | File path where the final RKNN model should be saved. |\n",
"| `--img_dataset_txt` (`-ds`) | `str` (default: `imgs.txt`) | File path to store the list of images used during quantization. |\n",
"| `--verbose` (`-vb`) | `bool` (default: `False`) | Enable detailed logging from the RKNN API during conversion. |\n",
"\n",
"\n",
"##### *Notes*\n",
"\n",
"1. This script is designed for use with [PhotonVision](https://photonvision.org), and by default sets the target platform for RKNN conversion to `RK3588`, a chipset commonly found in many variants of the Orange Pi 5 series (e.g., Orange Pi 5, 5 Pro, 5 Plus, 5 Max, etc.). You may modify the `DEFAULT_PLATFORM` value in the `create_rknn.py` script to match your specific hardware or deployment requirements if necessary.\n",
"\n",
"2. If you followed the Roboflow dataset download instructions from the previous section, the dataset will have been extracted to your **current working directory**. In that case, you can simply set `--img_dir` to \"`.`\" to reference the current directory."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b09656dd",
"metadata": {},
"outputs": [],
"source": [
"%run create_rknn.py --img_dir ./datasets --model_path weights.onnx"
]
},
{
"cell_type": "markdown",
"id": "5b3a6806",
"metadata": {},
"source": [
"And that’s it! Your RKNN model file is now ready for deployment on an Orange Pi."
]
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

617
scripts/rknn_conversion.ipynb Normal file
View File

@@ -0,0 +1,617 @@
{
"cells": [
{
"metadata": {},
"cell_type": "markdown",
"source": [
"# RKNN Conversion Guide\n",
"\n",
"----------------------------\n",
"\n",
"### Before you start\n",
"\n",
"If you are not using Google Colab, it is recommended to create a separate [Python virtual environment](https://docs.python.org/3/library/venv.html) before you run this project. This ensures that packages installed for the conversion process do not conflict with other packages you may already have set up."
],
"id": "65e9f457d12dcc6b"
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"### Setup\n",
"\n",
"This notebook requires the use of custom Python code. Please run the installation script below to import these external scripts if you do not have them already.\n",
"\n",
"**You may need to run the `Create ONNX/RKNN` cell multiple times when you restart your session. If you see a `create_onnx`\n",
"or `create_rknn` not found error, rerun the cell below and then retry.**\n",
"\n",
"**Do not modify the cells in this setup section unless you know what you’re doing or have a specific reason to.**"
],
"id": "500d656b7cc0ebd7"
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"#### Create ONNX/RKNN\n",
"\n",
"Please run the cell below to be able to use the `create_onnx` and `create_rknn` functions."
],
"id": "798298b1dbe33d2d"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": [
"import os.path\n",
"import subprocess\n",
"import sys\n",
"import random\n",
"\n",
"# This will work for all models that don't use anchors (e.g. all YOLO models except YOLOv5/v7)\n",
"# This includes YOLOv5u\n",
"yolo_non_anchor_repo = \"https://github.com/airockchip/ultralytics_yolo11\"\n",
"\n",
"# For original YOLOv5 models\n",
"yolov5_repo = \"https://github.com/airockchip/yolov5\"\n",
"\n",
"valid_yolo_versions = [\"yolov5\", \"yolov8\", \"yolov11\"]\n",
"\n",
"ultralytics_folder_name_yolov5 = \"airockchip_yolo_pkg_yolov5\"\n",
"ultralytics_default_folder_name = \"airockchip_yolo_pkg\"\n",
"\n",
"bad_model_msg = \"\"\"\n",
"This is usually due to passing in the wrong model version.\n",
"Please make sure you have the right model version and try again.\n",
"\"\"\"\n",
"\n",
"\n",
"def print_bad_model_msg(cause):\n",
" print(f\"{cause}{bad_model_msg}\")\n",
"\n",
"def run_and_exit_with_error(cmd, error_msg, enable_error_output=True):\n",
" try:\n",
" if enable_error_output:\n",
" subprocess.run(\n",
" cmd,\n",
" stderr=subprocess.STDOUT,\n",
" stdout=subprocess.PIPE,\n",
" universal_newlines=True,\n",
" ).check_returncode()\n",
" else:\n",
" subprocess.run(cmd).check_returncode()\n",
" except subprocess.CalledProcessError as e:\n",
" print(error_msg)\n",
"\n",
" if enable_error_output:\n",
" print(e.stdout)\n",
"\n",
" sys.exit(1)\n",
"\n",
"\n",
"def check_git_installed():\n",
" run_and_exit_with_error(\n",
" [\"git\", \"--version\"],\n",
" \"\"\"Git is not installed or not found in your PATH.\n",
"Please install Git from https://git-scm.com/downloads and try again.\"\"\",\n",
" )\n",
"\n",
"\n",
"def check_or_clone_rockchip_repo(repo_url, repo_name=ultralytics_default_folder_name):\n",
" if os.path.exists(repo_name):\n",
" print(\n",
" f'Existing Rockchip repo \"{repo_name}\" detected, skipping installation...'\n",
" )\n",
" else:\n",
" print(f'Cloning Rockchip repo to \"{repo_name}\"')\n",
" run_and_exit_with_error(\n",
" [\"git\", \"clone\", repo_url, repo_name],\n",
" \"Failed to clone Rockchip repo, please see error output\",\n",
" )\n",
"\n",
"\n",
"def run_pip_install_or_else_exit(args):\n",
" print(\"Running pip install...\")\n",
" run_and_exit_with_error(\n",
" [\"pip\", \"install\"] + args,\n",
" \"Pip install rockchip repo failed, please see error output\",\n",
" )\n",
"\n",
"\n",
"def run_onnx_conversion_yolov5(model_path):\n",
" check_or_clone_rockchip_repo(yolov5_repo, ultralytics_folder_name_yolov5)\n",
" run_pip_install_or_else_exit(\n",
" [\n",
" \"-r\",\n",
" os.path.join(ultralytics_folder_name_yolov5, \"requirements.txt\"),\n",
" \"torch<2.6.0\",\n",
" \"onnx==1.18.0\",\n",
" \"onnxscript\",\n",
" ]\n",
" )\n",
"\n",
" model_abspath = os.path.abspath(model_path)\n",
"\n",
" try:\n",
" subprocess.run(\n",
" [\n",
" \"python\",\n",
" f\"{ultralytics_folder_name_yolov5}/export.py\",\n",
" \"--weights\",\n",
" model_abspath,\n",
" \"--rknpu\",\n",
" \"--include\",\n",
" \"onnx\",\n",
" ],\n",
" stderr=subprocess.STDOUT,\n",
" stdout=subprocess.PIPE,\n",
" universal_newlines=True,\n",
" ).check_returncode()\n",
" except subprocess.CalledProcessError as e:\n",
" print(\"Failed to run YOLOv5 export, please see error output\")\n",
"\n",
" if \"ModuleNotFoundError\" in e.stdout and \"ultralytics\" in e.stdout:\n",
" print_bad_model_msg(\n",
" \"It seems the YOLOv5 repo could not find an ultralytics installation.\"\n",
" )\n",
" elif \"AttributeError\" in e.stdout and \"_register_detect_seperate\" in e.stdout:\n",
" print_bad_model_msg(\"It seems that you received a model attribute error.\")\n",
" else:\n",
" print(\"Unknown Error when converting:\")\n",
" print(e.stdout)\n",
"\n",
" sys.exit(1)\n",
"\n",
"\n",
"def run_onnx_conversion_no_anchor(model_path):\n",
" check_or_clone_rockchip_repo(yolo_non_anchor_repo)\n",
" run_pip_install_or_else_exit(\n",
" [\"-e\", ultralytics_default_folder_name, \"onnx==1.18.0\", \"onnxscript\"]\n",
" )\n",
"\n",
" sys.path.insert(0, os.path.abspath(ultralytics_default_folder_name))\n",
" model_abs_path = os.path.abspath(model_path)\n",
"\n",
" from ultralytics import YOLO\n",
"\n",
" try:\n",
" model = YOLO(model_abs_path)\n",
" model.export(format=\"rknn\")\n",
" except TypeError as e:\n",
" if \"originally trained\" in str(e):\n",
" print_bad_model_msg(\n",
" \"Ultralytics has detected that this model is a YOLOv5 model.\"\n",
" )\n",
" else:\n",
" raise e\n",
"\n",
" sys.exit(1)\n",
"\n",
"\n",
"def create_onnx(model_path: str, version: str):\n",
" check_git_installed()\n",
"\n",
" if not version in valid_yolo_versions:\n",
" print(f\"YOLO version \\\"{version}\\\" is not a valid version! Valid versions are: {\", \".join(valid_yolo_versions)}\")\n",
"\n",
" try:\n",
" if version.lower() == \"yolov5\":\n",
" run_onnx_conversion_yolov5(model_path)\n",
" else:\n",
" run_onnx_conversion_no_anchor(model_path)\n",
"\n",
" print(\n",
" \"Model export finished. Please use the generated ONNX file to convert to RKNN.\"\n",
" )\n",
" except SystemExit:\n",
" raise RuntimeError(\"Model export failed. Please see output above.\")\n",
"\n",
"# RKNN Conversion code\n",
"\n",
"image_extensions = (\".jpg\", \".jpeg\", \".png\", \".bmp\", \".gif\", \".tiff\", \".webp\")\n",
"DEFAULT_PLATFORM = \"rk3588\"\n",
"\n",
"\n",
"def list_img_dir(img_dir):\n",
" return [\n",
" os.path.abspath(os.path.join(img_dir, f))\n",
" for f in os.listdir(img_dir)\n",
" if f.lower().endswith(image_extensions)\n",
" ]\n",
"\n",
"\n",
"def sample_imgs(num, img_list):\n",
" if len(img_list) < num:\n",
" return img_list\n",
" else:\n",
" return random.sample(img_list, num)\n",
"\n",
"\n",
"def get_image_list_from_dataset(num_imgs, yaml_dir):\n",
" print(f\"Dataset detected with {yaml_dir} file\")\n",
" img_raw_paths = []\n",
"\n",
" with open(yaml_dir, \"r\") as yaml_file:\n",
" for line in yaml_file:\n",
" line = line.strip()\n",
" if (\n",
" line.startswith(\"train:\")\n",
" or line.startswith(\"val:\")\n",
" or line.startswith(\"test:\")\n",
" ):\n",
" img_raw_paths.append(line.split(\":\", 1)[1].strip())\n",
"\n",
" no_yaml_dir = yaml_dir.replace(\n",
" \"data.yaml\", \"dummy_dir\"\n",
" ) # data.yaml sets dirs one level up\n",
" img_set_paths = []\n",
"\n",
" for img_raw_path in img_raw_paths:\n",
" p = (\n",
" img_raw_path\n",
" if os.path.isabs(img_raw_path)\n",
" else os.path.realpath(os.path.join(no_yaml_dir, img_raw_path))\n",
" )\n",
"\n",
" if os.path.exists(p):\n",
" img_set_paths.append(p)\n",
"\n",
" if len(img_set_paths) < 1:\n",
" return None\n",
"\n",
" all_imgs = [list_img_dir(path) for path in img_set_paths]\n",
"\n",
" for imgs in all_imgs:\n",
" print(len(imgs))\n",
"\n",
" total_imgs = sum(len(group) for group in all_imgs)\n",
"\n",
" sampled_imgs = [\n",
" sample_imgs(round((len(group) / total_imgs) * num_imgs), group)\n",
" for group in all_imgs\n",
" ]\n",
"\n",
" return [img for group in sampled_imgs for img in group]\n",
"\n",
"\n",
"def get_image_list_from_img_dir(num_imgs, img_dir):\n",
" return sample_imgs(num_imgs, list_img_dir(img_dir))\n",
"\n",
"\n",
"def get_image_list(num_imgs, image_dir):\n",
" yaml_path = os.path.join(image_dir, \"data.yaml\")\n",
"\n",
" if os.path.exists(yaml_path):\n",
" return get_image_list_from_dataset(num_imgs, yaml_path)\n",
" else:\n",
" return get_image_list_from_img_dir(num_imgs, image_dir)\n",
"\n",
"\n",
"def run_rknn_conversion(\n",
" img_list_txt, disable_quant, model_path, rknn_output, verbose_logging\n",
"):\n",
" from rknn.api import RKNN\n",
"\n",
" rknn = RKNN(\n",
" verbose=verbose_logging,\n",
" verbose_file=(\"rknn_convert.log\" if verbose_logging else None),\n",
" )\n",
"\n",
" rknn.config(\n",
" mean_values=[[0, 0, 0]],\n",
" std_values=[[255, 255, 255]],\n",
" target_platform=DEFAULT_PLATFORM,\n",
" )\n",
"\n",
" print(\"Attempted RKNN load\")\n",
" ret = rknn.load_onnx(model=model_path)\n",
" if ret != 0:\n",
" print(\"Loading model failed!\")\n",
" exit(ret)\n",
"\n",
" print(\"Attempting RKNN build\")\n",
" ret = rknn.build(do_quantization=(not disable_quant), dataset=img_list_txt)\n",
" if ret != 0:\n",
" print(\"Building model failed!\")\n",
" exit(ret)\n",
"\n",
" print(\"Build succeeded! Starting export...\")\n",
" ret = rknn.export_rknn(rknn_output)\n",
" if ret != 0:\n",
" print(\"Exporting model failed!\")\n",
" exit(ret)\n",
" print(\"Finished export!\")\n",
"\n",
" # Release\n",
" rknn.release()\n",
"\n",
" print(f'Your model is in \"{rknn_output}\" and ready to use!')\n",
"\n",
"\n",
"def create_rknn(\n",
" model_path: str,\n",
" rknn_output: str = \"out.rknn\",\n",
" num_imgs: int = 300,\n",
" img_dir: str = None,\n",
" img_dataset_txt: str = \"imgs.txt\",\n",
" disable_quantize: bool = False,\n",
" verbose: bool = False,\n",
"):\n",
" if not rknn_output.endswith(\".rknn\"):\n",
" print(\"RKNN output path must end in .rknn!\")\n",
" return\n",
"\n",
" if not disable_quantize:\n",
" if img_dir is None or len(img_dir) < 1:\n",
" print(f\"Must specify list of images to use with --img_dir\")\n",
" return\n",
"\n",
" img_dir_abs = os.path.abspath(img_dir)\n",
"\n",
" img_list = get_image_list(num_imgs, img_dir_abs)\n",
" img_list_len = 0 if img_list is None else len(img_list)\n",
"\n",
" if img_list_len == 0:\n",
" print(f\"No images found in {img_dir_abs}\")\n",
" return\n",
" elif img_list_len < num_imgs:\n",
" print(\n",
" f\"Not enough images in your dataset/directory, you have {img_list_len} images, but need {num_imgs}\"\n",
" )\n",
" return\n",
"\n",
" if not img_dataset_txt.endswith(\".txt\"):\n",
" print(f\"Image dataset text file path must end in .txt\")\n",
" return\n",
"\n",
" with open(img_dataset_txt, \"w\") as set_file:\n",
" set_file.writelines(f\"{img}\\n\" for img in img_list)\n",
"\n",
" try:\n",
" run_rknn_conversion(\n",
" img_dataset_txt,\n",
" disable_quantize,\n",
" model_path,\n",
" rknn_output,\n",
" verbose,\n",
" )\n",
" except SystemExit:\n",
" print(\"RKNN Conversion failed, see output above\")\n"
],
"id": "ea6869140a61126d"
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"#### *Numpy Fix* - Important for Google Colab Users\n",
"\n",
"Google Colab comes with an incompatible version of Numpy installed. To fix this, please run the following cells below and **restart your session** when prompted."
],
"id": "b3a9e1a334bce144"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": [
"%pip uninstall numpy -y\n",
"%pip install \"numpy>=1.23.0,<2.0.0\""
],
"id": "7156e69495f48f49"
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"### Step 1: Convert to ONNX\n",
"\n",
"To convert to ONNX, simply run the `create_onnx` function, providing the path to your model weights and specifying the model version, as shown below."
],
"id": "332942d8582c9bae"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": "create_onnx(model_path=\"weights.pt\", version=\"yolov8\") # Valid versions are yolov5, yolov8, and yolov11",
"id": "408440b32de224ed"
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"### Step 2: Download RKNN API\n",
"\n",
"You can either use `pip` below to automatically detect and install the correct RKNN API Python library for you, or install it manually.\n",
"\n",
"#### Automatic installation\n",
"\n",
"Please run `pip` below. If it does not work, refer to the instructions for manual installation. You may need to restart your session after running the command below.\n"
],
"id": "ebd148faaa2b1933"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": "%pip install rknn-toolkit2",
"id": "7c7ef3010c663fc2"
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"#### Manual Installation (If Automatic Installation Fails)\n",
"Visit the [RKNN Toolkit 2](https://github.com/airockchip/rknn-toolkit2) Github repository, then click on rknn-toolkit2, followed by packages.\n",
"If you are running an x86_64 CPU (e.g., most Intel and AMD processors), select that option; otherwise, choose arm64 for ARM-based computers (e.g., M-series Macs or Snapdragon processors). If you're unsure which CPU you're using, check your system settings for processor architecture information.\n",
"\n",
"Once you've selected the correct CPU architecture, you'll see multiple packages. The file names will look something like:\n",
"`rknn_toolkit2-2.3.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl`.\n",
"The numbers after cp correspond to your Python version. For example, if you're using Python 3.10, look for a package with cp310 in the name. For Python 3.8, look for cp38; for Python 3.7, cp37, and so on.\n",
"\n",
"Once you've found the correct package, click the \"Raw\" button to download the .whl file. Then, run the following command in your terminal, replacing rknn_toolkit2.whl with the actual path to the file you downloaded:"
],
"id": "f44b71c5a820ab2"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": "%pip install rknn_toolkit2.whl",
"id": "c9b9d3da532eb916"
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"### Step 3: Convert to RKNN\n",
"\n",
"Please review the notes about quantization before running the RKNN conversion."
],
"id": "fc74ebc99b596f76"
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"### Quantization\n",
"\n",
"When performing quantization, it is critical to provide representative images of the objects or scenes you are trying to detect. These images are used to calibrate the model’s internal activations and greatly influence the final performance.\n",
"\n",
"It is recommended to use 300–500 representative images that reflect the real-world input your model will encounter. As the old saying goes, it’s quality over quantity — having a diverse, relevant set matters more than simply having many images.\n",
"\n",
"Quantization will cause some loss in model accuracy. However, if your calibration images are chosen wisely, this accuracy drop should be minimal and acceptable. If the sampled images are too uniform or unrelated, your quantized model's performance may worsen significantly.\n",
"\n",
"The script will automatically sample representative images randomly from the provided dataset. While this usually works well, please verify that the dataset contains diverse and relevant examples of your target objects. As a reminder, the images used to quantize the model are stored in the text file specified by `--img_dataset_txt`.\n"
],
"id": "80cf63a7f16f9af"
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"### Optional: Download a dataset from Roboflow for quantization\n",
"\n",
"If you do not already have a dataset or set of images containing the objects you want to detect, follow the steps below to download one from Roboflow Universe.\n",
"\n",
"#### **Step 1: Search for a Dataset**\n",
"\n",
"Go to [Roboflow Universe](https://universe.roboflow.com) and use the search bar to locate a dataset relevant to what you want to detect.\n",
"**Note:** The dataset must include the classes or object types you intend to detect.\n",
"\n",
"#### **Step 2: Access the Dataset Tab**\n",
"\n",
"After selecting a suitable project, navigate to the **Dataset** tab. Click the **\"Download Dataset\"** button. A prompt will appear with several options, including:\n",
"\n",
"- Train a model with this dataset\n",
"- Train from a portion of this dataset\n",
"- Download dataset\n",
"\n",
"Select **Download dataset**.\n",
"\n",
"#### **Step 3: Choose Format and View Download Code**\n",
"\n",
"- Under **Image and Annotation Format**, choose the version of YOLO you are using:\n",
" - For **YOLOv5**, choose `YOLOv5 PyTorch`\n",
" - For **YOLOv8**, choose `YOLOv8`\n",
" - For **YOLOv11**, choose `YOLOv11`\n",
"- If multiple annotation formats are listed for your model, always select the one ending in **\"PyTorch\"**.\n",
"\n",
"Then, under **Download Options**, click **\"Show Download Code\"** and continue.\n",
"\n",
"In the resulting screen, you will see three tabs:\n",
"- **Jupyter**\n",
"- **Terminal**\n",
"- **Raw URL**\n",
"\n",
"Select the **Terminal** tab and copy the provided command.\n",
"\n",
"#### **Step 4: Paste and Run**\n",
"\n",
"Paste the copied command into the notebook cell below and run it. This will download and extract the dataset into your environment, making it ready for use in the quantization process.\n",
"\n",
"Make sure to prefix the command with \"`!`\" so it executes properly in this Jupyter Notebook environment."
],
"id": "ae229bef78e3bc0c"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": "!curl -L \"https://universe.roboflow.com/ds/FaF3HbDmF7?key=iMoJR25O9H\" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip",
"id": "e16fb4f928fd956b"
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"### RKNN Conversion Script\n",
"To get started, run the `create_rknn` script below, replacing the arguments with your own values. Refer to the table below for detailed information on each argument’s purpose and usage. The `model_path` argument should point to your exported ONNX model from Step 1, and `img_dir` must reference a valid directory containing either a dataset or a set of images to be used for quantization.\n",
"\n",
"#### Overview of the `create_rknn` function\n",
"\n",
"This script converts a YOLO ONNX model to RKNN format using a set of calibration images. It's designed to work with either:\n",
"\n",
"- A flat directory of images (e.g. `train/images`), **or**\n",
"- A dataset directory containing a `data.yaml` file that defines `train`, `val`, and/or `test` folders.\n",
"\n",
"##### Arguments\n",
"\n",
"| Argument | Type | Description |\n",
"|----------|------|-----------------------------------------------------------------------------------------------------------------|\n",
"| `img_dir` | `str` (required) | Path to your image directory. This can either be a folder of images **or** a dataset folder with a `data.yaml`. |\n",
"| `model_path` | `str` (required) | Path to your YOLO ONNX model, created in Step 1. |\n",
"| `num_imgs` | `int` (default: `300`) | Number of images to use for quantization calibration. |\n",
"| `disable_quantize` | `bool` (default: `False`) | Set to `True` to skip quantization entirely. Not recommended for performance, and should not be used for deployment on PhotonVision, which requires quantization. |\n",
"| `rknn_output` | `str` (default: `out.rknn`) | File path where the final RKNN model should be saved. |\n",
"| `img_dataset_txt` | `str` (default: `imgs.txt`) | File path to store the list of images used during quantization. |\n",
"| `verbose` | `bool` (default: `False`) | Enable detailed logging from the RKNN API during conversion. |\n",
"\n",
"\n",
"##### *Notes*\n",
"\n",
"1. This script is designed for use with [PhotonVision](https://photonvision.org), and by default sets the target platform for RKNN conversion to `RK3588`, a chipset commonly found in many variants of the Orange Pi 5 series (e.g., Orange Pi 5, 5 Pro, 5 Plus, 5 Max, etc.). You may modify the `DEFAULT_PLATFORM` value in the setup cell to match your specific hardware or deployment requirements if necessary.\n",
"\n",
"2. If you followed the Roboflow dataset download instructions from the previous section, the dataset will have been extracted to your **current working directory**. In that case, you can simply set `img_dir` to \"`.`\" to reference the current directory."
],
"id": "f8f48b3139509618"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": "create_rknn(img_dir=\"./datasets\", model_path=\"weights.onnx\")",
"id": "2c48b133f5c93c7a"
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}