-
Notifications
You must be signed in to change notification settings - Fork 6
/
utils.py
107 lines (90 loc) · 4.17 KB
/
utils.py
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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
import numpy as np
import torch
from torch.utils.data import DataLoader
from torchvision import utils
import data_config
from datasets.CD_dataset import CDDataset, xBDataset, xBDatasetMulti
def get_loader(data_name, img_size=256, batch_size=8, split='test',
is_train=False, dataset='CDDataset', patch=None):
dataConfig = data_config.DataConfig().get_data_config(data_name)
root_dir = dataConfig.root_dir
label_transform = dataConfig.label_transform
print(dataConfig)
if dataset == 'CDDataset':
data_set = CDDataset(root_dir=root_dir, split=split,
img_size=img_size, is_train=is_train,
label_transform=label_transform, patch=patch)
elif dataset == 'xBDataset':
data_set = xBDataset(root_dir=root_dir, split=split,
img_size=img_size, is_train=is_train,
label_transform=label_transform)
elif dataset == 'xBDatasetMulti':
data_set = xBDatasetMulti(root_dir=root_dir, split=split,
img_size=img_size, is_train=is_train,
label_transform=label_transform)
else:
raise NotImplementedError(
'Wrong dataset name %s (choose one from [CDDataset])'
% dataset)
shuffle = is_train
dataloader = DataLoader(data_set, batch_size=batch_size,
shuffle=False, num_workers=4)
return dataloader
def get_loaders(args):
data_name = args.data_name
dataConfig = data_config.DataConfig().get_data_config(data_name)
root_dir = dataConfig.root_dir
label_transform = dataConfig.label_transform
split = args.split
split_val = 'val'
if hasattr(args, 'split_val'):
split_val = args.split_val
if args.dataset == 'CDDataset':
training_set = CDDataset(root_dir=root_dir, split=split,
img_size=args.img_size,is_train=True,
label_transform=label_transform)
val_set = CDDataset(root_dir=root_dir, split=split_val,
img_size=args.img_size,is_train=False,
label_transform=label_transform)
elif args.dataset == 'xBDataset':
training_set = xBDataset(root_dir=root_dir, split=split,
img_size=args.img_size,is_train=True,
label_transform=label_transform)
val_set = xBDataset(root_dir=root_dir, split=split_val,
img_size=args.img_size,is_train=False,
label_transform=label_transform)
elif args.dataset == 'xBDatasetMulti':
training_set = xBDatasetMulti(root_dir=root_dir, split=split,
img_size=args.img_size,is_train=True,
label_transform=label_transform)
val_set = xBDatasetMulti(root_dir=root_dir, split=split_val,
img_size=args.img_size,is_train=False,
label_transform=label_transform)
else:
raise NotImplementedError(
'Wrong dataset name %s (choose one from [CDDataset,])'
% args.dataset)
datasets = {'train': training_set, 'val': val_set}
dataloaders = {x: DataLoader(datasets[x], batch_size=args.batch_size,
shuffle=True, num_workers=args.num_workers)
for x in ['train', 'val']}
return dataloaders
def make_numpy_grid(tensor_data, pad_value=0,padding=0):
# tensor_data = tensor_data.detach()
vis = utils.make_grid(tensor_data, pad_value=pad_value,padding=padding)
vis = np.array(vis.cpu()).transpose((1,2,0))
if vis.shape[2] == 1:
vis = np.stack([vis, vis, vis], axis=-1)
return vis
def de_norm(tensor_data):
return tensor_data * 0.5 + 0.5
def get_device(args):
# set gpu ids
str_ids = args.gpu_ids.split(',')
args.gpu_ids = []
for str_id in str_ids:
id = int(str_id)
if id >= 0:
args.gpu_ids.append(id)
if len(args.gpu_ids) > 0:
torch.cuda.set_device(args.gpu_ids[0])