__init__.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447
  1. # Ultralytics YOLO 🚀, AGPL-3.0 license
  2. import contextlib
  3. import re
  4. import shutil
  5. import sys
  6. from difflib import get_close_matches
  7. from pathlib import Path
  8. from types import SimpleNamespace
  9. from typing import Dict, List, Union
  10. from ultralytics.utils import (ASSETS, DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_PATH, LOGGER, SETTINGS, SETTINGS_YAML,
  11. IterableSimpleNamespace, __version__, checks, colorstr, deprecation_warn, yaml_load,
  12. yaml_print)
  13. # Define valid tasks and modes
  14. MODES = 'train', 'val', 'predict', 'export', 'track', 'benchmark'
  15. TASKS = 'detect', 'segment', 'classify', 'pose'
  16. TASK2DATA = {'detect': 'coco8.yaml', 'segment': 'coco8-seg.yaml', 'classify': 'imagenet100', 'pose': 'coco8-pose.yaml'}
  17. TASK2MODEL = {
  18. 'detect': 'yolov8n.pt',
  19. 'segment': 'yolov8n-seg.pt',
  20. 'classify': 'yolov8n-cls.pt',
  21. 'pose': 'yolov8n-pose.pt'}
  22. TASK2METRIC = {
  23. 'detect': 'metrics/mAP50-95(B)',
  24. 'segment': 'metrics/mAP50-95(M)',
  25. 'classify': 'metrics/accuracy_top1',
  26. 'pose': 'metrics/mAP50-95(P)'}
  27. CLI_HELP_MSG = \
  28. f"""
  29. Arguments received: {str(['yolo'] + sys.argv[1:])}. Ultralytics 'yolo' commands use the following syntax:
  30. yolo TASK MODE ARGS
  31. Where TASK (optional) is one of {TASKS}
  32. MODE (required) is one of {MODES}
  33. ARGS (optional) are any number of custom 'arg=value' pairs like 'imgsz=320' that override defaults.
  34. See all ARGS at https://docs.ultralytics.com/usage/cfg or with 'yolo cfg'
  35. 1. Train a detection model for 10 epochs with an initial learning_rate of 0.01
  36. yolo train data=coco128.yaml model=yolov8n.pt epochs=10 lr0=0.01
  37. 2. Predict a YouTube video using a pretrained segmentation model at image size 320:
  38. yolo predict model=yolov8n-seg.pt source='https://youtu.be/Zgi9g1ksQHc' imgsz=320
  39. 3. Val a pretrained detection model at batch-size 1 and image size 640:
  40. yolo val model=yolov8n.pt data=coco128.yaml batch=1 imgsz=640
  41. 4. Export a YOLOv8n classification model to ONNX format at image size 224 by 128 (no TASK required)
  42. yolo export model=yolov8n-cls.pt format=onnx imgsz=224,128
  43. 5. Run special commands:
  44. yolo help
  45. yolo checks
  46. yolo version
  47. yolo settings
  48. yolo copy-cfg
  49. yolo cfg
  50. Docs: https://docs.ultralytics.com
  51. Community: https://community.ultralytics.com
  52. GitHub: https://github.com/ultralytics/ultralytics
  53. """
  54. # Define keys for arg type checks
  55. CFG_FLOAT_KEYS = 'warmup_epochs', 'box', 'cls', 'dfl', 'degrees', 'shear'
  56. CFG_FRACTION_KEYS = ('dropout', 'iou', 'lr0', 'lrf', 'momentum', 'weight_decay', 'warmup_momentum', 'warmup_bias_lr',
  57. 'label_smoothing', 'hsv_h', 'hsv_s', 'hsv_v', 'translate', 'scale', 'perspective', 'flipud',
  58. 'fliplr', 'mosaic', 'mixup', 'copy_paste', 'conf', 'iou', 'fraction') # fraction floats 0.0 - 1.0
  59. CFG_INT_KEYS = ('epochs', 'patience', 'batch', 'workers', 'seed', 'close_mosaic', 'mask_ratio', 'max_det', 'vid_stride',
  60. 'line_width', 'workspace', 'nbs', 'save_period')
  61. CFG_BOOL_KEYS = ('save', 'exist_ok', 'verbose', 'deterministic', 'single_cls', 'rect', 'cos_lr', 'overlap_mask', 'val',
  62. 'save_json', 'save_hybrid', 'half', 'dnn', 'plots', 'show', 'save_txt', 'save_conf', 'save_crop',
  63. 'show_labels', 'show_conf', 'visualize', 'augment', 'agnostic_nms', 'retina_masks', 'boxes', 'keras',
  64. 'optimize', 'int8', 'dynamic', 'simplify', 'nms', 'profile')
  65. def cfg2dict(cfg):
  66. """
  67. Convert a configuration object to a dictionary, whether it is a file path, a string, or a SimpleNamespace object.
  68. Args:
  69. cfg (str | Path | dict | SimpleNamespace): Configuration object to be converted to a dictionary.
  70. Returns:
  71. cfg (dict): Configuration object in dictionary format.
  72. """
  73. if isinstance(cfg, (str, Path)):
  74. cfg = yaml_load(cfg) # load dict
  75. elif isinstance(cfg, SimpleNamespace):
  76. cfg = vars(cfg) # convert to dict
  77. return cfg
  78. def get_cfg(cfg: Union[str, Path, Dict, SimpleNamespace] = DEFAULT_CFG_DICT, overrides: Dict = None):
  79. """
  80. Load and merge configuration data from a file or dictionary.
  81. Args:
  82. cfg (str | Path | Dict | SimpleNamespace): Configuration data.
  83. overrides (str | Dict | optional): Overrides in the form of a file name or a dictionary. Default is None.
  84. Returns:
  85. (SimpleNamespace): Training arguments namespace.
  86. """
  87. cfg = cfg2dict(cfg)
  88. # Merge overrides
  89. if overrides:
  90. overrides = cfg2dict(overrides)
  91. if 'save_dir' not in cfg:
  92. overrides.pop('save_dir', None) # special override keys to ignore
  93. check_dict_alignment(cfg, overrides)
  94. cfg = {**cfg, **overrides} # merge cfg and overrides dicts (prefer overrides)
  95. # Special handling for numeric project/name
  96. for k in 'project', 'name':
  97. if k in cfg and isinstance(cfg[k], (int, float)):
  98. cfg[k] = str(cfg[k])
  99. if cfg.get('name') == 'model': # assign model to 'name' arg
  100. cfg['name'] = cfg.get('model', '').split('.')[0]
  101. LOGGER.warning(f"WARNING ⚠️ 'name=model' automatically updated to 'name={cfg['name']}'.")
  102. # Type and Value checks
  103. for k, v in cfg.items():
  104. if v is not None: # None values may be from optional args
  105. if k in CFG_FLOAT_KEYS and not isinstance(v, (int, float)):
  106. raise TypeError(f"'{k}={v}' is of invalid type {type(v).__name__}. "
  107. f"Valid '{k}' types are int (i.e. '{k}=0') or float (i.e. '{k}=0.5')")
  108. elif k in CFG_FRACTION_KEYS:
  109. if not isinstance(v, (int, float)):
  110. raise TypeError(f"'{k}={v}' is of invalid type {type(v).__name__}. "
  111. f"Valid '{k}' types are int (i.e. '{k}=0') or float (i.e. '{k}=0.5')")
  112. if not (0.0 <= v <= 1.0):
  113. raise ValueError(f"'{k}={v}' is an invalid value. "
  114. f"Valid '{k}' values are between 0.0 and 1.0.")
  115. elif k in CFG_INT_KEYS and not isinstance(v, int):
  116. raise TypeError(f"'{k}={v}' is of invalid type {type(v).__name__}. "
  117. f"'{k}' must be an int (i.e. '{k}=8')")
  118. elif k in CFG_BOOL_KEYS and not isinstance(v, bool):
  119. raise TypeError(f"'{k}={v}' is of invalid type {type(v).__name__}. "
  120. f"'{k}' must be a bool (i.e. '{k}=True' or '{k}=False')")
  121. # Return instance
  122. return IterableSimpleNamespace(**cfg)
  123. def _handle_deprecation(custom):
  124. """Hardcoded function to handle deprecated config keys"""
  125. for key in custom.copy().keys():
  126. if key == 'hide_labels':
  127. deprecation_warn(key, 'show_labels')
  128. custom['show_labels'] = custom.pop('hide_labels') == 'False'
  129. if key == 'hide_conf':
  130. deprecation_warn(key, 'show_conf')
  131. custom['show_conf'] = custom.pop('hide_conf') == 'False'
  132. if key == 'line_thickness':
  133. deprecation_warn(key, 'line_width')
  134. custom['line_width'] = custom.pop('line_thickness')
  135. return custom
  136. def check_dict_alignment(base: Dict, custom: Dict, e=None):
  137. """
  138. This function checks for any mismatched keys between a custom configuration list and a base configuration list.
  139. If any mismatched keys are found, the function prints out similar keys from the base list and exits the program.
  140. Args:
  141. custom (dict): a dictionary of custom configuration options
  142. base (dict): a dictionary of base configuration options
  143. """
  144. custom = _handle_deprecation(custom)
  145. base_keys, custom_keys = (set(x.keys()) for x in (base, custom))
  146. mismatched = [k for k in custom_keys if k not in base_keys]
  147. if mismatched:
  148. string = ''
  149. for x in mismatched:
  150. matches = get_close_matches(x, base_keys) # key list
  151. matches = [f'{k}={base[k]}' if base.get(k) is not None else k for k in matches]
  152. match_str = f'Similar arguments are i.e. {matches}.' if matches else ''
  153. string += f"'{colorstr('red', 'bold', x)}' is not a valid YOLO argument. {match_str}\n"
  154. raise SyntaxError(string + CLI_HELP_MSG) from e
  155. def merge_equals_args(args: List[str]) -> List[str]:
  156. """
  157. Merges arguments around isolated '=' args in a list of strings.
  158. The function considers cases where the first argument ends with '=' or the second starts with '=',
  159. as well as when the middle one is an equals sign.
  160. Args:
  161. args (List[str]): A list of strings where each element is an argument.
  162. Returns:
  163. List[str]: A list of strings where the arguments around isolated '=' are merged.
  164. """
  165. new_args = []
  166. for i, arg in enumerate(args):
  167. if arg == '=' and 0 < i < len(args) - 1: # merge ['arg', '=', 'val']
  168. new_args[-1] += f'={args[i + 1]}'
  169. del args[i + 1]
  170. elif arg.endswith('=') and i < len(args) - 1 and '=' not in args[i + 1]: # merge ['arg=', 'val']
  171. new_args.append(f'{arg}{args[i + 1]}')
  172. del args[i + 1]
  173. elif arg.startswith('=') and i > 0: # merge ['arg', '=val']
  174. new_args[-1] += arg
  175. else:
  176. new_args.append(arg)
  177. return new_args
  178. def handle_yolo_hub(args: List[str]) -> None:
  179. """
  180. Handle Ultralytics HUB command-line interface (CLI) commands.
  181. This function processes Ultralytics HUB CLI commands such as login and logout.
  182. It should be called when executing a script with arguments related to HUB authentication.
  183. Args:
  184. args (List[str]): A list of command line arguments
  185. Example:
  186. ```bash
  187. python my_script.py hub login your_api_key
  188. ```
  189. """
  190. from ultralytics import hub
  191. if args[0] == 'login':
  192. key = args[1] if len(args) > 1 else ''
  193. # Log in to Ultralytics HUB using the provided API key
  194. hub.login(key)
  195. elif args[0] == 'logout':
  196. # Log out from Ultralytics HUB
  197. hub.logout()
  198. def handle_yolo_settings(args: List[str]) -> None:
  199. """
  200. Handle YOLO settings command-line interface (CLI) commands.
  201. This function processes YOLO settings CLI commands such as reset.
  202. It should be called when executing a script with arguments related to YOLO settings management.
  203. Args:
  204. args (List[str]): A list of command line arguments for YOLO settings management.
  205. Example:
  206. ```bash
  207. python my_script.py yolo settings reset
  208. ```
  209. """
  210. url = 'https://docs.ultralytics.com/quickstart/#ultralytics-settings' # help URL
  211. try:
  212. if any(args):
  213. if args[0] == 'reset':
  214. SETTINGS_YAML.unlink() # delete the settings file
  215. SETTINGS.reset() # create new settings
  216. LOGGER.info('Settings reset successfully') # inform the user that settings have been reset
  217. else: # save a new setting
  218. new = dict(parse_key_value_pair(a) for a in args)
  219. check_dict_alignment(SETTINGS, new)
  220. SETTINGS.update(new)
  221. LOGGER.info(f'💡 Learn about settings at {url}')
  222. yaml_print(SETTINGS_YAML) # print the current settings
  223. except Exception as e:
  224. LOGGER.warning(f"WARNING ⚠️ settings error: '{e}'. Please see {url} for help.")
  225. def parse_key_value_pair(pair):
  226. """Parse one 'key=value' pair and return key and value."""
  227. re.sub(r' *= *', '=', pair) # remove spaces around equals sign
  228. k, v = pair.split('=', 1) # split on first '=' sign
  229. assert v, f"missing '{k}' value"
  230. return k, smart_value(v)
  231. def smart_value(v):
  232. """Convert a string to an underlying type such as int, float, bool, etc."""
  233. if v.lower() == 'none':
  234. return None
  235. elif v.lower() == 'true':
  236. return True
  237. elif v.lower() == 'false':
  238. return False
  239. else:
  240. with contextlib.suppress(Exception):
  241. return eval(v)
  242. return v
  243. def entrypoint(debug=''):
  244. """
  245. This function is the ultralytics package entrypoint, it's responsible for parsing the command line arguments passed
  246. to the package.
  247. This function allows for:
  248. - passing mandatory YOLO args as a list of strings
  249. - specifying the task to be performed, either 'detect', 'segment' or 'classify'
  250. - specifying the mode, either 'train', 'val', 'test', or 'predict'
  251. - running special modes like 'checks'
  252. - passing overrides to the package's configuration
  253. It uses the package's default cfg and initializes it using the passed overrides.
  254. Then it calls the CLI function with the composed cfg
  255. """
  256. args = (debug.split(' ') if debug else sys.argv)[1:]
  257. if not args: # no arguments passed
  258. LOGGER.info(CLI_HELP_MSG)
  259. return
  260. special = {
  261. 'help': lambda: LOGGER.info(CLI_HELP_MSG),
  262. 'checks': checks.check_yolo,
  263. 'version': lambda: LOGGER.info(__version__),
  264. 'settings': lambda: handle_yolo_settings(args[1:]),
  265. 'cfg': lambda: yaml_print(DEFAULT_CFG_PATH),
  266. 'hub': lambda: handle_yolo_hub(args[1:]),
  267. 'login': lambda: handle_yolo_hub(args),
  268. 'copy-cfg': copy_default_cfg}
  269. full_args_dict = {**DEFAULT_CFG_DICT, **{k: None for k in TASKS}, **{k: None for k in MODES}, **special}
  270. # Define common mis-uses of special commands, i.e. -h, -help, --help
  271. special.update({k[0]: v for k, v in special.items()}) # singular
  272. special.update({k[:-1]: v for k, v in special.items() if len(k) > 1 and k.endswith('s')}) # singular
  273. special = {**special, **{f'-{k}': v for k, v in special.items()}, **{f'--{k}': v for k, v in special.items()}}
  274. overrides = {} # basic overrides, i.e. imgsz=320
  275. for a in merge_equals_args(args): # merge spaces around '=' sign
  276. if a.startswith('--'):
  277. LOGGER.warning(f"WARNING ⚠️ '{a}' does not require leading dashes '--', updating to '{a[2:]}'.")
  278. a = a[2:]
  279. if a.endswith(','):
  280. LOGGER.warning(f"WARNING ⚠️ '{a}' does not require trailing comma ',', updating to '{a[:-1]}'.")
  281. a = a[:-1]
  282. if '=' in a:
  283. try:
  284. k, v = parse_key_value_pair(a)
  285. if k == 'cfg': # custom.yaml passed
  286. LOGGER.info(f'Overriding {DEFAULT_CFG_PATH} with {v}')
  287. overrides = {k: val for k, val in yaml_load(checks.check_yaml(v)).items() if k != 'cfg'}
  288. else:
  289. overrides[k] = v
  290. except (NameError, SyntaxError, ValueError, AssertionError) as e:
  291. check_dict_alignment(full_args_dict, {a: ''}, e)
  292. elif a in TASKS:
  293. overrides['task'] = a
  294. elif a in MODES:
  295. overrides['mode'] = a
  296. elif a.lower() in special:
  297. special[a.lower()]()
  298. return
  299. elif a in DEFAULT_CFG_DICT and isinstance(DEFAULT_CFG_DICT[a], bool):
  300. overrides[a] = True # auto-True for default bool args, i.e. 'yolo show' sets show=True
  301. elif a in DEFAULT_CFG_DICT:
  302. raise SyntaxError(f"'{colorstr('red', 'bold', a)}' is a valid YOLO argument but is missing an '=' sign "
  303. f"to set its value, i.e. try '{a}={DEFAULT_CFG_DICT[a]}'\n{CLI_HELP_MSG}")
  304. else:
  305. check_dict_alignment(full_args_dict, {a: ''})
  306. # Check keys
  307. check_dict_alignment(full_args_dict, overrides)
  308. # Mode
  309. mode = overrides.get('mode')
  310. if mode is None:
  311. mode = DEFAULT_CFG.mode or 'predict'
  312. LOGGER.warning(f"WARNING ⚠️ 'mode' is missing. Valid modes are {MODES}. Using default 'mode={mode}'.")
  313. elif mode not in MODES:
  314. if mode not in ('checks', checks):
  315. raise ValueError(f"Invalid 'mode={mode}'. Valid modes are {MODES}.\n{CLI_HELP_MSG}")
  316. LOGGER.warning("WARNING ⚠️ 'yolo mode=checks' is deprecated. Use 'yolo checks' instead.")
  317. checks.check_yolo()
  318. return
  319. # Task
  320. task = overrides.pop('task', None)
  321. if task:
  322. if task not in TASKS:
  323. raise ValueError(f"Invalid 'task={task}'. Valid tasks are {TASKS}.\n{CLI_HELP_MSG}")
  324. if 'model' not in overrides:
  325. overrides['model'] = TASK2MODEL[task]
  326. # Model
  327. model = overrides.pop('model', DEFAULT_CFG.model)
  328. if model is None:
  329. model = 'yolov8n.pt'
  330. LOGGER.warning(f"WARNING ⚠️ 'model' is missing. Using default 'model={model}'.")
  331. overrides['model'] = model
  332. if 'rtdetr' in model.lower(): # guess architecture
  333. from ultralytics import RTDETR
  334. model = RTDETR(model) # no task argument
  335. elif 'fastsam' in model.lower():
  336. from ultralytics import FastSAM
  337. model = FastSAM(model)
  338. elif 'sam' in model.lower():
  339. from ultralytics import SAM
  340. model = SAM(model)
  341. else:
  342. from ultralytics import YOLO
  343. model = YOLO(model, task=task)
  344. if isinstance(overrides.get('pretrained'), str):
  345. model.load(overrides['pretrained'])
  346. # Task Update
  347. if task != model.task:
  348. if task:
  349. LOGGER.warning(f"WARNING ⚠️ conflicting 'task={task}' passed with 'task={model.task}' model. "
  350. f"Ignoring 'task={task}' and updating to 'task={model.task}' to match model.")
  351. task = model.task
  352. # Mode
  353. if mode in ('predict', 'track') and 'source' not in overrides:
  354. overrides['source'] = DEFAULT_CFG.source or ASSETS
  355. LOGGER.warning(f"WARNING ⚠️ 'source' is missing. Using default 'source={overrides['source']}'.")
  356. elif mode in ('train', 'val'):
  357. if 'data' not in overrides and 'resume' not in overrides:
  358. overrides['data'] = TASK2DATA.get(task or DEFAULT_CFG.task, DEFAULT_CFG.data)
  359. LOGGER.warning(f"WARNING ⚠️ 'data' is missing. Using default 'data={overrides['data']}'.")
  360. elif mode == 'export':
  361. if 'format' not in overrides:
  362. overrides['format'] = DEFAULT_CFG.format or 'torchscript'
  363. LOGGER.warning(f"WARNING ⚠️ 'format' is missing. Using default 'format={overrides['format']}'.")
  364. # Run command in python
  365. # getattr(model, mode)(**vars(get_cfg(overrides=overrides))) # default args using default.yaml
  366. getattr(model, mode)(**overrides) # default args from model
  367. # Special modes --------------------------------------------------------------------------------------------------------
  368. def copy_default_cfg():
  369. """Copy and create a new default configuration file with '_copy' appended to its name."""
  370. new_file = Path.cwd() / DEFAULT_CFG_PATH.name.replace('.yaml', '_copy.yaml')
  371. shutil.copy2(DEFAULT_CFG_PATH, new_file)
  372. LOGGER.info(f'{DEFAULT_CFG_PATH} copied to {new_file}\n'
  373. f"Example YOLO command with this new custom cfg:\n yolo cfg='{new_file}' imgsz=320 batch=8")
  374. if __name__ == '__main__':
  375. # Example: entrypoint(debug='yolo predict model=yolov8n.pt')
  376. entrypoint(debug='')