✨ [New] inference with coco format

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/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/tools/solver.py

@@ -1,9 +1,11 @@
 import json
 import os
+import sys
 import time
 
 import torch
 from loguru import logger
+from pycocotools.coco import COCO
 from torch import Tensor
 
 # TODO: We may can't use CUDA?
@@ -25,6 +27,7 @@ from yolo.utils.model_utils import (
     create_scheduler,
     predicts_to_json,
 )
+from yolo.utils.solver_utils import calculate_ap
 
 
 class ModelTrainer:
@@ -112,7 +115,7 @@ class ModelTrainer:
             epoch_loss = self.train_one_epoch(dataloader)
             self.progress.finish_one_epoch()
 
-            self.validator.solve(self.validation_dataloader)
+            self.validator.solve(self.validation_dataloader, epoch_idx=epoch)
 
 
 class ModelTester:
@@ -187,7 +190,12 @@ 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 = []
@@ -203,3 +211,7 @@ class ModelValidator:
         self.progress.finish_one_epoch()
         with open(self.json_path, "w") as f:
             json.dump(predict_json, f)
+
+        self.progress.run_coco()
+        result = calculate_ap(self.coco_gt, predict_json)
+        self.progress.finish_coco(result, epoch_idx)

yolo/utils/logging_utils.py

@@ -13,18 +13,26 @@ Example:
 
 import os
 import sys
+from collections import deque
 from typing import 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 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 +46,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,21 +72,32 @@ class ProgressLogger:
                 project="YOLO", resume="allow", mode="online", dir=self.save_path, id=None, name=exp_name
             )
 
+    def update_ap_table(self, ap_list, epoch_idx=-1):
+        ap_table, ap_main = make_ap_table(ap_list, self.ap_past_list, epoch_idx)
+        self.ap_past_list.append((epoch_idx, ap_main))
+        self.ap_table = ap_table
+
+        if self.use_wandb:
+            self.wandb.log({f"mAP/AP @ .5:.95": ap_main[1], f"mAP/AP @ .5": ap_main[3]})
+
+    def get_renderable(self):
+        return Group(*self.get_renderables(), self.ap_table)
+
     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("[cyan]Epochs  [white]| Loss | Box  | DFL  | BCE  |", total=num_epochs)
 
     def start_one_epoch(self, num_batches: int, optimizer: Optimizer = None, epoch_idx: int = None):
         self.num_batches = num_batches
-        if self.use_wandb:
+        if self.use_wandb and optimizer is not None:
             lr_values = [params["lr"] for params in optimizer.param_groups]
             lr_names = ["bias", "norm", "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)
+        self.batch_task = self.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")
+            self.update(self.batch_task, advance=1, description=f"[green]Validating")
             return
         if self.use_wandb:
             for loss_name, loss_value in loss_dict.items():
@@ -84,11 +107,20 @@ class ProgressLogger:
         for loss_name, loss_val in loss_dict.items():
             loss_str += f" {loss_val:2.2f} |"
 
-        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)
+        self.update(self.batch_task, advance=1, description=f"[green]Batches [white]{loss_str}")
+        self.update(self.task_epoch, advance=1 / self.num_batches)
+
+    def run_coco(self):
+        self.batch_task = self.add_task("[green]Run COCO", total=1)
+
+    def finish_coco(self, result, epoch_idx):
+        self.update_ap_table(result, epoch_idx)
+        self.update(self.batch_task, advance=1)
+        self.refresh()
+        self.remove_task(self.batch_task)
 
     def finish_one_epoch(self):
-        self.progress.remove_task(self.batch_task)
+        self.remove_task(self.batch_task)
 
     def finish_train(self):
         self.wandb.finish()

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