From 7cc40f684711cddd673a16de329bcfb35e0d8c39 Mon Sep 17 00:00:00 2001 From: Muhammad Rizwan Munawar Date: Wed, 2 Oct 2024 19:51:46 +0500 Subject: [PATCH] Auto-correct for export formats (#16625) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- ultralytics/engine/exporter.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index d4987f90..c2f0f6c6 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -178,6 +178,16 @@ class Exporter: if fmt in {"mlmodel", "mlpackage", "mlprogram", "apple", "ios", "coreml"}: # 'coreml' aliases fmt = "coreml" fmts = tuple(export_formats()["Argument"][1:]) # available export formats + if fmt not in fmts: + import difflib + + # Get the closest match if format is invalid + matches = difflib.get_close_matches(fmt, fmts, n=1, cutoff=0.6) # 60% similarity required to match + if closest_match: + LOGGER.warning(f"WARNING ⚠️ Invalid export format='{fmt}', updating to format='{matches[0]}'") + fmt = closest_match[0] + else: + raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}") flags = [x == fmt for x in fmts] if sum(flags) != 1: raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")