🔀 [Merge] branch 'main' into DATASET

demo/hf_demo.py

@@ -10,8 +10,8 @@ sys.path.append(str(Path(__file__).resolve().parent.parent))
 from yolo import (
     AugmentationComposer,
     NMSConfig,
+    PostProccess,
     Vec2Box,
-    bbox_nms,
     create_model,
     draw_bboxes,
 )
@@ -37,7 +37,7 @@ transform = AugmentationComposer([])
 
 
 def predict(model_name, image, nms_confidence, nms_iou):
-    global DEFAULT_MODEL, model, device, v2b, class_list
+    global DEFAULT_MODEL, model, device, v2b, class_list, post_proccess
     if model_name != DEFAULT_MODEL:
         model = load_model(model_name, device)
         v2b = Vec2Box(model, IMAGE_SIZE, device)
@@ -46,16 +46,15 @@ def predict(model_name, image, nms_confidence, nms_iou):
     image_tensor, _, rev_tensor = transform(image)
 
     image_tensor = image_tensor.to(device)[None]
-    rev_tensor = rev_tensor.to(device)
+    rev_tensor = rev_tensor.to(device)[None]
+
+    nms_config = NMSConfig(nms_confidence, nms_iou)
+    post_proccess = PostProccess(v2b, nms_config)
 
     with torch.no_grad():
         predict = model(image_tensor)
-        pred_class, _, pred_bbox = v2b(predict["Main"])
-
-    nms_config = NMSConfig(nms_confidence, nms_iou)
+        pred_bbox = post_proccess(predict, rev_tensor)
 
-    pred_bbox = pred_bbox / rev_tensor[0] - rev_tensor[None, None, 1:]
-    pred_bbox = bbox_nms(pred_class, pred_bbox, nms_config)
     result_image = draw_bboxes(image, pred_bbox, idx2label=class_list)
 
     return result_image

tests/test_tools/{test_module_helper.py → test_module_utils.py}

@@ -2,7 +2,6 @@ import sys
 from pathlib import Path
 
 import pytest
-import torch
 from torch import nn
 
 project_root = Path(__file__).resolve().parent.parent.parent

tests/test_utils/{test_dataaugment.py → test_data_augmentation.py}

@@ -54,7 +54,7 @@ def test_mosaic():
 
     # Mock parent with image_size and get_more_data method
     class MockParent:
-        image_size = 100
+        image_size = (100, 100)
 
         def get_more_data(self, num_images):
             return [(img, boxes) for _ in range(num_images)]

yolo/__init__.py

@@ -6,6 +6,7 @@ from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
 from yolo.utils.bounding_box_utils import Vec2Box, bbox_nms
 from yolo.utils.deploy_utils import FastModelLoader
 from yolo.utils.logging_utils import custom_logger
