From d0b0fe2592225ccc71132c47d61d16435105d644 Mon Sep 17 00:00:00 2001
From: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Tue, 29 Nov 2022 15:41:37 +0530
Subject: [PATCH] CLI updates (#58)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 setup.py                                    |  2 +-
 ultralytics/yolo/__init__.py                | 36 +---------------
 ultralytics/yolo/cli.py                     | 47 +++++++++++++++++++++
 ultralytics/yolo/engine/model.py            |  1 -
 ultralytics/yolo/utils/__init__.py          |  2 +-
 ultralytics/yolo/utils/configs/default.yaml |  2 +-
 6 files changed, 51 insertions(+), 39 deletions(-)
 create mode 100644 ultralytics/yolo/cli.py

diff --git a/setup.py b/setup.py
index 7085012..cd1c0db 100644
--- a/setup.py
+++ b/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',],})
diff --git a/ultralytics/yolo/__init__.py b/ultralytics/yolo/__init__.py
index d100fb9..886a846 100644
--- a/ultralytics/yolo/__init__.py
+++ b/ultralytics/yolo/__init__.py
@@ -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)
diff --git a/ultralytics/yolo/cli.py b/ultralytics/yolo/cli.py
new file mode 100644
index 0000000..49c0c26
--- /dev/null
+++ b/ultralytics/yolo/cli.py
@@ -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)
diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py
index d5e5ca4..3a70a2d 100644
--- a/ultralytics/yolo/engine/model.py
+++ b/ultralytics/yolo/engine/model.py
@@ -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
diff --git a/ultralytics/yolo/utils/__init__.py b/ultralytics/yolo/utils/__init__.py
index 21d7d06..0216ec3 100644
--- a/ultralytics/yolo/utils/__init__.py
+++ b/ultralytics/yolo/utils/__init__.py
@@ -1,5 +1,5 @@
 import contextlib
-import logging
+import logging.config
 import os
 import platform
 import sys
diff --git a/ultralytics/yolo/utils/configs/default.yaml b/ultralytics/yolo/utils/configs/default.yaml
index fdb5d39..99c1c10 100644
--- a/ultralytics/yolo/utils/configs/default.yaml
+++ b/ultralytics/yolo/utils/configs/default.yaml
@@ -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 -------------------------------------------------------------------------------------------------------