💬 [Update] logging -> loguru, consistency logging

yolo/model/module.py

@@ -3,10 +3,10 @@ from typing import Any, Dict, List, Optional, Tuple
 import torch
 import torch.nn.functional as F
 from einops import rearrange
-from loguru import logger
 from torch import Tensor, nn
 from torch.nn.common_types import _size_2_t
 
+from yolo.utils.logger import logger
 from yolo.utils.module_utils import auto_pad, create_activation_function, round_up
 
 

yolo/model/yolo.py

@@ -3,12 +3,12 @@ from pathlib import Path
 from typing import Dict, List, Union
 
 import torch
-from loguru import logger
 from omegaconf import ListConfig, OmegaConf
 from torch import nn
 
 from yolo.config.config import ModelConfig, YOLOLayer
 from yolo.tools.dataset_preparation import prepare_weight
+from yolo.utils.logger import logger
 from yolo.utils.module_utils import get_layer_map
 
 

yolo/tools/data_loader.py

@@ -5,7 +5,6 @@ from typing import Generator, List, Tuple, Union
 
 import numpy as np
 import torch
-from loguru import logger
 from PIL import Image
 from rich.progress import track
 from torch import Tensor
@@ -21,6 +20,7 @@ from yolo.utils.dataset_utils import (
     locate_label_paths,
     scale_segmentation,
 )
+from yolo.utils.logger import logger
 
 
 class YoloDataset(Dataset):
@@ -48,12 +48,12 @@ class YoloDataset(Dataset):
         cache_path = dataset_path / f"{phase_name}.cache"
 
         if not cache_path.exists():
-            logger.info("🏭 Generating {} cache", phase_name)
+            logger.info(f"🏭 Generating {phase_name} cache")
             data = self.filter_data(dataset_path, phase_name)
             torch.save(data, cache_path)
         else:
             data = torch.load(cache_path, weights_only=False)
-            logger.info("📦 Loaded {} cache", phase_name)
+            logger.info(f"📦 Loaded {phase_name} cache")
         return data
 
     def filter_data(self, dataset_path: Path, phase_name: str) -> list:

yolo/tools/dataset_preparation.py

@@ -3,10 +3,10 @@ from pathlib import Path
 from typing import Optional
 
 import requests
-from loguru import logger
 from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
 
 from yolo.config.config import DatasetConfig
+from yolo.utils.logger import logger
 
 
 def download_file(url, destination: Path):

yolo/tools/drawer.py

@@ -3,12 +3,12 @@ from typing import List, Optional, Union
 
 import numpy as np
 import torch
-from loguru import logger
 from PIL import Image, ImageDraw, ImageFont
 from torchvision.transforms.functional import to_pil_image
 
 from yolo.config.config import ModelConfig
 from yolo.model.yolo import YOLO
