diff --git a/yolo/YOLOv6 b/yolo/YOLOv6 index b495484..4364f29 160000 --- a/yolo/YOLOv6 +++ b/yolo/YOLOv6 @@ -1 +1 @@ -Subproject commit b4954849b4c562937afbc5df2ad0701e295e0464 +Subproject commit 4364f29bf3244f2e73d0c42a103cd7a9cbb16ca9 diff --git a/yolo/export_yolov5.py b/yolo/export_yolov5.py index 1fda88e..bb897ad 100644 --- a/yolo/export_yolov5.py +++ b/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 diff --git a/yolo/export_yolov5seg.py b/yolo/export_yolov5seg.py new file mode 100644 index 0000000..841d236 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/yolo/exporter.py b/yolo/exporter.py index f9f05ac..88b6e98 100644 --- a/yolo/exporter.py +++ b/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 \ No newline at end of file diff --git a/yolo/model_files/yolov5n-seg.pt b/yolo/model_files/yolov5n-seg.pt new file mode 100644 index 0000000..cb816e5 Binary files /dev/null and b/yolo/model_files/yolov5n-seg.pt differ diff --git a/yolo/test_seg_export.py b/yolo/test_seg_export.py new file mode 100644 index 0000000..40bfc72 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/yolo/ultralytics b/yolo/ultralytics index dce4efc..c340f84 160000 --- a/yolo/ultralytics +++ b/yolo/ultralytics @@ -1 +1 @@ -Subproject commit dce4efce48a05e028e6ec430045431c242e52484 +Subproject commit c340f84ce9325de720fbd9ada6523a28fc432651 diff --git a/yolo/yolov5 b/yolo/yolov5 index c7a2d6b..a199480 160000 --- a/yolo/yolov5 +++ b/yolo/yolov5 @@ -1 +1 @@ -Subproject commit c7a2d6bcf4f7e88db53f3d09a8484391dac7bc89 +Subproject commit a199480ba6bb527598df11abbc1d679ccda82670 diff --git a/yolov6r1/yolo/YOLOv6R1 b/yolov6r1/yolo/YOLOv6R1 index 5a2fa16..4364f29 160000 --- a/yolov6r1/yolo/YOLOv6R1 +++ b/yolov6r1/yolo/YOLOv6R1 @@ -1 +1 @@ -Subproject commit 5a2fa16120bdff4d8261cc07a84782de5945d448 +Subproject commit 4364f29bf3244f2e73d0c42a103cd7a9cbb16ca9