Instance Segmentation Update

yolo/export_yolov5.py

@@ -1,5 +1,6 @@
 import sys
 sys.path.append("./yolo/yolov5")
+# sys.path.append("./yolov5")
 import torch
 
 from yolov5.models.experimental import attempt_load

yolo/export_yolov5seg.py

@@ -0,0 +1,74 @@
+from typing import Any
+from pathlib import Path
+import onnx
+
+from export_yolov5 import YoloV5Exporter
+
+
+class YoloV5SegExporter(YoloV5Exporter):
+    def __init__(self, conv_path: Path, weights_filename: str, imgsz: int,
+                conv_id: Any, n_shaves: int = 6, use_legacy_frontend: str = 'false'):
+        super().__init__(conv_path, weights_filename, imgsz, conv_id, n_shaves, use_legacy_frontend)
+
+    # load the model
+    def load_model(self):
+        super().load_model()
+        m = self.model.module.model[-1] if hasattr(self.model, 'module') else self.model.model[-1]
+
+        self.num_masks = m.nm           # number of masks
+        self.num_out = m.no             # number of outputs per anchor
+        self.num_anch = m.na            # number of anchors
+        self.num_det_out = m.no-m.nm    # number of outputs without masks
+
+    # splitting layers into masks and the detection output
+    # using sigmoid on the detections, concatenating
+    def prune_model(self, onnx_model: onnx.ModelProto, inputs: list):
+        for i, inp in enumerate(inputs):
+            slice_out = onnx.helper.make_node(
+                'Split',
+                inputs=[inp],
+                outputs=[f'split_outputs{i+1}', f'split_masks{i+1}'],
+                split=[self.num_det_out, self.num_masks],
+                axis=2
+            )
+
+            sigmoid_out = onnx.helper.make_node(
+                'Sigmoid',
+                inputs=[f'split_outputs{i+1}'],
+                outputs=[f'sigmoid_outputs{i+1}']
+            )
+
+            out_node = onnx.helper.make_node(
+                'Concat',
+                inputs=[f'sigmoid_outputs{i+1}', f'split_masks{i+1}'],
+                outputs=[f'output{i+1}_yolov5seg'],
+                axis=2
+            )
+
+            onnx_model.graph.node.append(slice_out)
+            onnx_model.graph.node.append(sigmoid_out)
+            onnx_model.graph.node.append(out_node)
+
+    # exporting + preparing to prune the model 
+    def export_onnx(self) -> Path:
+        onnx_model, check = self.get_onnx()
+        assert check, 'assert check failed'
+
+        # get indices of convolustions (outputs), then getting connected reshapes
+        # from reshapes it's much easier to split outputs
+        conv_out = [n.output[0] for n in onnx_model.graph.node if "Conv" in n.name][-self.num_branches:]
+        inputs = [n.output[0] for n in onnx_model.graph.node
+                if "Reshape" in n.name and n.input[0] in conv_out]
+
+        # preparing the model for pruning in OpenVINO
+        self.prune_model(onnx_model, inputs)
+        onnx.checker.check_model(onnx_model)  # check onnx model
+
+        # save the simplified model
+        self.f_simplified = (self.conv_path / f"{self.model_name}-simplified.onnx").resolve()
+        onnx.save(onnx_model, self.f_simplified)
+        return self.f_simplified
+
+    # was facing an error without overriding the class, it should affect anything tho
+    def export_openvino(self, version: str = 'v5seg'):
+        return super().export_openvino(version)

yolo/exporter.py

@@ -32,11 +32,18 @@ def get_onnx(self):
         # export onnx model
         self.f_onnx = (self.conv_path / f"{self.model_name}.onnx").resolve()
         im = torch.zeros(1, 3, *self.imgsz[::-1])#.to(device)  # image size(1,3,320,192) BCHW iDetection
+
+        # check if it's a segmentation model
+        # works fine for yolov5 instance segmentation model, not sure about other models
+        out_names = ['output']
+        if 'Segmentation' in type(self.model).__name__:
+            out_names += ['protos_output']
+
         torch.onnx.export(self.model, im, self.f_onnx, verbose=False, opset_version=12,
                         training=torch.onnx.TrainingMode.EVAL,
                         do_constant_folding=True,
                         input_names=['images'],
-                        output_names=['output'],
+                        output_names=out_names,
                         dynamic_axes=None)
 
         # check if the arhcitecture is correct
@@ -53,12 +60,14 @@ def export_openvino(self, version):
             self.export_onnx()
 
         output_list = [f"output{i+1}_yolo{version}" for i in range(self.num_branches)]
+        if 'Segmentation' in type(self.model).__name__:
+            output_list += ["protos_output"]
         output_list = ",".join(output_list)
 
         # export to OpenVINO and prune the model in the process
-        cmd = f"mo --input_model '{self.f_simplified}' " \
-        f"--output_dir '{self.conv_path.resolve()}' " \
-        f"--model_name '{self.model_name}' " \
+        cmd = f'mo --input_model "{self.f_simplified}" ' \
+        f'--output_dir "{self.conv_path.resolve()}" ' \
+        f'--model_name "{self.model_name}" ' \
         '--data_type FP16 ' \
         '--reverse_input_channels ' \
         '--scale 255 ' \
@@ -153,4 +162,4 @@ def make_zip(self):
         zip_obj.close()
 
         self.f_zip = f_zip
-        return f_zip
+        return f_zip

yolo/test_seg_export.py

@@ -0,0 +1,10 @@
+from pathlib import Path
+from export_yolov5seg import YoloV5SegExporter
+
+
+conv_path = Path('model_files')
+weights_filename = 'yolov5n-seg.pt'
+img_size = 416
+exporter = YoloV5SegExporter(conv_path, weights_filename, img_size, None)
+
+f_path = exporter.make_zip()