🔀 [Merge] branch 'SETUP' into TEST

yolo/config/dataset/dev.yaml

@@ -1,5 +1,5 @@
 path: data/dev
 train: train
-validation: test
+validation: val
 
 auto_download:

yolo/config/general.yaml

@@ -9,7 +9,7 @@ out_path: runs
 exist_ok: True
 
 lucky_number: 10
-use_wandb: False
+use_wandb: True
 use_TensorBoard: False
 
 weight: True # Path to weight or True for auto, False for no pretrained weight

yolo/config/task/validation.yaml

@@ -8,5 +8,5 @@ data:
   pin_memory: True
   data_augment: {}
 nms:
-  min_confidence: 0.001
-  min_iou: 0.7
+  min_confidence: 0.05
+  min_iou: 0.9

yolo/lazy.py

@@ -28,18 +28,14 @@ def main(cfg: Config):
         model = model.to(device)
 
     vec2box = Vec2Box(model, cfg.image_size, device)
-
     if cfg.task.task == "train":
-        trainer = ModelTrainer(cfg, model, vec2box, progress, device, use_ddp)
-        trainer.solve(dataloader)
-
-    if cfg.task.task == "inference":
-        tester = ModelTester(cfg, model, vec2box, progress, device)
-        tester.solve(dataloader)
-
+        solver = ModelTrainer(cfg, model, vec2box, progress, device, use_ddp)
     if cfg.task.task == "validation":
-        valider = ModelValidator(cfg.task, model, vec2box, progress, device)
-        valider.solve(dataloader)
+        solver = ModelValidator(cfg.task, model, vec2box, progress, device)
+    if cfg.task.task == "inference":
+        solver = ModelTester(cfg, model, vec2box, progress, device)
+    progress.start()
+    solver.solve(dataloader)
 
 
 if __name__ == "__main__":

yolo/model/yolo.py

@@ -82,17 +82,14 @@ class YOLO(nn.Module):
         return output
 
     def get_out_channels(self, layer_type: str, layer_args: dict, output_dim: list, source: Union[int, list]):
-        # TODO refactor, check out_channels in layer_args, next check source is list, CBFuse|Concat
-        if any(module in layer_type for module in ["Conv", "ELAN", "ADown", "AConv"]):
+        if hasattr(layer_args, "out_channels"):
             return layer_args["out_channels"]
         if layer_type == "CBFuse":
             return output_dim[source[-1]]
-        if layer_type in ["Pool", "UpSample"]:
+        if isinstance(source, int):
             return output_dim[source]
-        if layer_type == "Concat":
+        if isinstance(source, list):
             return sum(output_dim[idx] for idx in source)
-        if layer_type == "IDetect":
-            return None
 
     def get_source_idx(self, source: Union[ListConfig, str, int], layer_idx: int):
         if isinstance(source, ListConfig):
