diff --git a/scripts/rknn-convert-tool/create_onnx.py b/scripts/rknn-convert-tool/create_onnx.py index 597b0315b..c568e2d57 100644 --- a/scripts/rknn-convert-tool/create_onnx.py +++ b/scripts/rknn-convert-tool/create_onnx.py @@ -16,33 +16,44 @@ 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. """ -# 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( @@ -50,23 +61,18 @@ def check_or_clone_rockchip_repo(repo_url, repo_name=ultralytics_default_folder_ ) 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): @@ -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) @@ -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) diff --git a/scripts/rknn-convert-tool/rknn_conversion.ipynb b/scripts/rknn-convert-tool/rknn_conversion.ipynb index 2ca06ed8e..80bcfe695 100644 --- a/scripts/rknn-convert-tool/rknn_conversion.ipynb +++ b/scripts/rknn-convert-tool/rknn_conversion.ipynb @@ -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", @@ -254,7 +254,7 @@ "| `--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", @@ -262,7 +262,7 @@ "\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." ]