ONNX compatibility updates
This commit is contained in:
Glenn Jocher
parent
29f2e80950
commit
8b34fbef33
|
|
@ -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()
|
||||||
Loading…
Reference in New Issue