π§βπ» [Add] Type Hint to dataloader
Browse files- 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:
|
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)
|