You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
215 lines
8.2 KiB
215 lines
8.2 KiB
# Ultralytics YOLO 🚀, GPL-3.0 license
|
|
import contextlib
|
|
import re
|
|
import shutil
|
|
import sys
|
|
from difflib import get_close_matches
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import Dict, Union
|
|
|
|
from ultralytics import __version__, yolo
|
|
from ultralytics.yolo.utils import (DEFAULT_CFG_DICT, DEFAULT_CFG_PATH, LOGGER, PREFIX, USER_CONFIG_DIR,
|
|
IterableSimpleNamespace, checks, colorstr, yaml_load, yaml_print)
|
|
|
|
CLI_HELP_MSG = \
|
|
"""
|
|
YOLOv8 'yolo' CLI commands use the following syntax:
|
|
|
|
yolo TASK MODE ARGS
|
|
|
|
Where TASK (optional) is one of [detect, segment, classify]
|
|
MODE (required) is one of [train, val, predict, export]
|
|
ARGS (optional) are any number of custom 'arg=value' pairs like 'imgsz=320' that override defaults.
|
|
See all ARGS at https://docs.ultralytics.com/cfg or with 'yolo cfg'
|
|
|
|
1. Train a detection model for 10 epochs with an initial learning_rate of 0.01
|
|
yolo detect train data=coco128.yaml model=yolov8n.pt epochs=10 lr0=0.01
|
|
|
|
2. Predict a YouTube video using a pretrained segmentation model at image size 320:
|
|
yolo segment predict model=yolov8n-seg.pt source=https://youtu.be/Zgi9g1ksQHc imgsz=320
|
|
|
|
3. Val a pretrained detection model at batch-size 1 and image size 640:
|
|
yolo detect val model=yolov8n.pt data=coco128.yaml batch=1 imgsz=640
|
|
|
|
4. Export a YOLOv8n classification model to ONNX format at image size 224 by 128 (no TASK required)
|
|
yolo export model=yolov8n-cls.pt format=onnx imgsz=224,128
|
|
|
|
5. Run special commands:
|
|
yolo help
|
|
yolo checks
|
|
yolo version
|
|
yolo settings
|
|
yolo copy-cfg
|
|
yolo cfg
|
|
|
|
Docs: https://docs.ultralytics.com/cli
|
|
Community: https://community.ultralytics.com
|
|
GitHub: https://github.com/ultralytics/ultralytics
|
|
"""
|
|
|
|
|
|
def cfg2dict(cfg):
|
|
"""
|
|
Convert a configuration object to a dictionary.
|
|
|
|
This function converts a configuration object to a dictionary, whether it is a file path, a string, or a SimpleNamespace object.
|
|
|
|
Inputs:
|
|
cfg (str) or (Path) or (SimpleNamespace): Configuration object to be converted to a dictionary.
|
|
|
|
Returns:
|
|
cfg (dict): Configuration object in dictionary format.
|
|
"""
|
|
if isinstance(cfg, (str, Path)):
|
|
cfg = yaml_load(cfg) # load dict
|
|
elif isinstance(cfg, SimpleNamespace):
|
|
cfg = vars(cfg) # convert to dict
|
|
return cfg
|
|
|
|
|
|
def get_cfg(cfg: Union[str, Path, Dict, SimpleNamespace], overrides: Dict = None):
|
|
"""
|
|
Load and merge configuration data from a file or dictionary.
|
|
|
|
Args:
|
|
cfg (str) or (Path) or (Dict) or (SimpleNamespace): Configuration data.
|
|
overrides (str) or (Dict), optional: Overrides in the form of a file name or a dictionary. Default is None.
|
|
|
|
Returns:
|
|
(SimpleNamespace): Training arguments namespace.
|
|
"""
|
|
cfg = cfg2dict(cfg)
|
|
|
|
# Merge overrides
|
|
if overrides:
|
|
overrides = cfg2dict(overrides)
|
|
check_cfg_mismatch(cfg, overrides)
|
|
cfg = {**cfg, **overrides} # merge cfg and overrides dicts (prefer overrides)
|
|
|
|
# Return instance
|
|
return IterableSimpleNamespace(**cfg)
|
|
|
|
|
|
def check_cfg_mismatch(base: Dict, custom: Dict):
|
|
"""
|
|
This function checks for any mismatched keys between a custom configuration list and a base configuration list.
|
|
If any mismatched keys are found, the function prints out similar keys from the base list and exits the program.
|
|
|
|
Inputs:
|
|
- custom (Dict): a dictionary of custom configuration options
|
|
- base (Dict): a dictionary of base configuration options
|
|
"""
|
|
base, custom = (set(x.keys()) for x in (base, custom))
|
|
mismatched = [x for x in custom if x not in base]
|
|
if mismatched:
|
|
for x in mismatched:
|
|
matches = get_close_matches(x, base, 3, 0.6)
|
|
match_str = f"Similar arguments are {matches}." if matches else 'There are no similar arguments.'
|
|
LOGGER.warning(f"'{colorstr('red', 'bold', x)}' is not a valid YOLO argument. {match_str}")
|
|
LOGGER.warning(CLI_HELP_MSG)
|
|
sys.exit()
|
|
|
|
|
|
def argument_error(arg):
|
|
return SyntaxError(f"'{arg}' is not a valid YOLO argument.\n{CLI_HELP_MSG}")
|
|
|
|
|
|
def entrypoint(debug=False):
|
|
"""
|
|
This function is the ultralytics package entrypoint, it's responsible for parsing the command line arguments passed
|
|
to the package.
|
|
|
|
This function allows for:
|
|
- passing mandatory YOLO args as a list of strings
|
|
- specifying the task to be performed, either 'detect', 'segment' or 'classify'
|
|
- specifying the mode, either 'train', 'val', 'test', or 'predict'
|
|
- running special modes like 'checks'
|
|
- passing overrides to the package's configuration
|
|
|
|
It uses the package's default cfg and initializes it using the passed overrides.
|
|
Then it calls the CLI function with the composed cfg
|
|
"""
|
|
args = ['train', 'predict', 'model=yolov8n.pt'] if debug else sys.argv[1:]
|
|
if not args: # no arguments passed
|
|
LOGGER.info(CLI_HELP_MSG)
|
|
return
|
|
|
|
tasks = 'detect', 'segment', 'classify'
|
|
modes = 'train', 'val', 'predict', 'export'
|
|
special = {
|
|
'help': lambda: LOGGER.info(CLI_HELP_MSG),
|
|
'checks': checks.check_yolo,
|
|
'version': lambda: LOGGER.info(__version__),
|
|
'settings': lambda: yaml_print(USER_CONFIG_DIR / 'settings.yaml'),
|
|
'cfg': lambda: yaml_print(DEFAULT_CFG_PATH),
|
|
'copy-cfg': copy_default_config}
|
|
|
|
overrides = {} # basic overrides, i.e. imgsz=320
|
|
for a in args:
|
|
if '=' in a:
|
|
try:
|
|
re.sub(r' *= *', '=', a) # remove spaces around equals sign
|
|
k, v = a.split('=')
|
|
if k == 'cfg': # custom.yaml passed
|
|
LOGGER.info(f"{PREFIX}Overriding {DEFAULT_CFG_PATH} with {v}")
|
|
overrides = {k: val for k, val in yaml_load(v).items() if k != 'cfg'}
|
|
else:
|
|
if v.lower() == 'none':
|
|
v = None
|
|
elif v.lower() == 'true':
|
|
v = True
|
|
elif v.lower() == 'false':
|
|
v = False
|
|
else:
|
|
with contextlib.suppress(Exception):
|
|
v = eval(v)
|
|
overrides[k] = v
|
|
except (NameError, SyntaxError, ValueError) as e:
|
|
raise argument_error(a) from e
|
|
|
|
elif a in tasks:
|
|
overrides['task'] = a
|
|
elif a in modes:
|
|
overrides['mode'] = a
|
|
elif a in special:
|
|
special[a]()
|
|
return
|
|
elif a in DEFAULT_CFG_DICT and DEFAULT_CFG_DICT[a] is False:
|
|
overrides[a] = True # auto-True for default False args, i.e. 'yolo show' sets show=True
|
|
elif a in DEFAULT_CFG_DICT:
|
|
raise SyntaxError(f"'{colorstr('red', 'bold', a)}' is a valid YOLO argument but is missing an '=' sign "
|
|
f"to set its value, i.e. try '{a}={DEFAULT_CFG_DICT[a]}'\n{CLI_HELP_MSG}")
|
|
else:
|
|
raise argument_error(a)
|
|
|
|
cfg = get_cfg(DEFAULT_CFG_DICT, overrides) # create CFG instance
|
|
|
|
# Mapping from task to module
|
|
module = {"detect": yolo.v8.detect, "segment": yolo.v8.segment, "classify": yolo.v8.classify}.get(cfg.task)
|
|
if not module:
|
|
raise SyntaxError(f"yolo task={cfg.task} is invalid. Valid tasks are: {', '.join(tasks)}\n{CLI_HELP_MSG}")
|
|
|
|
# Mapping from mode to function
|
|
func = {
|
|
"train": module.train,
|
|
"val": module.val,
|
|
"predict": module.predict,
|
|
"export": yolo.engine.exporter.export}.get(cfg.mode)
|
|
if not func:
|
|
raise SyntaxError(f"yolo mode={cfg.mode} is invalid. Valid modes are: {', '.join(modes)}\n{CLI_HELP_MSG}")
|
|
|
|
func(cfg)
|
|
|
|
|
|
# Special modes --------------------------------------------------------------------------------------------------------
|
|
def copy_default_config():
|
|
new_file = Path.cwd() / DEFAULT_CFG_PATH.name.replace('.yaml', '_copy.yaml')
|
|
shutil.copy2(DEFAULT_CFG_PATH, new_file)
|
|
LOGGER.info(f"{PREFIX}{DEFAULT_CFG_PATH} copied to {new_file}\n"
|
|
f"Example YOLO command with this new custom cfg:\n yolo cfg='{new_file}' imgsz=320 batch=8")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
entrypoint(debug=True)
|