HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Start export implementation (#110)

Browse source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2022-12-29 14:17:14 +01:00 • committed by GitHub
parent c1b38428bc
commit 92dad1c1b5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
32 changed files with 827 additions and 222 deletions
Download patch file Download diff file Expand all files Collapse all files

35
ultralytics/yolo/cli.py
View file

@ -1,43 +1,48 @@
import os
import shutil
from pathlib import Path
import hydra
import ultralytics
import ultralytics.yolo.v8 as yolo
from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG
from ultralytics import yolo
from .utils import LOGGER, colorstr
from .utils import DEFAULT_CONFIG, LOGGER, colorstr
@hydra.main(version_base=None, config_path="utils/configs", config_name="default")
@hydra.main(version_base=None, config_path="configs", config_name="default")
def cli(cfg):
cwd = Path().cwd()
LOGGER.info(f"{colorstr(f'Ultralytics YOLO v{ultralytics.__version__}')}")
task, mode = cfg.task.lower(), cfg.mode.lower()
if task == "init": # special case
shutil.copy2(DEFAULT_CONFIG, os.getcwd())
shutil.copy2(DEFAULT_CONFIG, cwd)
LOGGER.info(f"""
{colorstr("YOLO :")} configuration saved to {os.getcwd()}/{DEFAULT_CONFIG.name}.
{colorstr("YOLO:")} configuration saved to {cwd / DEFAULT_CONFIG.name}.
To run experiments using custom configuration:
yolo task='task' mode='mode' --config-name config_file.yaml
""")
return
elif task == "detect":
module_file = yolo.detect
module = yolo.v8.detect
elif task == "segment":
module_file = yolo.segment
module = yolo.v8.segment
elif task == "classify":
module_file = yolo.classify
module = yolo.v8.classify
elif task == "export":
func = yolo.trainer.exporter.export_model
else:
raise SyntaxError("task not recognized. Choices are `'detect', 'segment', 'classify'`")
if mode == "train":
module_function = module_file.train
func = module.train
elif mode == "val":
module_function = module_file.val
func = module.val
elif mode == "predict":
module_function = module_file.predict
func = module.predict
elif mode == "export":
func = yolo.trainer.exporter.export_model
else:
raise SyntaxError("mode not recognized. Choices are `'train', 'val', 'predict'`")
module_function(cfg)
raise SyntaxError("mode not recognized. Choices are `'train', 'val', 'predict', 'export'`")
func(cfg)