| 229 | |
| 230 | |
| 231 | def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train"): |
| 232 | if task == "inference": |
| 233 | return StreamDataLoader(data_cfg) |
| 234 | |
| 235 | if getattr(dataset_cfg, "auto_download", False): |
| 236 | prepare_dataset(dataset_cfg, task) |
| 237 | dataset = YoloDataset(data_cfg, dataset_cfg, task) |
| 238 | |
| 239 | return DataLoader( |
| 240 | dataset, |
| 241 | batch_size=data_cfg.batch_size, |
| 242 | num_workers=data_cfg.cpu_num, |
| 243 | pin_memory=data_cfg.pin_memory, |
| 244 | collate_fn=collate_fn, |
| 245 | ) |
| 246 | |
| 247 | |
| 248 | class StreamDataLoader: |