henry000 commited on
Commit
d252042
Β·
1 Parent(s): fce8aa7

πŸ§‘β€πŸ’» [Add] Type Hint to dataloader

Browse files
Files changed (1) hide show
  1. yolo/utils/dataloader.py +2 -1
yolo/utils/dataloader.py CHANGED
@@ -11,6 +11,7 @@ from torch.utils.data import DataLoader, Dataset
11
  from torchvision.transforms import functional as TF
12
  from tqdm.rich import tqdm
13
 
 
14
  from yolo.tools.dataset_helper import (
15
  create_image_info_dict,
16
  find_labels_path,
@@ -149,7 +150,7 @@ class YoloDataset(Dataset):
149
 
150
 
151
  class YoloDataLoader(DataLoader):
152
- def __init__(self, config: dict):
153
  """Initializes the YoloDataLoader with hydra-config files."""
154
  hyper = config.hyper.data
155
  dataset = YoloDataset(config)
 
11
  from torchvision.transforms import functional as TF
12
  from tqdm.rich import tqdm
13
 
14
+ from yolo.config.config import Config
15
  from yolo.tools.dataset_helper import (
16
  create_image_info_dict,
17
  find_labels_path,
 
150
 
151
 
152
  class YoloDataLoader(DataLoader):
153
+ def __init__(self, config: Config):
154
  """Initializes the YoloDataLoader with hydra-config files."""
155
  hyper = config.hyper.data
156
  dataset = YoloDataset(config)