+from yolo.utils.logger import logger
 
 
 def draw_bboxes(

yolo/tools/loss_functions.py

@@ -2,12 +2,12 @@ from typing import Any, Dict, List, Tuple
 
 import torch
 import torch.nn.functional as F
-from loguru import logger
 from torch import Tensor, nn
 from torch.nn import BCEWithLogitsLoss
 
 from yolo.config.config import Config, LossConfig
 from yolo.utils.bounding_box_utils import BoxMatcher, Vec2Box, calculate_iou
+from yolo.utils.logger import logger
 
 
 class BCELoss(nn.Module):

yolo/tools/solver.py

@@ -8,7 +8,6 @@ from pathlib import Path
 from typing import Dict, Optional
 
 import torch
-from loguru import logger
 from pycocotools.coco import COCO
 from torch import Tensor, distributed
 from torch.cuda.amp import GradScaler, autocast
@@ -22,6 +21,7 @@ from yolo.tools.drawer import draw_bboxes, draw_model
 from yolo.tools.loss_functions import create_loss_function
 from yolo.utils.bounding_box_utils import Vec2Box, calculate_map
 from yolo.utils.dataset_utils import locate_label_paths
+from yolo.utils.logger import logger
 from yolo.utils.logging_utils import ProgressLogger, log_model_structure
 from yolo.utils.model_utils import (
     ExponentialMovingAverage,

yolo/utils/bounding_box_utils.py

@@ -4,12 +4,12 @@ from typing import Dict, List, Optional, Tuple, Union
 import torch
 import torch.nn.functional as F
 from einops import rearrange
-from loguru import logger
 from torch import Tensor, arange, tensor
 from torchvision.ops import batched_nms
 
 from yolo.config.config import AnchorConfig, MatcherConfig, ModelConfig, NMSConfig
 from yolo.model.yolo import YOLO
+from yolo.utils.logger import logger
 
 
 def calculate_iou(bbox1, bbox2, metrics="iou") -> Tensor:

yolo/utils/dataset_utils.py

@@ -5,9 +5,9 @@ from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
-from loguru import logger
 
 from yolo.tools.data_conversion import discretize_categories
+from yolo.utils.logger import logger
 
 
 def locate_label_paths(dataset_path: Path, phase_name: Path) -> Tuple[Path, Path]:

yolo/utils/deploy_utils.py

@@ -1,11 +1,11 @@
 from pathlib import Path
 
 import torch
-from loguru import logger
 from torch import Tensor
 
 from yolo.config.config import Config
 from yolo.model.yolo import create_model
+from yolo.utils.logger import logger
 
 
 class FastModelLoader:

yolo/utils/logger.py

@@ -0,0 +1,10 @@
+import logging
+
+from rich.console import Console
+from rich.logging import RichHandler
+
+logger = logging.getLogger("YOLO_logger")
+logger.setLevel(logging.DEBUG)
+logger.propagate = False
+if not logger.hasHandlers():
+    logger.addHandler(RichHandler(console=Console(), show_level=True, show_path=True, show_time=True))

yolo/utils/logging_utils.py

@@ -15,6 +15,7 @@ import os
 import random
 import sys
 from collections import deque
+from logging import FileHandler
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -22,7 +23,6 @@ import numpy as np
 import torch
 import wandb
 import wandb.errors.term
-from loguru import logger
 from omegaconf import ListConfig
 from rich.console import Console, Group
 from rich.progress import (
@@ -41,18 +41,13 @@ from torchvision.transforms.functional import pil_to_tensor
 from yolo.config.config import Config, YOLOLayer
 from yolo.model.yolo import YOLO
 from yolo.tools.drawer import draw_bboxes
+from yolo.utils.logger import logger
 from yolo.utils.solver_utils import make_ap_table
 
 
 def custom_logger(quite: bool = False):
-    logger.remove()
     if quite:
-        return
-    logger.add(
-        sys.stderr,
-        colorize=True,
-        format="<fg #003385>[{time:MM/DD HH:mm:ss}]</> <level>{level: ^8}</level>| <level>{message}</level>",
-    )
+        logger.removeHandler("YOLO_logger")
 
 
 # TODO: should be moved to correct position
@@ -100,7 +95,7 @@ class ProgressLogger(Progress):
             from torch.utils.tensorboard import SummaryWriter
 
             self.tb_writer = SummaryWriter(log_dir=self.save_path / "tensorboard")
-            logger.opt(colors=True).info(f"📍 Enable TensorBoard locally at <blue><u>http://localhost:6006</></>")
+            logger.info(f"📍 Enable TensorBoard locally at <blue><u>http://localhost:6006</></>")
 
     def rank_check(logging_function):
         def wrapper(self, *args, **kwargs):
@@ -249,7 +244,7 @@ def custom_wandb_log(string="", level=int, newline=True, repeat=True, prefix=Tru
     if silent:
         return
     for line in string.split("\n"):
-        logger.opt(raw=not newline, colors=True).info("🌐 " + line)
+        logger.info("🌐 " + line)
 
 
 def log_model_structure(model: Union[ModuleList, YOLOLayer, YOLO]):
@@ -296,8 +291,8 @@ def validate_log_directory(cfg: Config, exp_name: str) -> Path:
             )
 
     save_path.mkdir(parents=True, exist_ok=True)
-    logger.opt(colors=True).info(f"📄 Created log folder: <u><fg #808080>{save_path}</></>")
-    logger.add(save_path / "output.log", mode="w", backtrace=True, diagnose=True)
+    logger.info(f"📄 Created log folder: [bold gray]{save_path}[/]", extra={"markup": True})
+    logger.addHandler(FileHandler(save_path / "output.log"))
     return save_path
 
 

yolo/utils/model_utils.py

@@ -4,7 +4,6 @@ from typing import List, Optional, Type, Union
 
 import torch
 import torch.distributed as dist
-from loguru import logger
 from omegaconf import ListConfig
 from torch import Tensor
 from torch.optim import Optimizer
@@ -13,6 +12,7 @@ from torch.optim.lr_scheduler import LambdaLR, SequentialLR, _LRScheduler
 from yolo.config.config import IDX_TO_ID, NMSConfig, OptimizerConfig, SchedulerConfig
 from yolo.model.yolo import YOLO
 from yolo.utils.bounding_box_utils import bbox_nms, transform_bbox
+from yolo.utils.logger import logger
 
 
 class ExponentialMovingAverage: