🔀 [Merge] branch 'SETUP' into MODELv2

yolo/lazy.py

@@ -13,13 +13,12 @@ from yolo.tools.data_loader import create_dataloader
 from yolo.tools.solver import ModelTester, ModelTrainer
 from yolo.utils.bounding_box_utils import Vec2Box
 from yolo.utils.deploy_utils import FastModelLoader
-from yolo.utils.logging_utils import custom_logger, validate_log_directory
+from yolo.utils.logging_utils import ProgressLogger
 
 
 @hydra.main(config_path="config", config_name="config", version_base=None)
 def main(cfg: Config):
-    custom_logger()
-    save_path = validate_log_directory(cfg, exp_name=cfg.name)
+    progress = ProgressLogger(cfg, exp_name=cfg.name)
     dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task)
     device = torch.device(cfg.device)
     if getattr(cfg.task, "fast_inference", False):
@@ -31,11 +30,11 @@ def main(cfg: Config):
     vec2box = Vec2Box(model, cfg.image_size, device)
 
     if cfg.task.task == "train":
-        trainer = ModelTrainer(cfg, model, vec2box, save_path, device)
+        trainer = ModelTrainer(cfg, model, vec2box, progress, device)
         trainer.solve(dataloader)
 
     if cfg.task.task == "inference":
-        tester = ModelTester(cfg, model, vec2box, save_path, device)
+        tester = ModelTester(cfg, model, vec2box, progress, device)
         tester.solve(dataloader)
 
 

yolo/tools/solver.py

@@ -14,7 +14,7 @@ from yolo.tools.data_loader import StreamDataLoader, create_dataloader
 from yolo.tools.drawer import draw_bboxes
 from yolo.tools.loss_functions import create_loss_function
 from yolo.utils.bounding_box_utils import Vec2Box, bbox_nms, calculate_map
-from yolo.utils.logging_utils import ProgressTracker
+from yolo.utils.logging_utils import ProgressLogger
 from yolo.utils.model_utils import (
     ExponentialMovingAverage,
     create_optimizer,
@@ -23,7 +23,7 @@ from yolo.utils.model_utils import (
 
 
 class ModelTrainer:
-    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, save_path: str, device):
+    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: ProgressLogger, device):
         train_cfg: TrainConfig = cfg.task
         self.model = model
         self.vec2box = vec2box
@@ -31,11 +31,11 @@ class ModelTrainer:
         self.optimizer = create_optimizer(model, train_cfg.optimizer)
         self.scheduler = create_scheduler(self.optimizer, train_cfg.scheduler)
         self.loss_fn = create_loss_function(cfg, vec2box)
-        self.progress = ProgressTracker(cfg.name, save_path, cfg.use_wandb)
+        self.progress = progress
         self.num_epochs = cfg.task.epoch
 
         self.validation_dataloader = create_dataloader(cfg.task.validation.data, cfg.dataset, cfg.task.validation.task)
-        self.validator = ModelValidator(cfg.task.validation, model, vec2box, save_path, device, self.progress)
+        self.validator = ModelValidator(cfg.task.validation, model, vec2box, progress, device, self.progress)
 
         if getattr(train_cfg.ema, "enabled", False):
             self.ema = ExponentialMovingAverage(model, decay=train_cfg.ema.decay)
@@ -102,14 +102,15 @@ class ModelTrainer:
 
 
 class ModelTester:
-    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, save_path: str, device):
+    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: ProgressLogger, device):
         self.model = model
         self.device = device
         self.vec2box = vec2box
-        self.progress = ProgressTracker(cfg, save_path, cfg.use_wandb)
+        self.progress = progress
 
         self.nms = cfg.task.nms
-        self.save_path = save_path
+        self.save_path = os.path.join(progress.save_path, "images")
+        os.makedirs(self.save_path, exist_ok=True)
         self.save_predict = getattr(cfg.task, "save_predict", None)
         self.idx2label = cfg.class_list
 
@@ -164,16 +165,13 @@ class ModelValidator:
         validation_cfg: ValidationConfig,
         model: YOLO,
         vec2box: Vec2Box,
-        save_path: str,
         device,
-        # TODO: think Progress?
-        progress: ProgressTracker,
+        progress: ProgressLogger,
     ):
         self.model = model
         self.vec2box = vec2box
         self.device = device
         self.progress = progress
-        self.save_path = save_path
 
         self.nms = validation_cfg.nms
 

yolo/utils/logging_utils.py

@@ -38,15 +38,18 @@ def custom_logger(quite: bool = False):
     )
 
 
-class ProgressTracker:
-    def __init__(self, exp_name: str, save_path: str, use_wandb: bool = False):
+class ProgressLogger:
+    def __init__(self, cfg: Config, exp_name: str):
+        custom_logger(getattr(cfg, "quite", False))
+        self.save_path = validate_log_directory(cfg, exp_name=cfg.name)
+
         self.progress = Progress(
             TextColumn("[progress.description]{task.description}"),
             BarColumn(bar_width=None),
             TextColumn("{task.completed:.0f}/{task.total:.0f}"),
             TimeRemainingColumn(),
         )
-        self.use_wandb = use_wandb
+        self.use_wandb = cfg.use_wandb
         if self.use_wandb:
             wandb.errors.term._log = custom_wandb_log
             self.wandb = wandb.init(