henry000 commited on
Commit
642e32f
Β·
1 Parent(s): 76821ac

πŸ§‘β€πŸ’» [add] custom logger, new format of logging

Browse files
Files changed (2) hide show
  1. train.py +3 -3
  2. 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.yaml", help="Name or path to the model configuration file."
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
+ )