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