+from yolo.utils.model_utils import PostProccess
 
 all = [
     "create_model",
@@ -22,4 +23,5 @@ all = [
     "ModelTester",
     "ModelTrainer",
     "ModelValidator",
+    "PostProccess",
 ]

yolo/config/config.py

@@ -142,6 +142,7 @@ class Config:
 
     class_num: int
     class_list: List[str]
+    class_idx_id: List[int]
     image_size: List[int]
 
     out_path: str
@@ -164,3 +165,87 @@ class YOLOLayer(nn.Module):
 
     def __post_init__(self):
         super().__init__()
+
+
+IDX_TO_ID = [
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    10,
+    11,
+    13,
+    14,
+    15,
+    16,
+    17,
+    18,
+    19,
+    20,
+    21,
+    22,
+    23,
+    24,
+    25,
+    27,
+    28,
+    31,
+    32,
+    33,
+    34,
+    35,
+    36,
+    37,
+    38,
+    39,
+    40,
+    41,
+    42,
+    43,
+    44,
+    46,
+    47,
+    48,
+    49,
+    50,
+    51,
+    52,
+    53,
+    54,
+    55,
+    56,
+    57,
+    58,
+    59,
+    60,
+    61,
+    62,
+    63,
+    64,
+    65,
+    67,
+    70,
+    72,
+    73,
+    74,
+    75,
+    76,
+    77,
+    78,
+    79,
+    80,
+    81,
+    82,
+    84,
+    85,
+    86,
+    87,
+    88,
+    89,
+    90,
+]

yolo/lazy.py

@@ -9,7 +9,7 @@ sys.path.append(str(project_root))
 from yolo.config.config import Config
 from yolo.model.yolo import create_model
 from yolo.tools.data_loader import create_dataloader
-from yolo.tools.solver import ModelTester, ModelTrainer
+from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
 from yolo.utils.bounding_box_utils import Vec2Box
 from yolo.utils.deploy_utils import FastModelLoader
 from yolo.utils.logging_utils import ProgressLogger
@@ -37,6 +37,10 @@ def main(cfg: Config):
         tester = ModelTester(cfg, model, vec2box, progress, device)
         tester.solve(dataloader)
 
+    if cfg.task.task == "validation":
+        valider = ModelValidator(cfg.task, model, vec2box, progress, device)
+        valider.solve(dataloader)
+
 
 if __name__ == "__main__":
     main()

yolo/tools/loss_functions.py

@@ -39,9 +39,9 @@ class BoxLoss(nn.Module):
 
 
 class DFLoss(nn.Module):
-    def __init__(self, anchors_norm: Tensor, reg_max: int) -> None:
+    def __init__(self, vec2box: Vec2Box, reg_max: int) -> None:
         super().__init__()
-        self.anchors_norm = anchors_norm
+        self.anchors_norm = (vec2box.anchor_grid / vec2box.scaler[:, None])[None]
         self.reg_max = reg_max
 
     def forward(
@@ -72,7 +72,7 @@ class YOLOLoss:
         self.vec2box = vec2box
 
         self.cls = BCELoss()
-        self.dfl = DFLoss(vec2box.anchor_norm, reg_max)
+        self.dfl = DFLoss(vec2box, reg_max)
         self.iou = BoxLoss()
 
         self.matcher = BoxMatcher(loss_cfg.matcher, self.class_num, vec2box.anchor_grid)

yolo/tools/solver.py

@@ -1,3 +1,4 @@
+import json
 import os
 import time
 
@@ -15,12 +16,14 @@ 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, bbox_nms, calculate_map
+from yolo.utils.bounding_box_utils import Vec2Box
 from yolo.utils.logging_utils import ProgressLogger, log_model_structure
 from yolo.utils.model_utils import (
     ExponentialMovingAverage,
+    PostProccess,
     create_optimizer,
     create_scheduler,
+    predicts_to_json,
 )
 
 
@@ -116,10 +119,9 @@ class ModelTester:
     def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: ProgressLogger, device):
         self.model = model
         self.device = device
-        self.vec2box = vec2box
         self.progress = progress
 
-        self.nms = cfg.task.nms
+        self.post_proccess = PostProccess(vec2box, cfg.task.nms)
         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)
@@ -141,9 +143,8 @@ class ModelTester:
                 rev_tensor = rev_tensor.to(self.device)
                 with torch.no_grad():
                     predicts = self.model(images)
-                    predicts = self.vec2box(predicts["Main"])
-                nms_out = bbox_nms(predicts[0], predicts[2], self.nms)
-                img = draw_bboxes(images, nms_out, idx2label=self.idx2label)
+                    predicts = self.post_proccess(predicts, rev_tensor)
+                img = draw_bboxes(origin_frame, predicts, idx2label=self.idx2label)
 
                 if dataloader.is_stream:
                     img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
@@ -176,32 +177,29 @@ class ModelValidator:
         validation_cfg: ValidationConfig,
         model: YOLO,
         vec2box: Vec2Box,
-        device,
         progress: ProgressLogger,
