✨ [Add] Dataset for loading image, labels

utils/dataloader.py

@@ -0,0 +1,105 @@
+from PIL import Image
+from os import path
+import os
+import hydra
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from loguru import logger
+from tqdm.rich import tqdm
+import diskcache as dc
+
+
+class YoloDataset(Dataset):
+    def __init__(self, dataset_cfg: dict, phase="train", transform=None, mixup=None):
+        phase_name = dataset_cfg.get(phase, phase)
+
+        self.transform = transform
+        self.mixup = mixup
+        self.data = self.load_data(dataset_cfg.path, phase_name)
+
+    def load_data(self, dataset_path, phase_name):
+        cache = dc.Cache(path.join(dataset_path, ".cache"))
+
+        if phase_name not in cache:
+            logger.info("Generate {} Cache", phase_name)
+
+            images_path = path.join(dataset_path, phase_name, "images")
+            labels_path = path.join(dataset_path, phase_name, "labels")
+
+            cache[phase_name] = self.filter_data(images_path, labels_path)
+
+        logger.info("Load {} Cache", phase_name)
+        data = cache[phase_name]
+        cache.close()
+
+        return data
+
+    def filter_data(self, images_path, labels_path):
+        data = []
+        valid_input = 0
+        images_list = os.listdir(images_path)
+        images_list.sort()
+        for image_name in tqdm(images_list):
+            if not image_name.lower().endswith((".jpg", ".jpeg", ".png")):
+                continue
+            img_path = path.join(images_path, image_name)
+            base_name, _ = path.splitext(image_name)
+            label_name = base_name + ".txt"
+            label_path = path.join(labels_path, label_name)
+
+            if not path.isfile(label_path):
+                # logger.warning(f"Warning: No label file for {label_path}")
+                continue
+
+            labels = self.load_valid_labels(label_path)
+            if labels is not None:
+                data.append((img_path, labels))
+                valid_input += 1
+        logger.info("Finish Record {}/{}", valid_input, len(os.listdir(images_path)))
+        return data
+
+    def load_valid_labels(self, label_path):
+        bboxes = []
+        with open(label_path, "r") as file:
+            for line in file:
+                segment = list(map(float, line.strip().split()))
+                cls = segment[0]
+                # Ensure parts length is odd and more than two points
+                if len(segment) % 2 != 1 or len(segment) < 5:
+                    logger.warning(f"Warning: Format error in {label_path}")
+                    continue
+                points = np.array(segment[1:]).reshape(-1, 2)  # change points to n x 2
+                valid_idx = np.any((points <= 1) | (points >= 0), axis=1)  # filter outlier points
+                points = points[valid_idx]  # only keep valid points
+
+                bbox = torch.tensor([cls, *points.max(axis=0), *points.min(axis=0)])
+                bboxes.append(bbox)
+        if not bboxes:
+            logger.warning(f"Warning: No valid BBox in {label_path}")
+            return None
+        return torch.stack(bboxes)
+
+    def __getitem__(self, idx):
+        img_path, bboxes = self.data[idx]
+        img = Image.open(img_path).convert("RGB")
+
+        return img, bboxes
+
+    def __len__(self):
+        return len(self.images)
+
+
+@hydra.main(config_path="../config/data", config_name="coco", version_base=None)
+def main(cfg):
+    dataset = YoloDataset(cfg)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.path.append("./")
+    from tools.log_helper import custom_logger
+
+    custom_logger()
+    main()