diff --git a/detect.py b/detect.py index 108f8f138052..cb3c4abd6425 100644 --- a/detect.py +++ b/detect.py @@ -203,7 +203,7 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s) def parse_opt(): parser = argparse.ArgumentParser() - parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)') + parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.mlmodel', help='model path(s)') parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob, 0 for webcam') parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w') parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold') diff --git a/models/common.py b/models/common.py index 9982a8d85c22..6fa7bdb33622 100644 --- a/models/common.py +++ b/models/common.py @@ -349,6 +349,9 @@ def forward(self, im, augment=False, visualize=False, val=False): y = self.model(im) if self.jit else self.model(im, augment=augment, visualize=visualize) return y if val else y[0] elif self.coreml: # CoreML *.mlmodel + im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3) + 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) @@ -361,7 +364,7 @@ def forward(self, im, augment=False, visualize=False, val=False): else: # ONNX Runtime y = self.session.run([self.session.get_outputs()[0].name], {self.session.get_inputs()[0].name: im})[0] else: # TensorFlow model (TFLite, pb, saved_model) - im = im.permute(0, 2, 3, 1).cpu().numpy() # TF format (1,h=640,w=640,3) + im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3) if self.pb: y = self.frozen_func(x=self.tf.constant(im)).numpy() elif self.saved_model: