Merge pull request #321 from dlawrences/master

Update CoreML model conversion to do pixel scaling as well
pull/104/head^2
Glenn Jocher 2020-07-08 00:04:46 -07:00 committed by GitHub
commit cd6609297c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions

View File

@ -61,7 +61,8 @@ if __name__ == '__main__':
import coremltools as ct
print('\nStarting CoreML export with coremltools %s...' % ct.__version__)
model = ct.convert(ts, inputs=[ct.ImageType(name='images', shape=img.shape)]) # convert
# convert model from torchscript and apply pixel scaling as per detect.py
model = ct.convert(ts, inputs=[ct.ImageType(name='images', shape=img.shape, scale=1/255.0, bias=[0, 0, 0])])
f = opt.weights.replace('.pt', '.mlmodel') # filename
model.save(f)
print('CoreML export success, saved as %s' % f)