-
Notifications
You must be signed in to change notification settings - Fork 2
/
get_loader.py
48 lines (40 loc) · 2.08 KB
/
get_loader.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
from .mydataset import ImageFolder, ImageFilelist
from .unaligned_data_loader import UnalignedDataLoader
import os
import torch
from torch.utils.data import DataLoader
def get_loader_test(source_path, target_path, evaluation_path, transforms, batch_size=32):
source_folder = ImageFolder(os.path.join(source_path),
transforms[source_path])
target_folder_train = ImageFolder(os.path.join(target_path),
transform=transforms[target_path],
return_paths=False,train=True)
eval_folder_test = ImageFilelist(os.path.join(evaluation_path),
'/data/ugui0/ksaito/VISDA_tmp/image_list_val.txt',
transform=transforms[evaluation_path],
return_paths=True)
train_loader = UnalignedDataLoader()
train_loader.initialize(source_folder, target_folder_train, batch_size)
test_loader = torch.utils.data.DataLoader(
eval_folder_test,
batch_size=batch_size,
shuffle=False,
num_workers=4)
return train_loader, test_loader
def get_loader(source_path, target_path, evaluation_path, transforms, batch_size=32):
source_folder = ImageFolder(os.path.join(source_path),
transforms[source_path])
target_folder_train = ImageFolder(os.path.join(target_path),
transform=transforms[target_path],
return_paths=False)
eval_folder_test = ImageFolder(os.path.join(evaluation_path),
transform=transforms[evaluation_path],
return_paths=True)
train_loader = UnalignedDataLoader()
train_loader.initialize(source_folder, target_folder_train, batch_size)
test_loader = torch.utils.data.DataLoader(
eval_folder_test,
batch_size=batch_size,
shuffle=False,
num_workers=4)
return train_loader, test_loader