Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 42 additions & 37 deletions scripts/rknn-convert-tool/create_onnx.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,57 +16,63 @@
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.
"""


# idk how else to make Google Colab display this nicely
class IncorrectModelError(Exception):
def __init__(self, message):
self.message = message
super().__init__(self.message)


def print_bad_model_msg(cause):
print(f"{cause}{bad_model_msg}")


def check_git_installed():
def run_and_exit_with_error(cmd, error_msg, enable_error_output=True):
try:
subprocess.run(["git", "--version"]).check_returncode()
except:
print("Git is not installed or not found in your PATH.")
print("Please install Git from https://git-scm.com/downloads and try again.")
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}"')
try:
subprocess.run(["git", "clone", repo_url, repo_name]).check_returncode()
except subprocess.CalledProcessError as e:
print("Failed to clone Rockchip repo, see error output below")
print(e.output)
sys.exit(1)
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...")

try:
subprocess.run(["pip", "install"] + args).check_returncode()
except subprocess.CalledProcessError as e:
print("Pip install rockchip repo failed, see error output")
print(e.output)
sys.exit(1)
run_and_exit_with_error(
["pip", "install"] + args,
"Pip install rockchip repo failed, please see error output",
)


def run_onnx_conversion_yolov5(model_path):
Expand All @@ -93,23 +99,22 @@ def run_onnx_conversion_yolov5(model_path):
"--include",
"onnx",
],
capture_output=True,
text=True,
stderr=subprocess.STDOUT,
stdout=subprocess.PIPE,
universal_newlines=True,
).check_returncode()
except subprocess.CalledProcessError as e:
print("Failed to run YOLOv5 export, see output below")
output_string = (e.stdout or "") + (e.stderr or "")
print(output_string)
print("Failed to run YOLOv5 export, please see error output")

if "ModuleNotFoundError" in output_string and "ultralytics" in output_string:
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 output_string
and "_register_detect_seperate" in output_string
):
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)

Expand All @@ -132,7 +137,7 @@ def run_onnx_conversion_no_anchor(model_path):
"Ultralytics has detected that this model is a YOLOv5 model."
)
else:
print(e)
raise e

sys.exit(1)

Expand Down
8 changes: 4 additions & 4 deletions scripts/rknn-convert-tool/rknn_conversion.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,11 @@
"# DO NOT modify the filenames\n",
"scripts = [\n",
" {\n",
" \"url\": \"https://raw.githubusercontent.com/boomermath/photonvision_rknn_fork/refs/heads/rknn_conversion_tool/scripts/rknn-convert-tool/create_onnx.py\",\n",
" \"url\": \"https://raw.githubusercontent.com/PhotonVision/photonvision/ba1c0db7e19db090ca04a8375255b00db2e0babd/scripts/rknn-convert-tool/create_onnx.py\",\n",
" \"filename\": \"create_onnx.py\" # CREATE_ONNX_SCRIPT\n",
" },\n",
" {\n",
" \"url\": \"https://raw.githubusercontent.com/boomermath/photonvision_rknn_fork/refs/heads/rknn_conversion_tool/scripts/rknn-convert-tool/create_rknn.py\",\n",
" \"url\": \"https://raw.githubusercontent.com/PhotonVision/photonvision/ba1c0db7e19db090ca04a8375255b00db2e0babd/scripts/rknn-convert-tool/create_rknn.py\",\n",
" \"filename\": \"create_rknn.py\" # CREATE_RKNN_SCRIPT\n",
" }\n",
"]\n",
Expand Down Expand Up @@ -254,15 +254,15 @@
"| `--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. |\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 `TARGET_PLATFORM` value in the `create_onnx.py` script to match your specific hardware or deployment requirements if necessary.\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."
]
Expand Down
Loading