+        device,
     ):
         self.model = model
-        self.vec2box = vec2box
         self.device = device
         self.progress = progress
 
-        self.nms = validation_cfg.nms
+        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):
         # logger.info("🧪 Start Validation!")
         self.model.eval()
-        # TODO: choice mAP metrics?
-        iou_thresholds = torch.arange(0.5, 1.0, 0.05)
-        map_all = []
+        predict_json = []
         self.progress.start_one_epoch(len(dataloader))
         for 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.vec2box(predicts["Main"])
-            nms_out = bbox_nms(predicts[0], predicts[2], self.nms)
-            for idx, predict in enumerate(nms_out):
-                map_value = calculate_map(predict, targets[idx], iou_thresholds)
-                map_all.append(map_value[0])
-            self.progress.one_batch(mapp=torch.Tensor(map_all).mean())
+                predicts = self.post_proccess(predicts, rev_tensor)
+            self.progress.one_batch()
 
+            predict_json.extend(predicts_to_json(img_paths, predicts))
         self.progress.finish_one_epoch()
+        with open(self.json_path, "w") as f:
+            json.dump(predict_json, f)

yolo/utils/bounding_box_utils.py

@@ -108,12 +108,13 @@ def transform_bbox(bbox: Tensor, indicator="xywh -> xyxy"):
     return bbox.to(dtype=data_type)
 
 
-def generate_anchors(image_size: List[int], anchors_list: List[Tuple[int]]):
+def generate_anchors(image_size: List[int], strides: List[int]):
     """
     Find the anchor maps for each w, h.
 
     Args:
-        anchors_list List[[w1, h1], [w2, h2], ...]: the anchor num for each predicted anchor
+        image_size List: the image size of augmented image size
+        strides List[8, 16, 32, ...]: the stride size for each predicted layer
 
     Returns:
         all_anchors [HW x 2]:
@@ -122,15 +123,14 @@ def generate_anchors(image_size: List[int], anchors_list: List[Tuple[int]]):
     W, H = image_size
     anchors = []
     scaler = []
-    for anchor_wh in anchors_list:
-        stride = W // anchor_wh[0]
-        anchor_num = anchor_wh[0] * anchor_wh[1]
+    for stride in strides:
+        anchor_num = W // stride * H // stride
         scaler.append(torch.full((anchor_num,), stride))
         shift = stride // 2
-        x = torch.arange(0, W, stride) + shift
-        y = torch.arange(0, H, stride) + shift
-        anchor_x, anchor_y = torch.meshgrid(x, y, indexing="ij")
-        anchor = torch.stack([anchor_y.flatten(), anchor_x.flatten()], dim=-1)
+        h = torch.arange(0, H, stride) + shift
+        w = torch.arange(0, W, stride) + shift
+        anchor_h, anchor_w = torch.meshgrid(h, w, indexing="ij")
+        anchor = torch.stack([anchor_w.flatten(), anchor_h.flatten()], dim=-1)
         anchors.append(anchor)
     all_anchors = torch.cat(anchors, dim=0)
     all_scalers = torch.cat(scaler, dim=0)
@@ -172,6 +172,7 @@ class BoxMatcher:
         Returns:
             [batch x targets x anchors]: The probabilities from `pred_cls` corresponding to the class indices specified in `target_cls`.
         """
+        # TODO: Turn 8400 to HW
         target_cls = target_cls.expand(-1, -1, 8400)
         predict_cls = predict_cls.transpose(1, 2)
         cls_probabilities = torch.gather(predict_cls, 1, target_cls)
@@ -266,24 +267,34 @@ class BoxMatcher:
 
 class Vec2Box:
     def __init__(self, model: YOLO, image_size, device):
