From cc248950d8d7b528c215cd35c1c2760feb2c3e4e Mon Sep 17 00:00:00 2001 From: HonzaCuhel Date: Mon, 5 Aug 2024 16:49:05 +0200 Subject: [PATCH] Add V9 support --- pyproject.toml | 2 +- tools/utils/version_detection.py | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 49b6853..c8915d8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "tools" -version = "0.1.2" +version = "0.1.3" description = "Converter for YOLO models into .ONNX format." readme = "README.md" requires-python = ">=3.8" diff --git a/tools/utils/version_detection.py b/tools/utils/version_detection.py index 145f9c2..1a134ab 100644 --- a/tools/utils/version_detection.py +++ b/tools/utils/version_detection.py @@ -45,11 +45,15 @@ def detect_version(path: str, debug: bool = False) -> str: print(data.decode(errors="replace")) content = data.decode("latin1") - if ( + if "yolov10" in content or "v10DetectLoss" in content: + return YOLOV10_CONVERSION + elif ( "YOLOv5u" in content or "YOLOv8" in content or "yolov8" in content or ("v8DetectionLoss" in content and "ultralytics" in content) + or "yolov9" in content + or ("v9-model" and "ultralytics" in content) ): return YOLOV8_CONVERSION elif "yolov6" in content: @@ -62,8 +66,6 @@ def detect_version(path: str, debug: bool = False) -> str: return YOLOV6R3_CONVERSION elif "yolov7" in content: return YOLOV7_CONVERSION - elif "yolov10" in content: - return YOLOV10_CONVERSION elif ( "SPPF" in content or "yolov5" in content