glenn-jocher commited on
Commit
b5b56a3
1 Parent(s): 7b31a53

Add CoreML inference (#6195)

Browse files

* Add Apple CoreML inference

* Cleanup

Files changed (4) hide show
  1. detect.py +1 -1
  2. export.py +1 -2
  3. models/common.py +6 -3
  4. val.py +1 -1
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 (under development)
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 (under development)
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
- box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels
424
- conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)
425
- y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
 
 
 
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 (under development)
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