From 2fd43d3f8039ed2e780f3abb964415d4d6f1b089 Mon Sep 17 00:00:00 2001 From: HonzaCuhel Date: Wed, 22 Jan 2025 11:09:33 +0100 Subject: [PATCH 1/3] Add anchors to V5 & V7 export --- pyproject.toml | 2 +- requirements.txt | 2 +- tools/modules/exporter.py | 4 ++++ tools/yolo/yolov5_exporter.py | 12 +++++++++--- tools/yolov7/yolov7_exporter.py | 9 ++++++--- 5 files changed, 21 insertions(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0998029..cfc00f6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "tools" -version = "0.2.0" +version = "0.2.1" description = "Converter for YOLO models into .ONNX format." readme = "README.md" requires-python = ">=3.8" diff --git a/requirements.txt b/requirements.txt index 0b6b353..425c9ac 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ torchvision>=0.10.1 Pillow>=7.1.2 PyYAML>=5.3.1 gcsfs -luxonis-ml[data,nn_archive]>=0.4.1 +luxonis-ml[data,nn_archive]>=0.5.1 # luxonis-ml[data,nn_archive] @ git+https://github.com/luxonis/luxonis-ml@dev onnx>=1.10.1 onnxruntime diff --git a/tools/modules/exporter.py b/tools/modules/exporter.py index 3ffd910..6877395 100644 --- a/tools/modules/exporter.py +++ b/tools/modules/exporter.py @@ -47,6 +47,7 @@ def __init__( self.model_path = model_path self.imgsz = imgsz self.model_name = os.path.basename(self.model_path).split(".")[0] + self.model = None # Set up file paths self.f_onnx = None self.f_nn_archive = None @@ -113,6 +114,7 @@ def make_nn_archive( n_prototypes: Optional[int] = None, n_keypoints: Optional[int] = None, is_softmax: Optional[bool] = None, + anchors: Optional[List[List[List[float]]]] = None, output_kwargs: Optional[dict] = None, ): """Export the model to NN archive format. @@ -129,6 +131,7 @@ def make_nn_archive( n_prototypes (Optional[int], optional): Number of prototypes. Defaults to None. n_keypoints (Optional[int], optional): Number of keypoints. Defaults to None. is_softmax (Optional[bool], optional): Whether to use softmax. Defaults to None. + anchors (Optional[List[List[List[float]]]], optional): Anchors. Defaults to None. output_kwargs (Optional[dict], optional): Output keyword arguments. Defaults to None. """ self.f_nn_archive = (self.output_folder / f"{self.model_name}.tar.xz").resolve() @@ -185,6 +188,7 @@ def make_nn_archive( n_prototypes=n_prototypes, n_keypoints=n_keypoints, is_softmax=is_softmax, + anchors=anchors, **output_kwargs, ), outputs=self.all_output_names, diff --git a/tools/yolo/yolov5_exporter.py b/tools/yolo/yolov5_exporter.py index b5571fc..e4d1ee3 100644 --- a/tools/yolo/yolov5_exporter.py +++ b/tools/yolo/yolov5_exporter.py @@ -80,8 +80,8 @@ def load_model(self): except Exception as e: print(f"Error while getting number of channels: {e}") - m = model.module.model[-1] if hasattr(model, "module") else model.model[-1] - self.num_branches = len(m.anchor_grid) + self.m = model.module.model[-1] if hasattr(model, "module") else model.model[-1] + self.num_branches = len(self.m.anchor_grid) def export_nn_archive(self, class_names: Optional[List[str]] = None): """ @@ -98,4 +98,10 @@ def export_nn_archive(self, class_names: Optional[List[str]] = None): ), f"Number of the given class names {len(class_names)} does not match number of classes {self.model.nc} provided in the model!" names = class_names - self.make_nn_archive(names, self.model.nc) + anchors = [ + self.m.anchor_grid[i][0, :, 0, 0].numpy().tolist() + for i in range(self.num_branches) + ] + self.make_nn_archive( + names, self.model.nc, parser="YOLOExtendedParser", anchors=anchors + ) diff --git a/tools/yolov7/yolov7_exporter.py b/tools/yolov7/yolov7_exporter.py index 254521f..07252d6 100644 --- a/tools/yolov7/yolov7_exporter.py +++ b/tools/yolov7/yolov7_exporter.py @@ -64,8 +64,8 @@ def load_model(self): self.model = model - m = model.module.model[-1] if hasattr(model, "module") else model.model[-1] - self.num_branches = len(m.anchor_grid) + self.m = model.module.model[-1] if hasattr(model, "module") else model.model[-1] + self.num_branches = len(self.m.anchor_grid) def export_nn_archive(self, class_names: Optional[List[str]] = None): """ @@ -82,4 +82,7 @@ def export_nn_archive(self, class_names: Optional[List[str]] = None): ), f"Number of the given class names {len(class_names)} does not match number of classes {self.model.nc} provided in the model!" names = class_names - self.make_nn_archive(names, self.model.nc) + anchors = self.m.anchor_grid[:, 0, :, 0, 0].numpy().tolist() + self.make_nn_archive( + names, self.model.nc, parser="YOLOExtendedParser", anchors=anchors + ) From 7f60fcd470e9ac013cd551fe82024b7b3b97a251 Mon Sep 17 00:00:00 2001 From: HonzaCuhel Date: Fri, 31 Jan 2025 16:07:15 +0100 Subject: [PATCH 2/3] Update NN Archive & Classification --- tools/modules/exporter.py | 1 - tools/yolo/yolov8_exporter.py | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/tools/modules/exporter.py b/tools/modules/exporter.py index 6877395..d899b67 100644 --- a/tools/modules/exporter.py +++ b/tools/modules/exporter.py @@ -162,7 +162,6 @@ def make_nn_archive( "preprocessing": { "mean": [0, 0, 0], "scale": [255, 255, 255], - "reverse_channels": True, }, } ], diff --git a/tools/yolo/yolov8_exporter.py b/tools/yolo/yolov8_exporter.py index e59596e..36f3573 100644 --- a/tools/yolo/yolov8_exporter.py +++ b/tools/yolo/yolov8_exporter.py @@ -235,7 +235,6 @@ def make_cls_nn_archive(self, class_list: List[str], n_classes: int): "preprocessing": { "mean": [0, 0, 0], "scale": [255, 255, 255], - "reverse_channels": True, }, } ], @@ -248,7 +247,7 @@ def make_cls_nn_archive(self, class_list: List[str], n_classes: int): ], "heads": [ Head( - parser="Classification", + parser="ClassificationParser", metadata=HeadClassificationMetadata( is_softmax=False, n_classes=n_classes, From 776dbfe23b9cb76e5fd73533b8d37e11ff44a19d Mon Sep 17 00:00:00 2001 From: HonzaCuhel Date: Sun, 2 Feb 2025 15:11:01 +0100 Subject: [PATCH 3/3] Rename V8 obj det output names --- tools/yolo/yolov8_exporter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/yolo/yolov8_exporter.py b/tools/yolo/yolov8_exporter.py index 36f3573..deb96a8 100644 --- a/tools/yolo/yolov8_exporter.py +++ b/tools/yolo/yolov8_exporter.py @@ -47,7 +47,7 @@ def get_output_names(mode: int) -> List[str]: List[str]: List of output names """ if mode == DETECT_MODE: - return ["output1_yolov8", "output2_yolov8", "output3_yolov8"] + return ["output1_yolov6r2", "output2_yolov6r2", "output3_yolov6r2"] elif mode == SEGMENT_MODE: return [ "output1_yolov8",