coreml-converter.py 494 B

123456789101112131415161718
  1. #!/usr/bin/env python
  2. import os
  3. import sys
  4. file = sys.argv[1];
  5. base, extension = os.path.splitext(file)
  6. if extension == '.h5':
  7. import coremltools
  8. coreml_model = coremltools.converters.keras.convert(file)
  9. coreml_model.save(base + '.mlmodel')
  10. elif extension == '.pkl':
  11. import coremltools
  12. import sklearn
  13. sklearn_model = sklearn.externals.joblib.load(file)
  14. coreml_model = coremltools.converters.sklearn.convert(sklearn_model)
  15. coreml_model.save(base + '.mlmodel')