lucytuan commited on
Commit
3154a8a
Β·
1 Parent(s): dc3d6a2

πŸ› [Fix] dataset structure in dataloader.py

Browse files

Also fix the draw_bboxes untuple in dataloader.py

Files changed (1) hide show
  1. utils/dataloader.py +4 -4
utils/dataloader.py CHANGED
@@ -16,7 +16,7 @@ from utils.drawer import draw_bboxes
16
 
17
 
18
  class YoloDataset(Dataset):
19
- def __init__(self, config: dict, phase: str = "train", image_size: int = 640):
20
  dataset_cfg = config.data
21
  augment_cfg = config.augmentation
22
  phase_name = dataset_cfg.get(phase, phase)
@@ -44,8 +44,8 @@ class YoloDataset(Dataset):
44
 
45
  if data is None:
46
  logger.info("Generating {} cache", phase_name)
47
- images_path = path.join(dataset_path, phase_name, "images")
48
- labels_path = path.join(dataset_path, phase_name, "labels")
49
  data = self.filter_data(images_path, labels_path)
50
  cache[phase_name] = data
51
 
@@ -173,7 +173,7 @@ def get_dataloader(config):
173
  @hydra.main(config_path="../config", config_name="config", version_base=None)
174
  def main(cfg):
175
  dataloader = get_dataloader(cfg)
176
- draw_bboxes(next(iter(dataloader)))
177
 
178
 
179
  if __name__ == "__main__":
 
16
 
17
 
18
  class YoloDataset(Dataset):
19
+ def __init__(self, config: dict, phase: str = "train2017", image_size: int = 640):
20
  dataset_cfg = config.data
21
  augment_cfg = config.augmentation
22
  phase_name = dataset_cfg.get(phase, phase)
 
44
 
45
  if data is None:
46
  logger.info("Generating {} cache", phase_name)
47
+ images_path = path.join(dataset_path, "images", phase_name)
48
+ labels_path = path.join(dataset_path, "label", phase_name)
49
  data = self.filter_data(images_path, labels_path)
50
  cache[phase_name] = data
51
 
 
173
  @hydra.main(config_path="../config", config_name="config", version_base=None)
174
  def main(cfg):
175
  dataloader = get_dataloader(cfg)
176
+ draw_bboxes(*next(iter(dataloader)))
177
 
178
 
179
  if __name__ == "__main__":