-        if getattr(model, "strides", None) is None:
-            logger.info("🧸 Found no anchor, Make a dummy test for auto-anchor size")
-            dummy_input = torch.zeros(1, 3, *image_size).to(device)
-            dummy_output = model(dummy_input)
-            anchors_num = []
-            for predict_head in dummy_output["Main"]:
-                _, _, *anchor_num = predict_head[2].shape
-                anchors_num.append(anchor_num)
+        self.device = device
+
+        if getattr(model, "strides"):
+            logger.info(f"🈶 Found stride of model {model.strides}")
+            self.strides = model.strides
         else:
-            logger.info(f"🈶 Found anchor {model.strides}")
-            anchors_num = [[image_size[0] // stride, image_size[0] // stride] for stride in model.strides]
+            logger.info("🧸 Found no stride of model, performed a dummy test for auto-anchor size")
+            self.strides = self.create_auto_anchor(model, image_size)
 
+        # TODO: this is a exception of onnx, remove it when onnx device if fixed
         if not isinstance(model, YOLO):
             device = torch.device("cpu")
 
-        anchor_grid, scaler = generate_anchors(image_size, anchors_num)
+        anchor_grid, scaler = generate_anchors(image_size, self.strides)
         self.anchor_grid, self.scaler = anchor_grid.to(device), scaler.to(device)
-        self.anchor_norm = (anchor_grid / scaler[:, None])[None].to(device)
+
+    def create_auto_anchor(self, model: YOLO, image_size):
+        dummy_input = torch.zeros(1, 3, *image_size).to(self.device)
+        dummy_output = model(dummy_input)
+        strides = []
+        for predict_head in dummy_output["Main"]:
+            _, _, *anchor_num = predict_head[2].shape
+            strides.append(image_size[1] // anchor_num[1])
+        return strides
+
+    def update(self, image_size):
+        anchor_grid, scaler = generate_anchors(image_size, self.strides)
+        self.anchor_grid, self.scaler = anchor_grid.to(self.device), scaler.to(self.device)
 
     def __call__(self, predicts):
         preds_cls, preds_anc, preds_box = [], [], []

yolo/utils/logging_utils.py

@@ -72,9 +72,9 @@ class ProgressLogger:
                 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, mapp=None):
+    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]Batches [white]{mapp:.2%}")
+            self.progress.update(self.batch_task, advance=1, description=f"[green]Validating")
             return
         if self.use_wandb:
             for loss_name, loss_value in loss_dict.items():

yolo/utils/model_utils.py

@@ -1,17 +1,18 @@
 import os
-from typing import List, Type, Union
+from pathlib import Path
+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 nn
-from torch.nn.parallel import DistributedDataParallel as DDP
+from torch import Tensor
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LambdaLR, SequentialLR, _LRScheduler
 
-from yolo.config.config import OptimizerConfig, SchedulerConfig
+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
 
 
 class ExponentialMovingAverage:
@@ -93,3 +94,40 @@ def get_device(device_spec: Union[str, int, List[int]]) -> torch.device:
         device_spec = initialize_distributed()
     device = torch.device(device_spec)
     return device, ddp_flag
+
+
+class PostProccess:
+    """
+    TODO: function document
+    scale back the prediction and do nms for pred_bbox
+    """
+
+    def __init__(self, vec2box, nms_cfg: NMSConfig) -> None:
+        self.vec2box = vec2box
+        self.nms = nms_cfg
+
+    def __call__(self, predict, rev_tensor: Optional[Tensor]):
+        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]
+        pred_bbox = bbox_nms(pred_class, pred_bbox, self.nms)
+        return pred_bbox
+
+
+def predicts_to_json(img_paths, predicts):
+    """
+    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):
+        bboxes[:, 1:5] = transform_bbox(bboxes[:, 1:5], "xyxy -> xywh")
+        for cls, *pos, conf in bboxes:
+            bbox = {
+                "image_id": int(Path(img_path).stem),
+                "category_id": IDX_TO_ID[int(cls)],
+                "bbox": [float(p) for p in pos],
+                "score": float(conf),
+            }
+            batch_json.append(bbox)
+    return batch_json