Add anchors to V5 & V7 export (#128)

* Add anchors to V5 & V7 export

* Update NN Archive & Classification

* Rename V8 obj det output names

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"

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

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()
@@ -159,7 +162,6 @@ def make_nn_archive(
                             "preprocessing": {
                                 "mean": [0, 0, 0],
                                 "scale": [255, 255, 255],
-                                "reverse_channels": True,
                             },
                         }
                     ],
@@ -185,6 +187,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,

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
+        )

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",
@@ -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,

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
+        )