Auto-correct for export formats (#16625)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/16638/head
Muhammad Rizwan Munawar 5 months ago committed by GitHub
parent 0e7d95cb29
commit 7cc40f6847
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 10
      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}")

Loading…
Cancel
Save