✨ [Finish] Dataloder and get_dataloader

config/config.yaml

@@ -7,4 +7,5 @@ defaults:
  - download: ../data/download
  - augmentation: ../data/augmentation
  - model: v7-base
+ - hyper: default
  - _self_

config/hyper/default.yaml

@@ -0,0 +1,5 @@
+data:
+  batch_size: 4 
+  shuffle: True
+  num_workers: 4
+  pin_memory: True

utils/dataloader.py

@@ -1,5 +1,5 @@
 from os import listdir, path
-from typing import Union
+from typing import List, Tuple, Union
 
 import diskcache as dc
 import hydra
@@ -7,16 +7,11 @@ import numpy as np
 import torch
 from loguru import logger
 from PIL import Image
-from torch.utils.data import Dataset
+from torch.utils.data import DataLoader, Dataset
+from torchvision.transforms import functional as TF
 from tqdm.rich import tqdm
 
-from utils.data_augment import (
-    Compose,
-    MixUp,
-    Mosaic,
-    RandomHorizontalFlip,
-    RandomVerticalFlip,
-)
+from utils.data_augment import Compose, HorizontalFlip, MixUp, Mosaic, VerticalFlip
 from utils.drawer import draw_bboxes
 
 
@@ -130,16 +125,55 @@ class YoloDataset(Dataset):
         img, bboxes = self.get_data(idx)
         if self.transform:
             img, bboxes = self.transform(img, bboxes)
+        img = TF.to_tensor(img)
         return img, bboxes
 
     def __len__(self) -> int:
         return len(self.data)
 
 
+class YoloDataLoader(DataLoader):
+    def __init__(self, config: dict):
+        """Initializes the YoloDataLoader with hydra-config files."""
+        hyper = config.hyper.data
+        dataset = YoloDataset(config)
+
+        super().__init__(
+            dataset,
+            batch_size=hyper.batch_size,
+            shuffle=hyper.shuffle,
+            num_workers=hyper.num_workers,
+            pin_memory=hyper.pin_memory,
+            collate_fn=self.collate_fn,
+        )
+
+    def collate_fn(self, batch: List[Tuple[torch.Tensor, torch.Tensor]]) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        """
+        A collate function to handle batching of images and their corresponding targets.
+
+        Args:
+            batch (list of tuples): Each tuple contains:
+                - image (torch.Tensor): The image tensor.
+                - labels (torch.Tensor): The tensor of labels for the image.
+
+        Returns:
+            Tuple[torch.Tensor, List[torch.Tensor]]: A tuple containing:
+                - A tensor of batched images.
+                - A list of tensors, each corresponding to bboxes for each image in the batch.
+        """
+        images = torch.stack([item[0] for item in batch])
+        targets = [item[1] for item in batch]
+        return images, targets
+
+
+def get_dataloader(config):
+    return YoloDataLoader(config)
+
+
 @hydra.main(config_path="../config", config_name="config", version_base=None)
 def main(cfg):
-    dataset = YoloDataset(cfg)
-    draw_bboxes(*dataset[0])
+    dataloader = get_dataloader(cfg)
+    draw_bboxes(next(iter(dataloader)))
 
 
 if __name__ == "__main__":

utils/drawer.py

@@ -1,23 +1,31 @@
+from typing import List, Union
+
+import torch
+from loguru import logger
 from PIL import Image, ImageDraw, ImageFont
+from torchvision.transforms.functional import to_pil_image
 
 
-def draw_bboxes(img, bboxes):
+def draw_bboxes(img: Union[Image.Image, torch.Tensor], bboxes: List[List[Union[int, float]]]):
     """
     Draw bounding boxes on an image.
 
     Args:
-    - image_path (str): Path to the image file.
-    - bboxes (list of lists/tuples): Bounding boxes with [x_min, y_min, x_max, y_max, class_id].
+    - img (PIL Image or torch.Tensor): Image on which to draw the bounding boxes.
+    - bboxes (List of Lists/Tensors): Bounding boxes with [class_id, x_min, y_min, x_max, y_max],
+      where coordinates are normalized [0, 1].
     """
-    # Load an image
-    draw = ImageDraw.Draw(img)
+    # Convert tensor image to PIL Image if necessary
+    if isinstance(img, torch.Tensor):
+        if img.dim() > 3:
+            logger.info("Multi-frame tensor detected, using the first image.")
+            img = img[0]
+            bboxes = bboxes[0]
+        img = to_pil_image(img)
 
-    # Font for class_id (optional)
-    try:
-        font = ImageFont.truetype("arial.ttf", 30)
-    except IOError:
-        font = ImageFont.load_default(30)
+    draw = ImageDraw.Draw(img)
     width, height = img.size
+    font = ImageFont.load_default(30)
 
     for bbox in bboxes:
         class_id, x_min, y_min, x_max, y_max = bbox
@@ -26,7 +34,8 @@ def draw_bboxes(img, bboxes):
         y_min = y_min * height
         y_max = y_max * height
         shape = [(x_min, y_min), (x_max, y_max)]
-        draw.rectangle(shape, outline="red", width=2)
+        draw.rectangle(shape, outline="red", width=3)
         draw.text((x_min, y_min), str(int(class_id)), font=font, fill="blue")
 
-    img.save("output.jpg")
+    img.save("visualize.jpg")  # Save the image with annotations
+    logger.info("Saved visualize image at visualize.png")