diff --git a/utils/onnx2coreml.py b/utils/onnx2coreml.py new file mode 100644 index 00000000..3b9af6bd --- /dev/null +++ b/utils/onnx2coreml.py @@ -0,0 +1,137 @@ +import os +from onnx import onnx_pb +from onnx_coreml import convert +import glob + + +# https://github.com/onnx/onnx-coreml +# http://machinethink.net/blog/mobilenet-ssdlite-coreml/ + +def main(): + os.system('rm -rf saved_models && mkdir saved_models') + files = glob.glob('saved_models/*.onnx') + glob.glob('../yolov3/weights/*.onnx') + + for f in files: + # 1. ONNX to CoreML + name = 'saved_models/' + f.split('/')[-1].replace('.onnx', '') + + model_file = open(f, 'rb') + model_proto = onnx_pb.ModelProto() + model_proto.ParseFromString(model_file.read()) + coreml_model = convert(model_proto, image_input_names=['0']) + # coreml_model.save(model_out) + + # 2. Reduce model to FP16, change outputs to DOUBLE and save + import coremltools + + spec = coreml_model.get_spec() + for i in range(2): + spec.description.output[i].type.multiArrayType.dataType = \ + coremltools.proto.FeatureTypes_pb2.ArrayFeatureType.ArrayDataType.Value('DOUBLE') + + spec = coremltools.utils.convert_neural_network_spec_weights_to_fp16(spec) + coreml_model = coremltools.models.MLModel(spec) + + num_classes = 80 + num_anchors = 507 + spec.description.output[0].type.multiArrayType.shape.append(num_classes) + spec.description.output[0].type.multiArrayType.shape.append(num_anchors) + + spec.description.output[1].type.multiArrayType.shape.append(4) + spec.description.output[1].type.multiArrayType.shape.append(num_anchors) + coreml_model.save(name + '.mlmodel') + print(spec.description) + + # 3. Create NMS protobuf + import numpy as np + + nms_spec = coremltools.proto.Model_pb2.Model() + nms_spec.specificationVersion = 3 + + for i in range(2): + decoder_output = coreml_model._spec.description.output[i].SerializeToString() + + nms_spec.description.input.add() + nms_spec.description.input[i].ParseFromString(decoder_output) + + nms_spec.description.output.add() + nms_spec.description.output[i].ParseFromString(decoder_output) + + nms_spec.description.output[0].name = 'confidence' + nms_spec.description.output[1].name = 'coordinates' + + output_sizes = [num_classes, 4] + for i in range(2): + ma_type = nms_spec.description.output[i].type.multiArrayType + ma_type.shapeRange.sizeRanges.add() + ma_type.shapeRange.sizeRanges[0].lowerBound = 0 + ma_type.shapeRange.sizeRanges[0].upperBound = -1 + ma_type.shapeRange.sizeRanges.add() + ma_type.shapeRange.sizeRanges[1].lowerBound = output_sizes[i] + ma_type.shapeRange.sizeRanges[1].upperBound = output_sizes[i] + del ma_type.shape[:] + + nms = nms_spec.nonMaximumSuppression + nms.confidenceInputFeatureName = '133' # 1x507x80 + nms.coordinatesInputFeatureName = '134' # 1x507x4 + nms.confidenceOutputFeatureName = 'confidence' + nms.coordinatesOutputFeatureName = 'coordinates' + nms.iouThresholdInputFeatureName = 'iouThreshold' + nms.confidenceThresholdInputFeatureName = 'confidenceThreshold' + + nms.iouThreshold = 0.6 + nms.confidenceThreshold = 0.4 + nms.pickTop.perClass = True + + labels = np.loadtxt('../yolov3/data/coco.names', dtype=str, delimiter='\n') + nms.stringClassLabels.vector.extend(labels) + + nms_model = coremltools.models.MLModel(nms_spec) + nms_model.save(name + '_nms.mlmodel') + + # 4. Pipeline models togethor + from coremltools.models import datatypes + # from coremltools.models import neural_network + from coremltools.models.pipeline import Pipeline + + input_features = [('image', datatypes.Array(3, 416, 416)), + ('iouThreshold', datatypes.Double()), + ('confidenceThreshold', datatypes.Double())] + + output_features = ['confidence', 'coordinates'] + + pipeline = Pipeline(input_features, output_features) + + # Add 3rd dimension of size 1 (apparently not needed, produces error on compile) + # ssd_output = coreml_model._spec.description.output + # ssd_output[0].type.multiArrayType.shape[:] = [num_classes, num_anchors, 1] + # ssd_output[1].type.multiArrayType.shape[:] = [4, num_anchors, 1] + + # And now we can add the three models, in order: + pipeline.add_model(coreml_model) + pipeline.add_model(nms_model) + + # Correct datatypes + pipeline.spec.description.input[0].ParseFromString(coreml_model._spec.description.input[0].SerializeToString()) + pipeline.spec.description.output[0].ParseFromString(nms_model._spec.description.output[0].SerializeToString()) + pipeline.spec.description.output[1].ParseFromString(nms_model._spec.description.output[1].SerializeToString()) + + # Update metadata + pipeline.spec.description.metadata.versionString = 'yolov3-tiny.pt imported from PyTorch' + pipeline.spec.description.metadata.shortDescription = 'https://github.com/ultralytics/yolov3' + pipeline.spec.description.metadata.author = 'glenn.jocher@ultralytics.com' + pipeline.spec.description.metadata.license = 'https://github.com/ultralytics/yolov3' + + user_defined_metadata = {'classes': ','.join(labels), + 'iou_threshold': str(nms.iouThreshold), + 'confidence_threshold': str(nms.confidenceThreshold)} + pipeline.spec.description.metadata.userDefined.update(user_defined_metadata) + + # Save the model + pipeline.spec.specificationVersion = 3 + final_model = coremltools.models.MLModel(pipeline.spec) + final_model.save((name + '_pipelined.mlmodel')) + + +if __name__ == '__main__': + main()