diff --git a/detect.py b/detect.py index a4a1fb69b42e..2d1963ad6f86 100644 --- a/detect.py +++ b/detect.py @@ -17,7 +17,7 @@ yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn yolov5s.xml # OpenVINO yolov5s.engine # TensorRT - yolov5s.mlmodel # CoreML (under development) + yolov5s.mlmodel # CoreML (MacOS-only) yolov5s_saved_model # TensorFlow SavedModel yolov5s.pb # TensorFlow GraphDef yolov5s.tflite # TensorFlow Lite diff --git a/export.py b/export.py index a0cb5fdc5678..3447fc6ed1ab 100644 --- a/export.py +++ b/export.py @@ -25,7 +25,7 @@ yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn yolov5s.xml # OpenVINO yolov5s.engine # TensorRT - yolov5s.mlmodel # CoreML (under development) + yolov5s.mlmodel # CoreML (MacOS-only) yolov5s_saved_model # TensorFlow SavedModel yolov5s.pb # TensorFlow GraphDef yolov5s.tflite # TensorFlow Lite @@ -156,7 +156,6 @@ def export_coreml(model, im, file, prefix=colorstr('CoreML:')): LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...') f = file.with_suffix('.mlmodel') - model.train() # CoreML exports should be placed in model.train() mode ts = torch.jit.trace(model, im, strict=False) # TorchScript model ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])]) ct_model.save(f) diff --git a/models/common.py b/models/common.py index 836314568f67..d8d5423a16e0 100644 --- a/models/common.py +++ b/models/common.py @@ -420,9 +420,12 @@ def forward(self, im, augment=False, visualize=False, val=False): im = Image.fromarray((im[0] * 255).astype('uint8')) # im = im.resize((192, 320), Image.ANTIALIAS) y = self.model.predict({'image': im}) # coordinates are xywh normalized - box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels - conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float) - y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1) + if 'confidence' in y: + box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels + conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float) + y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1) + else: + y = y[list(y)[-1]] # last output else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU) im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3) if self.saved_model: # SavedModel diff --git a/val.py b/val.py index 4d707f62bffa..4709f67511bb 100644 --- a/val.py +++ b/val.py @@ -11,7 +11,7 @@ yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn yolov5s.xml # OpenVINO yolov5s.engine # TensorRT - yolov5s.mlmodel # CoreML (under development) + yolov5s.mlmodel # CoreML (MacOS-only) yolov5s_saved_model # TensorFlow SavedModel yolov5s.pb # TensorFlow GraphDef yolov5s.tflite # TensorFlow Lite