🔨 [Add] the PostProccess class to convert predicts

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/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