YOLO / train.py
henry000's picture
🔨 [Init] the base train code, and load yaml cfg
b80dc1e
raw
history blame
831 Bytes
import argparse
import sys
from loguru import logger
from model.yolo import get_model
from utils.tools import load_model_cfg
def parse_arguments() -> argparse.Namespace:
"""
Parse command-line arguments to get the model configuration file.
Returns:
argparse.Namespace: The command-line arguments object with 'config' attribute.
"""
parser = argparse.ArgumentParser(description="Load a YOLO model configuration and display the model.")
parser.add_argument(
"--model-config", type=str, default="v7-base.yaml", help="Name or path to the model configuration file."
)
return parser.parse_args()
if __name__ == "__main__":
args = parse_arguments()
model_cfg = load_model_cfg(args.model_config)
model = get_model(model_cfg)
logger.info("Success load model: {}", model)