builder.py
1.34 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
import copy
from torch.utils.data import DataLoader
from utils.registery import DATASET_REGISTRY
from .CoordinatesData import CoordinatesData
def build_dataset(cfg):
dataset_cfg = copy.deepcopy(cfg)
try:
dataset_cfg = dataset_cfg['dataset']
except Exception:
raise 'should contain {dataset}!'
train_cfg = copy.deepcopy(dataset_cfg)
val_cfg = copy.deepcopy(dataset_cfg)
train_cfg['args']['anno_file'] = train_cfg['args'].pop('train_anno_file')
train_cfg['args'].pop('val_anno_file', None)
train_cfg['args']['phase'] = 'train'
val_cfg['args']['anno_file'] = val_cfg['args'].pop('val_anno_file')
val_cfg['args'].pop('train_anno_file', None)
val_cfg['args']['phase'] = 'valid'
train_data = DATASET_REGISTRY.get(cfg['dataset']['name'])(**train_cfg['args'])
val_data = DATASET_REGISTRY.get(cfg['dataset']['name'])(**val_cfg['args'])
return train_data, val_data
def build_dataloader(cfg):
dataloader_cfg = copy.deepcopy(cfg)
try:
dataloader_cfg = cfg['dataloader']
except Exception:
raise 'should contain {dataloader}!'
train_ds, val_ds = build_dataset(cfg)
train_loader = DataLoader(train_ds,
**dataloader_cfg)
val_loader = DataLoader(val_ds,
**dataloader_cfg)
return train_loader, val_loader