glenn-jocher
commited on
Commit
•
b5b56a3
1
Parent(s):
7b31a53
Add CoreML inference (#6195)
Browse files* Add Apple CoreML inference
* Cleanup
detect.py
CHANGED
@@ -17,7 +17,7 @@ Usage - formats:
|
|
17 |
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
18 |
yolov5s.xml # OpenVINO
|
19 |
yolov5s.engine # TensorRT
|
20 |
-
yolov5s.mlmodel # CoreML (
|
21 |
yolov5s_saved_model # TensorFlow SavedModel
|
22 |
yolov5s.pb # TensorFlow GraphDef
|
23 |
yolov5s.tflite # TensorFlow Lite
|
|
|
17 |
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
18 |
yolov5s.xml # OpenVINO
|
19 |
yolov5s.engine # TensorRT
|
20 |
+
yolov5s.mlmodel # CoreML (MacOS-only)
|
21 |
yolov5s_saved_model # TensorFlow SavedModel
|
22 |
yolov5s.pb # TensorFlow GraphDef
|
23 |
yolov5s.tflite # TensorFlow Lite
|
export.py
CHANGED
@@ -25,7 +25,7 @@ Inference:
|
|
25 |
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
26 |
yolov5s.xml # OpenVINO
|
27 |
yolov5s.engine # TensorRT
|
28 |
-
yolov5s.mlmodel # CoreML (
|
29 |
yolov5s_saved_model # TensorFlow SavedModel
|
30 |
yolov5s.pb # TensorFlow GraphDef
|
31 |
yolov5s.tflite # TensorFlow Lite
|
@@ -156,7 +156,6 @@ def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
|
|
156 |
LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')
|
157 |
f = file.with_suffix('.mlmodel')
|
158 |
|
159 |
-
model.train() # CoreML exports should be placed in model.train() mode
|
160 |
ts = torch.jit.trace(model, im, strict=False) # TorchScript model
|
161 |
ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
|
162 |
ct_model.save(f)
|
|
|
25 |
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
26 |
yolov5s.xml # OpenVINO
|
27 |
yolov5s.engine # TensorRT
|
28 |
+
yolov5s.mlmodel # CoreML (MacOS-only)
|
29 |
yolov5s_saved_model # TensorFlow SavedModel
|
30 |
yolov5s.pb # TensorFlow GraphDef
|
31 |
yolov5s.tflite # TensorFlow Lite
|
|
|
156 |
LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')
|
157 |
f = file.with_suffix('.mlmodel')
|
158 |
|
|
|
159 |
ts = torch.jit.trace(model, im, strict=False) # TorchScript model
|
160 |
ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
|
161 |
ct_model.save(f)
|
models/common.py
CHANGED
@@ -420,9 +420,12 @@ class DetectMultiBackend(nn.Module):
|
|
420 |
im = Image.fromarray((im[0] * 255).astype('uint8'))
|
421 |
# im = im.resize((192, 320), Image.ANTIALIAS)
|
422 |
y = self.model.predict({'image': im}) # coordinates are xywh normalized
|
423 |
-
|
424 |
-
|
425 |
-
|
|
|
|
|
|
|
426 |
else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
|
427 |
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
|
428 |
if self.saved_model: # SavedModel
|
|
|
420 |
im = Image.fromarray((im[0] * 255).astype('uint8'))
|
421 |
# im = im.resize((192, 320), Image.ANTIALIAS)
|
422 |
y = self.model.predict({'image': im}) # coordinates are xywh normalized
|
423 |
+
if 'confidence' in y:
|
424 |
+
box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels
|
425 |
+
conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)
|
426 |
+
y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
|
427 |
+
else:
|
428 |
+
y = y[list(y)[-1]] # last output
|
429 |
else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
|
430 |
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
|
431 |
if self.saved_model: # SavedModel
|
val.py
CHANGED
@@ -11,7 +11,7 @@ Usage - formats:
|
|
11 |
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
12 |
yolov5s.xml # OpenVINO
|
13 |
yolov5s.engine # TensorRT
|
14 |
-
yolov5s.mlmodel # CoreML (
|
15 |
yolov5s_saved_model # TensorFlow SavedModel
|
16 |
yolov5s.pb # TensorFlow GraphDef
|
17 |
yolov5s.tflite # TensorFlow Lite
|
|
|
11 |
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
12 |
yolov5s.xml # OpenVINO
|
13 |
yolov5s.engine # TensorRT
|
14 |
+
yolov5s.mlmodel # CoreML (MacOS-only)
|
15 |
yolov5s_saved_model # TensorFlow SavedModel
|
16 |
yolov5s.pb # TensorFlow GraphDef
|
17 |
yolov5s.tflite # TensorFlow Lite
|