🔀 [Merge] branch 'INFERENCE' into DEPLOY

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/data_augmentation.py

@@ -10,7 +10,7 @@ class AugmentationComposer:
     def __init__(self, transforms, image_size: int = [640, 640]):
         self.transforms = transforms
         # TODO: handle List of image_size [640, 640]
-        self.image_size = image_size[0]
+        self.image_size = image_size
         self.pad_resize = PadAndResize(self.image_size)
 
         for transform in self.transforms:
@@ -29,27 +29,32 @@ class AugmentationComposer:
 
 
 class PadAndResize:
-    def __init__(self, image_size):
+    def __init__(self, image_size, background_color=(128, 128, 128)):
         """Initialize the object with the target image size."""
-        self.image_size = image_size
+        self.target_width, self.target_height = image_size
+        self.background_color = background_color
 
-    def __call__(self, image, boxes):
-        original_size = max(image.size)
-        scale = self.image_size / original_size
-        square_img = Image.new("RGB", (original_size, original_size), (128, 128, 128))
-        left = (original_size - image.width) // 2
-        top = (original_size - image.height) // 2
-        square_img.paste(image, (left, top))
+    def __call__(self, image: Image, boxes):
+        img_width, img_height = image.size
+        scale = min(self.target_width / img_width, self.target_height / img_height)
+        new_width, new_height = int(img_width * scale), int(img_height * scale)
+
+        resized_image = image.resize((new_width, new_height), Image.LANCZOS)
 
-        resized_img = square_img.resize((self.image_size, self.image_size))
+        pad_left = (self.target_width - new_width) // 2
+        pad_top = (self.target_height - new_height) // 2
+        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] = (boxes[:, 1] * image.width + left) / self.image_size * scale
-        boxes[:, 2] = (boxes[:, 2] * image.height + top) / self.image_size * scale
-        boxes[:, 3] = (boxes[:, 3] * image.width + left) / self.image_size * scale
-        boxes[:, 4] = (boxes[:, 4] * image.height + top) / self.image_size * scale
+        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
 
-        rev_tensor = torch.tensor([scale, left, top, left, top])
-        return resized_img, boxes, rev_tensor
+        transform_info = torch.tensor([scale, pad_left, pad_top, pad_left, pad_top])
+        return padded_image, boxes, transform_info
 
 
 class HorizontalFlip:
@@ -94,7 +99,7 @@ class Mosaic:
 
         assert self.parent is not None, "Parent is not set. Mosaic cannot retrieve image size."
 
-        img_sz = self.parent.image_size  # Assuming `image_size` is defined in parent
+        img_sz = self.parent.image_size[0]  # Assuming `image_size` is defined in parent
         more_data = self.parent.get_more_data(3)  # get 3 more images randomly
 
         data = [(image, boxes)] + more_data

yolo/tools/data_loader.py

@@ -141,16 +141,16 @@ class YoloDataset(Dataset):
     def get_data(self, idx):
         img_path, bboxes = self.data[idx]
         img = Image.open(img_path).convert("RGB")
-        return img, bboxes
+        return img, bboxes, img_path
 
     def get_more_data(self, num: int = 1):
         indices = torch.randint(0, len(self), (num,))
-        return [self.get_data(idx) for idx in indices]
+        return [self.get_data(idx)[:2] for idx in indices]
 
     def __getitem__(self, idx) -> Union[Image.Image, torch.Tensor]:
-        img, bboxes = self.get_data(idx)
-        img, bboxes, _ = self.transform(img, bboxes)
-        return img, bboxes
+        img, bboxes, img_path = self.get_data(idx)
+        img, bboxes, rev_tensor = self.transform(img, bboxes)
+        return img, bboxes, rev_tensor, img_path
 
     def __len__(self) -> int:
         return len(self.data)
@@ -195,9 +195,11 @@ class YoloDataLoader(DataLoader):
             batch_targets[idx, :target_size] = batch[idx][1]
         batch_targets[:, :, 1:] *= self.image_size
 
-        batch_images = torch.stack([item[0] for item in batch])
+        batch_images, _, batch_reverse, batch_path = zip(*batch)
+        batch_images = torch.stack(batch_images)
+        batch_reverse = torch.stack(batch_reverse)
 
-        return batch_images, batch_targets
+        return batch_images, batch_targets, batch_reverse, batch_path
 
 
 def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train", use_ddp: bool = False):
@@ -261,12 +263,14 @@ class StreamDataLoader:
         if isinstance(frame, np.ndarray):
             frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
             frame = Image.fromarray(frame)
+        origin_frame = frame
         frame, _, rev_tensor = self.transform(frame, torch.zeros(0, 5))
         frame = frame[None]
+        rev_tensor = rev_tensor[None]
         if not self.is_stream:
-            self.queue.put(frame)
+            self.queue.put((frame, rev_tensor, origin_frame))
         else:
-            self.current_frame = frame
+            self.current_frame = (frame, rev_tensor, origin_frame)
 
     def __iter__(self) -> Generator[Tensor, None, None]:
         return self

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,
 )
 
 
@@ -72,7 +75,7 @@ class ModelTrainer:
         self.model.train()
         total_loss = 0
 
-        for images, targets in dataloader:
+        for images, targets, *_ in dataloader:
             loss, loss_each = self.train_one_batch(images, targets)
 
             total_loss += loss
@@ -135,8 +138,9 @@ class ModelTester:
 
             last_time = time.time()
         try:
-            for idx, images in enumerate(dataloader):
+            for idx, (images, rev_tensor, origin_frame) in enumerate(dataloader):
                 images = images.to(self.device)
+                rev_tensor = rev_tensor.to(self.device)
                 with torch.no_grad():
                     predicts = self.model(images)
                     predicts = self.post_proccess(predicts, rev_tensor)
@@ -173,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 in dataloader:
-            images, targets = images.to(self.device), targets.to(self.device)
+        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/logging_utils.py

@@ -70,9 +70,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