diff --git a/utils/datasets.py b/utils/datasets.py index cfcb19b1..5ec67a44 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -15,7 +15,9 @@ from utils.utils import xyxy2xywh class load_images(): # for inference def __init__(self, path, batch_size=1, img_size=416): if os.path.isdir(path): + image_format = ['.jpg', '.jpeg', '.png', '.tif'] self.files = sorted(glob.glob('%s/*.*' % path)) + self.files = list(filter(lambda x: os.path.splitext(x)[1].lower() in image_format, self.files)) elif os.path.isfile(path): self.files = [path] diff --git a/utils/onnx2coreml.py b/utils/onnx2coreml.py index 2e18a763..4c00c20a 100644 --- a/utils/onnx2coreml.py +++ b/utils/onnx2coreml.py @@ -42,8 +42,11 @@ def main(): spec = coremltools.utils.convert_neural_network_spec_weights_to_fp16(spec) yolov3_model = coremltools.models.MLModel(spec) + name_out0 = spec.description.output[0].name + name_out1 = spec.description.output[1].name + num_classes = 80 - num_anchors = 507 + num_anchors = 507 # 507 for yolov3-tiny, spec.description.output[0].type.multiArrayType.shape.append(num_anchors) spec.description.output[0].type.multiArrayType.shape.append(num_classes) # spec.description.output[0].type.multiArrayType.shape.append(1) @@ -75,7 +78,7 @@ def main(): from PIL import Image img = Image.open('../yolov3/data/samples/zidane_416.jpg') out = yolov3_model.predict({'0': img}, useCPUOnly=True) - print(out['148'].shape, out['150'].shape) + print(out[name_out0].shape, out[name_out1].shape) # 3. Create NMS protobuf import numpy as np @@ -107,15 +110,15 @@ def main(): del ma_type.shape[:] nms = nms_spec.nonMaximumSuppression - nms.confidenceInputFeatureName = '148' # 1x507x80 - nms.coordinatesInputFeatureName = '150' # 1x507x4 + nms.confidenceInputFeatureName = name_out0 # 1x507x80 + nms.coordinatesInputFeatureName = name_out1 # 1x507x4 nms.confidenceOutputFeatureName = 'confidence' nms.coordinatesOutputFeatureName = 'coordinates' nms.iouThresholdInputFeatureName = 'iouThreshold' nms.confidenceThresholdInputFeatureName = 'confidenceThreshold' - nms.iouThreshold = 0.6 - nms.confidenceThreshold = 0.3 + nms.iouThreshold = 0.4 + nms.confidenceThreshold = 0.5 nms.pickTop.perClass = True labels = np.loadtxt('../yolov3/data/coco.names', dtype=str, delimiter='\n')