diff --git a/export.py b/export.py index e29d62743..01e66839b 100644 --- a/export.py +++ b/export.py @@ -520,7 +520,7 @@ def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")): @try_export -def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")): +def export_coreml(model, im, file, int8, half, nms, mlmodel, prefix=colorstr("CoreML:")): """ Export a YOLOv5 model to CoreML format with optional NMS, INT8, and FP16 support. @@ -531,6 +531,7 @@ def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")): int8 (bool): Flag indicating whether to use INT8 quantization (default is False). half (bool): Flag indicating whether to use FP16 quantization (default is False). nms (bool): Flag indicating whether to include Non-Maximum Suppression (default is False). + mlmodel (bool): Flag indicating whether to export as older *.mlmodel format (default is False). prefix (str): Prefix string for logging purposes (default is 'CoreML:'). Returns: @@ -548,12 +549,11 @@ def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")): model = Model(cfg, ch=3, nc=80) im = torch.randn(1, 3, 640, 640) file = Path("yolov5s_coreml") - export_coreml(model, im, file, int8=False, half=False, nms=True) + export_coreml(model, im, file, int8=False, half=False, nms=True, mlmodel=False) ``` """ check_requirements("coremltools") import coremltools as ct - mlmodel = False LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...") if mlmodel: @@ -1084,7 +1084,7 @@ def add_tflite_metadata(file, metadata, num_outputs): tmp_file.unlink() -def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")): +def pipeline_coreml(model, im, file, names, y, mlmodel, prefix=colorstr("CoreML Pipeline:")): """ Convert a PyTorch YOLOv5 model to CoreML format with Non-Maximum Suppression (NMS), handling different input/output shapes, and saving the model. @@ -1096,6 +1096,7 @@ def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline: file (Path): Path to save the converted CoreML model. names (dict[int, str]): Dictionary mapping class indices to class names. y (torch.Tensor): Output tensor from the PyTorch model's forward pass. + mlmodel (bool): Flag indicating whether to export as older *.mlmodel format (default is False). prefix (str): Custom prefix for logging messages. Returns: @@ -1128,7 +1129,6 @@ def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline: import coremltools as ct from PIL import Image - mlmodel = False if mlmodel: f = file.with_suffix(".mlmodel") # filename else: @@ -1280,6 +1280,7 @@ def run( per_tensor=False, # TF per tensor quantization dynamic=False, # ONNX/TF/TensorRT: dynamic axes simplify=False, # ONNX: simplify model + mlmodel=False, # CoreML: Export in *.mlmodel format opset=12, # ONNX: opset version verbose=False, # TensorRT: verbose log workspace=4, # TensorRT: workspace size (GB) @@ -1317,6 +1318,7 @@ def run( topk_all (int): Top-K boxes for all classes to keep for TensorFlow.js NMS. Default is 100. iou_thres (float): IoU threshold for NMS. Default is 0.45. conf_thres (float): Confidence threshold for NMS. Default is 0.25. + mlmodel (bool): Flag to use *.mlmodel for CoreML export. Default is False. Returns: None @@ -1344,6 +1346,7 @@ def run( simplify=False, opset=12, verbose=False, + mlmodel=False, workspace=4, nms=False, agnostic_nms=False, @@ -1407,9 +1410,9 @@ def run( if xml: # OpenVINO f[3], _ = export_openvino(file, metadata, half, int8, data) if coreml: # CoreML - f[4], ct_model = export_coreml(model, im, file, int8, half, nms) + f[4], ct_model = export_coreml(model, im, file, int8, half, nms, mlmodel) if nms: - pipeline_coreml(ct_model, im, file, model.names, y) + pipeline_coreml(ct_model, im, file, model.names, y, mlmodel) if any((saved_model, pb, tflite, edgetpu, tfjs)): # TensorFlow formats assert not tflite or not tfjs, "TFLite and TF.js models must be exported separately, please pass only one type." assert not isinstance(model, ClassificationModel), "ClassificationModel export to TF formats not yet supported." @@ -1497,6 +1500,7 @@ def parse_opt(known=False): parser.add_argument("--per-tensor", action="store_true", help="TF per-tensor quantization") parser.add_argument("--dynamic", action="store_true", help="ONNX/TF/TensorRT: dynamic axes") parser.add_argument("--simplify", action="store_true", help="ONNX: simplify model") + parser.add_argument("--mlmodel", action="store_true", help="CoreML: Export in *.mlmodel format") parser.add_argument("--opset", type=int, default=17, help="ONNX: opset version") parser.add_argument("--verbose", action="store_true", help="TensorRT: verbose log") parser.add_argument("--workspace", type=int, default=4, help="TensorRT: workspace size (GB)")