π§βπ» [add] custom logger, new format of logging
Browse files- train.py +3 -3
- utils/tools.py +9 -4
train.py
CHANGED
@@ -1,8 +1,7 @@
|
|
1 |
import argparse
|
2 |
-
import sys
|
3 |
from loguru import logger
|
4 |
from model.yolo import get_model
|
5 |
-
from utils.tools import load_model_cfg
|
6 |
|
7 |
|
8 |
def parse_arguments() -> argparse.Namespace:
|
@@ -14,12 +13,13 @@ def parse_arguments() -> argparse.Namespace:
|
|
14 |
"""
|
15 |
parser = argparse.ArgumentParser(description="Load a YOLO model configuration and display the model.")
|
16 |
parser.add_argument(
|
17 |
-
"--model-config", type=str, default="v7-base
|
18 |
)
|
19 |
return parser.parse_args()
|
20 |
|
21 |
|
22 |
if __name__ == "__main__":
|
|
|
23 |
args = parse_arguments()
|
24 |
model_cfg = load_model_cfg(args.model_config)
|
25 |
model = get_model(model_cfg)
|
|
|
1 |
import argparse
|
|
|
2 |
from loguru import logger
|
3 |
from model.yolo import get_model
|
4 |
+
from utils.tools import load_model_cfg, custom_logger
|
5 |
|
6 |
|
7 |
def parse_arguments() -> argparse.Namespace:
|
|
|
13 |
"""
|
14 |
parser = argparse.ArgumentParser(description="Load a YOLO model configuration and display the model.")
|
15 |
parser.add_argument(
|
16 |
+
"--model-config", type=str, default="v7-base", help="Name or path to the model configuration file."
|
17 |
)
|
18 |
return parser.parse_args()
|
19 |
|
20 |
|
21 |
if __name__ == "__main__":
|
22 |
+
custom_logger()
|
23 |
args = parse_arguments()
|
24 |
model_cfg = load_model_cfg(args.model_config)
|
25 |
model = get_model(model_cfg)
|
utils/tools.py
CHANGED
@@ -1,4 +1,5 @@
|
|
1 |
import os
|
|
|
2 |
import yaml
|
3 |
from loguru import logger
|
4 |
from typing import Dict, Any
|
@@ -49,10 +50,6 @@ def load_model_cfg(file_path: str) -> Dict[str, Any]:
|
|
49 |
model_cfg["nc"] = 80
|
50 |
logger.warning("'nc' not found in the YAML file. Setting default 'nc' to 80.")
|
51 |
|
52 |
-
if "anchor" not in model_cfg:
|
53 |
-
logger.error("'anchor' is missing in the configuration file.")
|
54 |
-
raise ValueError("Missing required key: 'anchor'")
|
55 |
-
|
56 |
if "model" not in model_cfg:
|
57 |
logger.error("'model' is missing in the configuration file.")
|
58 |
raise ValueError("Missing required key: 'model'")
|
@@ -65,3 +62,11 @@ def load_model_cfg(file_path: str) -> Dict[str, Any]:
|
|
65 |
except yaml.YAMLError as e:
|
66 |
logger.error(f"Error parsing YAML file: {e}")
|
67 |
raise
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
import os
|
2 |
+
import sys
|
3 |
import yaml
|
4 |
from loguru import logger
|
5 |
from typing import Dict, Any
|
|
|
50 |
model_cfg["nc"] = 80
|
51 |
logger.warning("'nc' not found in the YAML file. Setting default 'nc' to 80.")
|
52 |
|
|
|
|
|
|
|
|
|
53 |
if "model" not in model_cfg:
|
54 |
logger.error("'model' is missing in the configuration file.")
|
55 |
raise ValueError("Missing required key: 'model'")
|
|
|
62 |
except yaml.YAMLError as e:
|
63 |
logger.error(f"Error parsing YAML file: {e}")
|
64 |
raise
|
65 |
+
|
66 |
+
|
67 |
+
def custom_logger():
|
68 |
+
logger.remove()
|
69 |
+
logger.add(
|
70 |
+
sys.stderr,
|
71 |
+
format="<green>{time:MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
|
72 |
+
)
|