🐛 [Fix] quite mode bugs, enable to disable :)

tests/conftest.py

@@ -67,7 +67,8 @@ def model_v7(inference_v7_cfg: Config, device) -> YOLO:
 
 @pytest.fixture(scope="session")
 def solver(train_cfg: Config) -> Trainer:
-    callbacks, loggers = setup(train_cfg)
+    train_cfg.use_wandb = False
+    callbacks, loggers, save_path = setup(train_cfg)
     trainer = Trainer(
         accelerator="cuda",
         max_epochs=getattr(train_cfg.task, "epoch", None),
@@ -77,6 +78,7 @@ def solver(train_cfg: Config) -> Trainer:
         log_every_n_steps=1,
         gradient_clip_val=10,
         deterministic=True,
+        default_root_dir=save_path,
     )
     return trainer
 

yolo/lazy.py

@@ -14,7 +14,7 @@ from yolo.utils.logging_utils import setup
 
 @hydra.main(config_path="config", config_name="config", version_base=None)
 def main(cfg: Config):
-    callbacks, loggers = setup(cfg)
+    callbacks, loggers, save_path = setup(cfg)
 
     trainer = Trainer(
         accelerator="cuda",
@@ -25,6 +25,8 @@ def main(cfg: Config):
         log_every_n_steps=1,
         gradient_clip_val=10,
         deterministic=True,
+        enable_progress_bar=not getattr(cfg, "quite", False),
+        default_root_dir=save_path,
     )
 
     match cfg.task.task:

yolo/tools/solver.py

@@ -149,6 +149,6 @@ class InferenceModel(BaseModel):
         return fps
 
     def _save_image(self, img, batch_idx):
-        save_image_path = Path(self.logger.save_dir) / f"frame{batch_idx:03d}.png"
+        save_image_path = Path(self.trainer.default_root_dir) / f"frame{batch_idx:03d}.png"
         img.save(save_image_path)
         print(f"💾 Saved visualize image at {save_image_path}")

yolo/utils/logging_utils.py

@@ -223,7 +223,7 @@ class ImageLogger(Callback):
                 logger.log_image("Prediction", images, step=step, boxes=[log_bbox(pred_boxes)])
 
 
-def setup_logger(logger_name):
+def setup_logger(logger_name, quite=False):
     class EmojiFormatter(logging.Formatter):
         def format(self, record, emoji=":high_voltage:"):
             return f"{emoji} {super().format(record)}"
@@ -234,16 +234,14 @@ def setup_logger(logger_name):
     if rich_logger:
         rich_logger.handlers.clear()
         rich_logger.addHandler(rich_handler)
+        if quite:
+            rich_logger.setLevel(logging.ERROR)
 
 
 def setup(cfg: Config):
-    # seed_everything(cfg.lucky_number)
-    if hasattr(cfg, "quite"):
-        logger.removeHandler("YOLO_logger")
-        return
-
-    setup_logger("lightning.fabric")
-    setup_logger("lightning.pytorch")
+    quite = hasattr(cfg, "quite")
+    setup_logger("lightning.fabric", quite=quite)
+    setup_logger("lightning.pytorch", quite=quite)
 
     def custom_wandb_log(string="", level=int, newline=True, repeat=True, prefix=True, silent=False):
         if silent:
@@ -256,6 +254,11 @@ def setup(cfg: Config):
     save_path = validate_log_directory(cfg, cfg.name)
 
     progress, loggers = [], []
+
+    if quite:
+        logger.setLevel(logging.ERROR)
+        return progress, loggers, save_path
+
     progress.append(YOLORichProgressBar())
     progress.append(YOLORichModelSummary())
     progress.append(ImageLogger())
@@ -264,7 +267,7 @@ def setup(cfg: Config):
     if cfg.use_wandb:
         loggers.append(WandbLogger(project="YOLO", name=cfg.name, save_dir=save_path, id=None))
 
-    return progress, loggers
+    return progress, loggers, save_path
 
 
 def log_model_structure(model: Union[ModuleList, YOLOLayer, YOLO]):
@@ -312,7 +315,8 @@ def validate_log_directory(cfg: Config, exp_name: str) -> Path:
             )
 
     save_path.mkdir(parents=True, exist_ok=True)
-    logger.info(f"📄 Created log folder: [blue b u]{save_path}[/]")
+    if not getattr(cfg, "quite", False):
+        logger.info(f"📄 Created log folder: [blue b u]{save_path}[/]")
     logger.addHandler(FileHandler(save_path / "output.log"))
     return save_path