@@ -128,7 +125,6 @@ def create_model(model_cfg: ModelConfig, weight_path: Union[bool, str] = True, c
     Returns:
         YOLO: An instance of the model defined by the given configuration.
     """
-    # TODO: "weight_path -> weight = [True|None-False|Path]: True should be default of model name?"
     OmegaConf.set_struct(model_cfg, False)
     model = YOLO(model_cfg, class_num)
     if weight_path:
@@ -138,7 +134,6 @@ def create_model(model_cfg: ModelConfig, weight_path: Union[bool, str] = True, c
             logger.info(f"🌐 Weight {weight_path} not found, try downloading")
             prepare_weight(weight_path=weight_path)
         if os.path.exists(weight_path):
-            # TODO: fix map_location
             model.model.load_state_dict(torch.load(weight_path, map_location=torch.device("cpu")), strict=False)
             logger.info("✅ Success load model & weight")
     else:

yolo/tools/data_augmentation.py

@@ -46,12 +46,8 @@ class PadAndResize:
         padded_image = Image.new("RGB", (self.target_width, self.target_height), self.background_color)
         padded_image.paste(resized_image, (pad_left, pad_top))
 
-        boxes[:, 1] *= scale  # xmin
-        boxes[:, 2] *= scale  # ymin
-        boxes[:, 3] *= scale  # xmax
-        boxes[:, 4] *= scale  # ymax
-        boxes[:, [1, 3]] += pad_left
-        boxes[:, [2, 4]] += pad_top
+        boxes[:, [1, 3]] = (boxes[:, [1, 3]] * new_width + pad_left) / self.target_width
+        boxes[:, [2, 4]] = (boxes[:, [2, 4]] * new_height + pad_top) / self.target_height
 
         transform_info = torch.tensor([scale, pad_left, pad_top, pad_left, pad_top])
         return padded_image, boxes, transform_info

yolo/tools/data_loader.py

@@ -199,7 +199,7 @@ class YoloDataLoader(DataLoader):
         batch_images = torch.stack(batch_images)
         batch_reverse = torch.stack(batch_reverse)
 
-        return batch_images, batch_targets, batch_reverse, batch_path
+        return batch_size, batch_images, batch_targets, batch_reverse, batch_path
 
 
 def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train", use_ddp: bool = False):

yolo/tools/solver.py

@@ -1,9 +1,12 @@
 import json
 import os
+import sys
 import time
+from collections import defaultdict
 
 import torch
 from loguru import logger
+from pycocotools.coco import COCO
 from torch import Tensor
 
 # TODO: We may can't use CUDA?
@@ -16,7 +19,7 @@ from yolo.model.yolo import YOLO
 from yolo.tools.data_loader import StreamDataLoader, create_dataloader
 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
+from yolo.utils.bounding_box_utils import Vec2Box, calculate_map
 from yolo.utils.logging_utils import ProgressLogger, log_model_structure
 from yolo.utils.model_utils import (
     ExponentialMovingAverage,
@@ -25,6 +28,7 @@ from yolo.utils.model_utils import (
     create_scheduler,
     predicts_to_json,
 )
+from yolo.utils.solver_utils import calculate_ap
 
 
 class ModelTrainer:
@@ -69,22 +73,28 @@ class ModelTrainer:
         self.scaler.step(self.optimizer)
         self.scaler.update()
 
-        return loss.item(), loss_item
+        return loss_item
 
     def train_one_epoch(self, dataloader):
         self.model.train()
-        total_loss = 0
+        total_loss = defaultdict(lambda: torch.tensor(0.0, device=self.device))
+        total_samples = 0
 
-        for images, targets, *_ in dataloader:
-            loss, loss_each = self.train_one_batch(images, targets)
+        for batch_size, images, targets, *_ in dataloader:
+            loss_each = self.train_one_batch(images, targets)
 
-            total_loss += loss
+            for loss_name, loss_val in loss_each.items():
+                total_loss[loss_name] += loss_val * batch_size
+            total_samples += batch_size
             self.progress.one_batch(loss_each)
 
+        for loss_val in total_loss.values():
+            loss_val /= total_samples
+
         if self.scheduler:
             self.scheduler.step()
 
-        return total_loss / len(dataloader)
+        return total_loss
 
     def save_checkpoint(self, epoch: int, filename="checkpoint.pt"):
         checkpoint = {
@@ -107,12 +117,11 @@ class ModelTrainer:
             if self.use_ddp:
                 dataloader.sampler.set_epoch(epoch)
 
-            self.progress.start_one_epoch(len(dataloader), self.optimizer, epoch)
-            # TODO: calculate epoch loss
+            self.progress.start_one_epoch(len(dataloader), "Train", self.optimizer, epoch)
             epoch_loss = self.train_one_epoch(dataloader)
-            self.progress.finish_one_epoch()
+            self.progress.finish_one_epoch(epoch_loss, epoch)
 
-            self.validator.solve(self.validation_dataloader)
+            self.validator.solve(self.validation_dataloader, epoch_idx=epoch)
 
 
 class ModelTester:
@@ -187,19 +196,30 @@ class ModelValidator:
         self.post_proccess = PostProccess(vec2box, validation_cfg.nms)
         self.json_path = os.path.join(self.progress.save_path, f"predict.json")
 
-    def solve(self, dataloader):
+        sys.stdout = open(os.devnull, "w")
+        # TODO: load with config file
+        self.coco_gt = COCO("data/coco/annotations/instances_val2017.json")
+        sys.stdout = sys.__stdout__
+
+    def solve(self, dataloader, epoch_idx=-1):
         # logger.info("🧪 Start Validation!")
         self.model.eval()
-        predict_json = []
-        self.progress.start_one_epoch(len(dataloader))
-        for images, targets, rev_tensor, img_paths in dataloader:
+        mAPs, predict_json = [], []
+        self.progress.start_one_epoch(len(dataloader), task="Validate")
+        for batch_size, images, targets, rev_tensor, img_paths in dataloader:
             images, targets, rev_tensor = images.to(self.device), targets.to(self.device), rev_tensor.to(self.device)
             with torch.no_grad():
                 predicts = self.model(images)
-                predicts = self.post_proccess(predicts, rev_tensor)
-            self.progress.one_batch()
+                predicts = self.post_proccess(predicts)
+                for idx, predict in enumerate(predicts):
+                    mAPs.append(calculate_map(predict, targets[idx]))
+            self.progress.one_batch(Tensor(mAPs))
 
-            predict_json.extend(predicts_to_json(img_paths, predicts))
-        self.progress.finish_one_epoch()
+            predict_json.extend(predicts_to_json(img_paths, predicts, rev_tensor))
+        self.progress.finish_one_epoch(Tensor(mAPs), epoch_idx=epoch_idx)
         with open(self.json_path, "w") as f:
             json.dump(predict_json, f)
+
+        self.progress.start_pycocotools()
+        result = calculate_ap(self.coco_gt, predict_json)
+        self.progress.finish_pycocotools(result, epoch_idx)

yolo/utils/bounding_box_utils.py

@@ -5,7 +5,7 @@ import torch
 import torch.nn.functional as F
 from einops import rearrange
 from loguru import logger
-from torch import Tensor
+from torch import Tensor, arange
 from torchvision.ops import batched_nms
 
 from yolo.config.config import MatcherConfig, ModelConfig, NMSConfig
@@ -338,8 +338,8 @@ def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig):
     return predicts_nms
 
 
-def calculate_map(predictions, ground_truths, iou_thresholds):
-    # TODO: Refactor this block
+def calculate_map(predictions, ground_truths, iou_thresholds=arange(0.5, 1, 0.05)):
+    # TODO: Refactor this block, Flexible for calculate different mAP condition?
     device = predictions.device
     n_preds = predictions.size(0)
     n_gts = (ground_truths[:, 0] != -1).sum()
@@ -369,13 +369,16 @@ def calculate_map(predictions, ground_truths, iou_thresholds):
         precision = tp_cumsum / (tp_cumsum + fp_cumsum + 1e-6)
         recall = tp_cumsum / (n_gts + 1e-6)
 
-        recall_thresholds = torch.arange(0, 1, 0.1)
-        precision_at_recall = torch.zeros_like(recall_thresholds)
-        for i, r in enumerate(recall_thresholds):
-            precision_at_recall[i] = precision[recall >= r].max().item() if torch.any(recall >= r) else 0
+        precision = torch.cat([torch.ones(1, device=device), precision, torch.zeros(1, device=device)])
+        recall = torch.cat([torch.zeros(1, device=device), recall, torch.ones(1, device=device)])
+
+        precision, _ = torch.cummax(precision.flip(0), dim=0)
+        precision = precision.flip(0)
+
+        indices = (recall[1:] != recall[:-1]).nonzero(as_tuple=True)[0]
+        ap = torch.sum((recall[indices + 1] - recall[indices]) * precision[indices + 1])
 
-        ap = precision_at_recall.mean()
         aps.append(ap)
 
     mean_ap = torch.mean(torch.stack(aps))
-    return mean_ap, aps
+    return mean_ap, aps[0]

yolo/utils/logging_utils.py

@@ -13,18 +13,27 @@ Example:
 
 import os
 import sys
-from typing import Dict, List
+from collections import deque
+from typing import Any, Dict, List
 
 import wandb
 import wandb.errors.term
 from loguru import logger
-from rich.console import Console
-from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
+from omegaconf import ListConfig
+from rich.console import Console, Group
+from rich.progress import (
+    BarColumn,
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeRemainingColumn,
+)
 from rich.table import Table
 from torch import Tensor
 from torch.optim import Optimizer
 
 from yolo.config.config import Config, YOLOLayer
+from yolo.utils.solver_utils import make_ap_table
 
 
 def custom_logger(quite: bool = False):
@@ -38,20 +47,24 @@ def custom_logger(quite: bool = False):
     )
 
 
-class ProgressLogger:
-    def __init__(self, cfg: Config, exp_name: str):
+class ProgressLogger(Progress):
+    def __init__(self, cfg: Config, exp_name: str, *args, **kwargs):
         local_rank = int(os.getenv("LOCAL_RANK", "0"))
         self.quite_mode = local_rank or getattr(cfg, "quite", False)
         custom_logger(self.quite_mode)
         self.save_path = validate_log_directory(cfg, exp_name=cfg.name)
 
-        self.progress = Progress(
+        progress_bar = (
+            SpinnerColumn(),
             TextColumn("[progress.description]{task.description}"),
             BarColumn(bar_width=None),
             TextColumn("{task.completed:.0f}/{task.total:.0f}"),
             TimeRemainingColumn(),
         )
-        self.progress.start()
+        self.ap_table = Table()
+        # TODO: load maxlen by config files
+        self.ap_past_list = deque(maxlen=5)
+        super().__init__(*args, *progress_bar, **kwargs)
 
         self.use_wandb = cfg.use_wandb
         if self.use_wandb:
@@ -60,35 +73,71 @@ class ProgressLogger:
                 project="YOLO", resume="allow", mode="online", dir=self.save_path, id=None, name=exp_name
             )
 
+    def get_renderable(self):
+        renderable = Group(*self.get_renderables(), self.ap_table)
+        return renderable
+
     def start_train(self, num_epochs: int):
-        self.task_epoch = self.progress.add_task("[cyan]Epochs  [white]| Loss | Box  | DFL  | BCE  |", total=num_epochs)
+        self.task_epoch = self.add_task(f"[cyan]Start Training {num_epochs} epochs", total=num_epochs)
 
-    def start_one_epoch(self, num_batches: int, optimizer: Optimizer = None, epoch_idx: int = None):
+    def start_one_epoch(
+        self, num_batches: int, task: str = "Train", optimizer: Optimizer = None, epoch_idx: int = None
+    ):
         self.num_batches = num_batches
-        if self.use_wandb:
+        self.task = task
+        if hasattr(self, "task_epoch"):
+            self.update(self.task_epoch, description=f"[cyan] Preparing Data")
+
+        if self.use_wandb and optimizer is not None:
             lr_values = [params["lr"] for params in optimizer.param_groups]
-            lr_names = ["bias", "norm", "conv"]
+            lr_names = ["Learning Rate/bias", "Learning Rate/norm", "Learning Rate/conv"]
             for lr_name, lr_value in zip(lr_names, lr_values):
-                self.wandb.log({f"Learning Rate/{lr_name}": lr_value}, step=epoch_idx)
-        self.batch_task = self.progress.add_task("[green]Batches", total=num_batches)
-
-    def one_batch(self, loss_dict: Dict[str, Tensor] = None):
-        if loss_dict is None:
-            self.progress.update(self.batch_task, advance=1, description=f"[green]Validating")
-            return
+                self.wandb.log({lr_name: lr_value}, step=epoch_idx)
+        self.batch_task = self.add_task(f"[green] Phase: {task}", total=num_batches)
+
+    def one_batch(self, batch_info: Dict[str, Tensor] = None):
+        epoch_descript = "[cyan]" + self.task + "[white] |"
+        batch_descript = "|"
+        if self.task == "Train":
+            self.update(self.task_epoch, advance=1 / self.num_batches)
+        elif self.task == "Validate":
+            batch_info = {
+                "mAP.5": batch_info.mean(dim=0)[0],
+                "mAP.5:.95": batch_info.mean(dim=0)[1],
+            }
+        for info_name, info_val in batch_info.items():
+            epoch_descript += f"{info_name: ^9}|"
+            batch_descript += f"   {info_val:2.2f}  |"
+        self.update(self.batch_task, advance=1, description=f"[green]{self.task} [white]{batch_descript}")
+        if hasattr(self, "task_epoch"):
+            self.update(self.task_epoch, description=epoch_descript)
+
+    def finish_one_epoch(self, batch_info: Dict[str, Any] = None, epoch_idx: int = -1):
+        if self.task == "Train":
+            for loss_name in batch_info.keys():
+                batch_info["Loss/" + loss_name] = batch_info.pop(loss_name)
+        elif self.task == "Validate":
+            batch_info = {
+                "Metrics/mAP.5": batch_info.mean(dim=0)[0],
+                "Metrics/mAP.5:.95": batch_info.mean(dim=0)[1],
+            }
         if self.use_wandb:
-            for loss_name, loss_value in loss_dict.items():
-                self.wandb.log({f"Loss/{loss_name}": loss_value})
+            self.wandb.log(batch_info, step=epoch_idx)
+        self.remove_task(self.batch_task)
 
-        loss_str = "| -.-- |"
-        for loss_name, loss_val in loss_dict.items():
-            loss_str += f" {loss_val:2.2f} |"
+    def start_pycocotools(self):
+        self.batch_task = self.add_task("[green] run pycocotools", total=1)
 
-        self.progress.update(self.batch_task, advance=1, description=f"[green]Batches [white]{loss_str}")
-        self.progress.update(self.task_epoch, advance=1 / self.num_batches)
+    def finish_pycocotools(self, result, epoch_idx=-1):
+        ap_table, ap_main = make_ap_table(result, self.ap_past_list, epoch_idx)
+        self.ap_past_list.append((epoch_idx, ap_main))
+        self.ap_table = ap_table
 
-    def finish_one_epoch(self):
-        self.progress.remove_task(self.batch_task)
+        if self.use_wandb:
+            self.wandb.log({"PyCOCO/AP @ .5:.95": ap_main[1], "PyCOCO/AP @ .5": ap_main[3]})
+        self.update(self.batch_task, advance=1)
+        self.refresh()
+        self.remove_task(self.batch_task)
 
     def finish_train(self):
         self.wandb.finish()
@@ -115,7 +164,11 @@ def log_model_structure(model: List[YOLOLayer]):
         layer_param = sum(x.numel() for x in layer.parameters())  # number parameters
         in_channels, out_channels = getattr(layer, "in_c", None), getattr(layer, "out_c", None)
         if in_channels and out_channels:
-            channels = f"{in_channels:4} -> {out_channels:4}"
+            if isinstance(in_channels, (list, ListConfig)):
+                in_channels = "M"
+            if isinstance(out_channels, (list, ListConfig)):
+                out_channels = "M"
+            channels = f"{str(in_channels): >4} -> {str(out_channels): >4}"
         else:
             channels = "-"
         table.add_row(str(idx), layer.layer_type, layer.tags, f"{layer_param:,}", channels)

yolo/utils/model_utils.py

@@ -106,7 +106,7 @@ class PostProccess:
         self.vec2box = vec2box
         self.nms = nms_cfg
 
-    def __call__(self, predict, rev_tensor: Optional[Tensor]):
+    def __call__(self, predict, rev_tensor: Optional[Tensor] = None):
         pred_class, _, pred_bbox = self.vec2box(predict["Main"])
         if rev_tensor is not None:
             pred_bbox = (pred_bbox - rev_tensor[:, None, 1:]) / rev_tensor[:, 0:1, None]
@@ -114,13 +114,15 @@ class PostProccess:
         return pred_bbox
 
 
-def predicts_to_json(img_paths, predicts):
+def predicts_to_json(img_paths, predicts, rev_tensor):
     """
     TODO: function document
     turn a batch of imagepath and predicts(n x 6 for each image) to a List of diction(Detection output)
     """
     batch_json = []
-    for img_path, bboxes in zip(img_paths, predicts):
+    for img_path, bboxes, box_reverse in zip(img_paths, predicts, rev_tensor):
+        scale, shift = box_reverse.split([1, 4])
+        bboxes[:, 1:5] = (bboxes[:, 1:5] - shift[None]) / scale[None]
         bboxes[:, 1:5] = transform_bbox(bboxes[:, 1:5], "xyxy -> xywh")
         for cls, *pos, conf in bboxes:
             bbox = {

yolo/utils/solver_utils.py

@@ -0,0 +1,46 @@
+import os
+import sys
+
+from pycocotools.coco import COCO
+from pycocotools.cocoeval import COCOeval
+from rich.table import Table
+
+
+def calculate_ap(coco_gt: COCO, pd_path):
+    sys.stdout = open(os.devnull, "w")
+    coco_dt = coco_gt.loadRes(pd_path)
+    coco_eval = COCOeval(coco_gt, coco_dt, "bbox")
+    coco_eval.evaluate()
+    coco_eval.accumulate()
+    coco_eval.summarize()
+    sys.stdout = sys.__stdout__
+    return coco_eval.stats
+
+
+def make_ap_table(score, past_result=[], epoch=-1):
+    ap_table = Table()
+    ap_table.add_column("Epoch", justify="center", style="white", width=5)
+    ap_table.add_column("Avg. Precision", justify="left", style="cyan")
+    ap_table.add_column("", justify="right", style="green", width=5)
+    ap_table.add_column("Avg. Recall", justify="left", style="cyan")
+    ap_table.add_column("", justify="right", style="green", width=5)
+
+    for eps, (ap_name1, ap_value1, ap_name2, ap_value2) in past_result:
+        ap_table.add_row(f"{eps: 3d}", ap_name1, f"{ap_value1:.2f}", ap_name2, f"{ap_value2:.2f}")
+    if past_result:
+        ap_table.add_row()
+
+    this_ap = ("AP @ .5:.95", score[0], "AP @        .5", score[1])
+    metrics = [
+        ("AP @ .5:.95", score[0], "AR maxDets   1", score[6]),
+        ("AP @     .5", score[1], "AR maxDets  10", score[7]),
+        ("AP @    .75", score[2], "AR maxDets 100", score[8]),
+        ("AP  (small)", score[3], "AR     (small)", score[9]),
+        ("AP (medium)", score[4], "AR    (medium)", score[10]),
+        ("AP  (large)", score[5], "AR     (large)", score[11]),
+    ]
+
+    for ap_name, ap_value, ar_name, ar_value in metrics:
+        ap_table.add_row(f"{epoch: 3d}", ap_name, f"{ap_value:.2f}", ar_name, f"{ar_value:.2f}")
+
+    return ap_table, this_ap