mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-23 05:24:22 +08:00
CLI updates (#58)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
parent
c5f5b80c04
commit
d0b0fe2592
2
setup.py
2
setup.py
@ -49,4 +49,4 @@ setup(
|
||||
keywords="machine-learning, deep-learning, vision, ML, DL, AI, YOLO, YOLOv3, YOLOv5, YOLOv8, HUB, Ultralytics",
|
||||
entry_points={
|
||||
'console_scripts': [
|
||||
'yolo = ultralytics.yolo.__init__:cli',],})
|
||||
'yolo = ultralytics.yolo.cli:cli',],})
|
||||
|
@ -1,39 +1,5 @@
|
||||
import hydra
|
||||
|
||||
import ultralytics
|
||||
import ultralytics.yolo.v8 as yolo
|
||||
|
||||
from .engine.model import YOLO
|
||||
from .engine.trainer import DEFAULT_CONFIG, BaseTrainer
|
||||
from .engine.trainer import BaseTrainer
|
||||
from .engine.validator import BaseValidator
|
||||
from .utils import LOGGER
|
||||
|
||||
__all__ = ["BaseTrainer", "BaseValidator", "YOLO"] # allow simpler import
|
||||
|
||||
|
||||
@hydra.main(version_base=None, config_path="utils/configs", config_name="default")
|
||||
def cli(cfg):
|
||||
LOGGER.info(f"using Ultralytics YOLO v{ultralytics.__version__}")
|
||||
module_file = None
|
||||
if cfg.task.lower() == "detect":
|
||||
module_file = yolo.detect
|
||||
elif cfg.task.lower() == "segment":
|
||||
module_file = yolo.segment
|
||||
elif cfg.task.lower() == "classify":
|
||||
module_file = yolo.classify
|
||||
|
||||
if not module_file:
|
||||
raise Exception("task not recognized. Choices are `'detect', 'segment', 'classify'`")
|
||||
|
||||
module_function = None
|
||||
|
||||
if cfg.mode.lower() == "train":
|
||||
module_function = module_file.train
|
||||
elif cfg.mode.lower() == "val":
|
||||
module_function = module_file.val
|
||||
elif cfg.mode.lower() == "infer":
|
||||
module_function = module_file.infer
|
||||
|
||||
if not module_function:
|
||||
raise Exception("mode not recognized. Choices are `'train', 'val', 'infer'`")
|
||||
module_function(cfg)
|
||||
|
47
ultralytics/yolo/cli.py
Normal file
47
ultralytics/yolo/cli.py
Normal file
@ -0,0 +1,47 @@
|
||||
import os
|
||||
import shutil
|
||||
|
||||
import hydra
|
||||
|
||||
import ultralytics
|
||||
import ultralytics.yolo.v8 as yolo
|
||||
from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG
|
||||
|
||||
from .utils import LOGGER, colorstr
|
||||
|
||||
|
||||
@hydra.main(version_base=None, config_path="utils/configs", config_name="default")
|
||||
def cli(cfg):
|
||||
LOGGER.info(f"{colorstr(f'Ultralytics YOLO v{ultralytics.__version__}')}")
|
||||
|
||||
module_file = None
|
||||
if cfg.task.lower() == "init": # special case
|
||||
shutil.copy2(DEFAULT_CONFIG, os.getcwd())
|
||||
LOGGER.info(f"""
|
||||
{colorstr("YOLO :")} configuration saved to {os.getcwd()}/{DEFAULT_CONFIG.name}.
|
||||
To run experiments using custom configuration:
|
||||
yolo task='task' mode='mode' --config-name config_file.yaml
|
||||
""")
|
||||
return
|
||||
elif cfg.task.lower() == "detect":
|
||||
module_file = yolo.detect
|
||||
elif cfg.task.lower() == "segment":
|
||||
module_file = yolo.segment
|
||||
elif cfg.task.lower() == "classify":
|
||||
module_file = yolo.classify
|
||||
|
||||
if not module_file:
|
||||
raise Exception("task not recognized. Choices are `'detect', 'segment', 'classify'`")
|
||||
|
||||
module_function = None
|
||||
|
||||
if cfg.mode.lower() == "train":
|
||||
module_function = module_file.train
|
||||
elif cfg.mode.lower() == "val":
|
||||
module_function = module_file.val
|
||||
elif cfg.mode.lower() == "infer":
|
||||
module_function = module_file.infer
|
||||
|
||||
if not module_function:
|
||||
raise Exception("mode not recognized. Choices are `'train', 'val', 'infer'`")
|
||||
module_function(cfg)
|
@ -3,7 +3,6 @@ Top-level YOLO model interface. First principle usage example - https://github.c
|
||||
"""
|
||||
import yaml
|
||||
|
||||
import ultralytics.yolo as yolo
|
||||
from ultralytics.yolo.utils import LOGGER
|
||||
from ultralytics.yolo.utils.checks import check_yaml
|
||||
from ultralytics.yolo.utils.modeling import get_model
|
||||
|
@ -1,5 +1,5 @@
|
||||
import contextlib
|
||||
import logging
|
||||
import logging.config
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
|
@ -2,7 +2,7 @@
|
||||
# Default training settings and hyperparameters for medium-augmentation COCO training
|
||||
|
||||
# Task and Mode
|
||||
task: "classify" # choices=['detect', 'segment', 'classify']
|
||||
task: "classify" # choices=['detect', 'segment', 'classify', 'init'] # init is a special case
|
||||
mode: "train" # choice=['train', 'val', 'infer']
|
||||
|
||||
# Train settings -------------------------------------------------------------------------------------------------------
|
||||
|
Loading…
x
Reference in New Issue
Block a user