🔀 [Merge] remote-tracking branch 'origin/DDP_BUGS' into Lightning

yolo/tools/data_loader.py

@@ -22,6 +22,34 @@ from yolo.utils.dataset_utils import (
 from yolo.utils.logger import logger
 
 
+def tensorlize(data):
+    # TODO Move Tensorlize to helper
+    img_paths, bboxes = zip(*data)
+    max_box = max(bbox.size(0) for bbox in bboxes)
+    padded_bbox_list = []
+    for bbox in bboxes:
+        padding = torch.full((max_box, 5), -1, dtype=torch.float32)
+        padding[: bbox.size(0)] = bbox
+        padded_bbox_list.append(padding)
+    bboxes = np.stack(padded_bbox_list)
+    img_paths = np.array(img_paths)
+    return img_paths, bboxes
+
+
+def tensorlize(data):
+    # TODO Move Tensorlize to helper
+    img_paths, bboxes = zip(*data)
+    max_box = max(bbox.size(0) for bbox in bboxes)
+    padded_bbox_list = []
+    for bbox in bboxes:
+        padding = torch.full((max_box, 5), -1, dtype=torch.float32)
+        padding[: bbox.size(0)] = bbox
+        padded_bbox_list.append(padding)
+    bboxes = np.stack(padded_bbox_list)
+    img_paths = np.array(img_paths)
+    return img_paths, bboxes
+
+
 class YoloDataset(Dataset):
     def __init__(self, data_cfg: DataConfig, dataset_cfg: DatasetConfig, phase: str = "train2017"):
         augment_cfg = data_cfg.data_augment
@@ -31,7 +59,8 @@ class YoloDataset(Dataset):
         transforms = [eval(aug)(prob) for aug, prob in augment_cfg.items()]
         self.transform = AugmentationComposer(transforms, self.image_size)
         self.transform.get_more_data = self.get_more_data
-        self.data = self.load_data(Path(dataset_cfg.path), phase_name)
+        img_paths, bboxes = tensorlize(self.load_data(Path(dataset_cfg.path), phase_name))
+        self.img_paths, self.bboxes = img_paths, bboxes
 
     def load_data(self, dataset_path: Path, phase_name: str):
         """
@@ -132,9 +161,10 @@ class YoloDataset(Dataset):
 
     def get_data(self, idx):
         img_path, bboxes = self.data[idx]
+        valid_mask = bboxes[:, 0] != -1
         with Image.open(img_path) as img:
             img = img.convert("RGB")
-        return img, bboxes, img_path
+        return img, torch.from_numpy(bboxes[valid_mask]), img_path
 
     def get_more_data(self, num: int = 1):
         indices = torch.randint(0, len(self), (num,))
@@ -148,7 +178,7 @@ class YoloDataset(Dataset):
         return img, bboxes, rev_tensor, img_path
 
     def __len__(self) -> int:
-        return len(self.data)
+        return len(self.bboxes)
 
 
 def collate_fn(batch: List[Tuple[Tensor, Tensor]]) -> Tuple[Tensor, List[Tensor]]: