File size: 850 Bytes
b80dc1e 642e32f b80dc1e 642e32f b80dc1e 642e32f b80dc1e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
import argparse
from loguru import logger
from model.yolo import get_model
from utils.tools import load_model_cfg, custom_logger
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", help="Name or path to the model configuration file."
)
return parser.parse_args()
if __name__ == "__main__":
custom_logger()
args = parse_arguments()
model_cfg = load_model_cfg(args.model_config)
model = get_model(model_cfg)
logger.info("Success load model: {}", model)
|