def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
### input A (label maps)
dir_A = '_A' if self.opt.label_nc == 0 else '_label'
self.dir_A = os.path.join(opt.dataroot, opt.phase + dir_A)
self.A_paths = sorted(make_dataset(self.dir_A))
### input B (real images)
if opt.isTrain:
dir_B = '_B' if self.opt.label_nc == 0 else '_img'
self.dir_B = os.path.join(opt.dataroot, opt.phase + dir_B)
self.B_paths = sorted(make_dataset(self.dir_B))
### instance maps
if not opt.no_instance:
self.dir_inst = os.path.join(opt.dataroot, opt.phase + '_inst')
self.inst_paths = sorted(make_dataset(self.dir_inst))
### load precomputed instance-wise encoded features
if opt.load_features:
self.dir_feat = os.path.join(opt.dataroot, opt.phase + '_feat')
print('----------- loading features from %s ----------' % self.dir_feat)
self.feat_paths = sorted(make_dataset(self.dir_feat))
self.dataset_size = len(self.A_paths)
def initialize(self, dataroots, load_size=64):
if(not isinstance(dataroots,list)):
dataroots = [dataroots,]
self.roots = dataroots
self.load_size = load_size
# image directory
self.dir_ref = [os.path.join(root, 'ref') for root in self.roots]
self.ref_paths = make_dataset(self.dir_ref)
self.ref_paths = sorted(self.ref_paths)
self.dir_p0 = [os.path.join(root, 'p0') for root in self.roots]
self.p0_paths = make_dataset(self.dir_p0)
self.p0_paths = sorted(self.p0_paths)
self.dir_p1 = [os.path.join(root, 'p1') for root in self.roots]
self.p1_paths = make_dataset(self.dir_p1)
self.p1_paths = sorted(self.p1_paths)
transform_list = []
transform_list.append(transforms.Scale(load_size))
transform_list += [transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),(0.5, 0.5, 0.5))]
self.transform = transforms.Compose(transform_list)
# judgement directory
self.dir_J = [os.path.join(root, 'judge') for root in self.roots]
self.judge_paths = make_dataset(self.dir_J,mode='np')
self.judge_paths = sorted(self.judge_paths)
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot)
self.dir_img = os.path.join(self.dir_A, 'img')
self.dir_prior = os.path.join(self.dir_A, 'prior')
self.img_paths = sorted(make_dataset(self.dir_img))
self.prior_paths = sorted(make_dataset(self.dir_prior))
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_AB = os.path.join(opt.dataroot)
self.dir_img = os.path.join(self.dir_AB, 'img')
self.dir_priors = os.path.join(self.dir_AB, 'prior')
self.dir_parsing = os.path.join(self.dir_AB, 'parsing')
self.img_paths = sorted(make_dataset(self.dir_img))
self.parsing_paths = sorted(make_dataset(self.dir_parsing))
assert(opt.resize_or_crop == 'resize_and_crop')
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')
self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B')
self.A_paths = make_dataset(self.dir_A)
self.B_paths = make_dataset(self.dir_B)
self.A_paths = sorted(self.A_paths)
self.B_paths = sorted(self.B_paths)
self.A_size = len(self.A_paths)
self.B_size = len(self.B_paths)
self.transform = get_transform(opt)
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
if opt.phase == 'test' :
self.dir_AB = opt.dataroot
else:
self.dir_AB = os.path.join(opt.dataroot, opt.phase)
self.AB_paths = sorted(make_dataset(self.dir_AB))
assert(opt.resize_or_crop == 'resize_and_crop')
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir = os.path.join(opt.dataroot, opt.phase)
self.paths = make_dataset(self.dir)
self.paths = sorted(self.paths)
self.size = len(self.paths)
self.fineSize = opt.fineSize
self.transform = get_transform(opt)
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')
self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B')
self.dir_C = os.path.join(opt.dataroot, opt.phase + 'C')
self.A_paths = make_dataset(self.dir_A)
self.B_paths = make_dataset(self.dir_B)
self.C_paths = make_dataset(self.dir_C)
self.A_paths = sorted(self.A_paths)
self.B_paths = sorted(self.B_paths)
self.C_paths = sorted(self.C_paths)
self.A_size = len(self.A_paths)
self.B_size = len(self.B_paths)
self.C_size = len(self.C_paths)
self.transform = get_transform(opt)
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.A_paths = sorted(make_dataset(opt.dataroot, opt.max_dataset_size))
self.transform = get_transform(opt, grayscale=(self.opt.input_nc == 1))
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot.format(phase=opt.phase, setname='A'))
self.dir_B = os.path.join(opt.dataroot.format(phase=opt.phase, setname='B'))
self.A_paths = make_dataset(self.dir_A)
self.B_paths = make_dataset(self.dir_B)
self.A_paths = sorted(self.A_paths)
self.B_paths = sorted(self.B_paths)
self.A_size = len(self.A_paths)
self.B_size = len(self.B_paths)
# self.transform = get_transform(opt)
transform_list = [transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5))]
self.transform = transforms.Compose(transform_list)
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.dir_A = os.path.join(opt.dataroot, opt.phase+'A') # get the image directory
self.dir_B = os.path.join(opt.dataroot, opt.phase+'B')
self.A_paths = sorted(make_dataset(self.dir_A, opt.max_dataset_size)) # get image paths
self.B_paths = sorted(make_dataset(self.dir_B, opt.max_dataset_size))
# self.A_paths = make_dataset(self.dir_A, opt.max_dataset_size)
# self.B_paths = make_dataset(self.dir_B, opt.max_dataset_size)
# self.A_paths.sort(key=lambda x: int(x.rstrip("_gt_2d.tif").split("/")[-1]))
# self.B_paths.sort(key=lambda x: int(x.rstrip(".tif").split("/")[-1]))
assert(self.opt.load_size >= self.opt.crop_size) # crop_size should be smaller than the size of loaded image
self.input_nc = self.opt.output_nc if self.opt.direction == 'BtoA' else self.opt.input_nc
self.output_nc = self.opt.input_nc if self.opt.direction == 'BtoA' else self.opt.output_nc
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.Color_Input = opt.Color_Input
self.Color_Output = opt.Color_Output
### input A (label maps)
dir_A = '_A'
self.dir_A = os.path.join(opt.dataroot, opt.phase + dir_A)
self.A_paths = sorted(make_dataset(self.dir_A))
### input B (real images)
dir_B = '_B'
self.dir_B = os.path.join(opt.dataroot, opt.phase + dir_B)
self.B_paths = sorted(make_dataset(self.dir_B))
self.dataset_size = len(self.A_paths)
def make_align_img(dir_A, dir_B, dir_AB):
print("Align data floder creating!")
num = 0
imgs_A = make_dataset(dir_A)
imgs_B = make_dataset(dir_B)
for img_A in tqdm(imgs_A):
img_inner = get_inner_path(img_A, dir_A)
if os.path.join(dir_B, img_inner) in imgs_B:
photo_A = cv2.imread(img_A)
photo_B = cv2.imread(os.path.join(dir_B, img_inner))
if photo_A.shape == photo_B.shape:
photo_AB = np.concatenate([photo_A, photo_B], 1)
img_AB = os.path.join(dir_AB, img_inner)
if not os.path.isdir(os.path.split(img_AB)[0]):
os.makedirs(os.path.split(img_AB)[0])
cv2.imwrite(img_AB, photo_AB)
num += 1
print("Align data floder created! %d img was processed" % num)
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.dir_A = os.path.join(
opt.dataroot,
opt.data_type + 'A') # create a path '/path/to/data/trainA' : RGB
self.dir_B = os.path.join(
opt.dataroot, opt.data_type +
'B') # create a path '/path/to/data/trainB': Thermal
self.A_paths = sorted(make_dataset(
self.dir_A,
opt.max_dataset_size)) # load images from '/path/to/data/trainA'
self.B_paths = sorted(make_dataset(
self.dir_B,
opt.max_dataset_size)) # load images from '/path/to/data/trainB'
self.A_size = len(self.A_paths) # get the size of dataset A
self.B_size = len(self.B_paths) # get the size of dataset B
btoA = self.opt.direction == 'BtoA'
input_nc = self.opt.output_nc if btoA else self.opt.input_nc # get the number of channels of input image
output_nc = self.opt.input_nc if btoA else self.opt.output_nc # get the number of channels of output image
self.transform_A = get_transform(self.opt, grayscale=(input_nc == 1))
self.transform_B = get_transform(self.opt, grayscale=(output_nc == 1))
if 'M' in opt.loss_type:
if opt.data_type == 'train':
gps_name = 'AM09_gps.txt'
else:
gps_name = 'AM05_gps.txt'
with open(os.path.join('./datasets', gps_name), 'r') as fp:
gps = [[
float(x.rstrip().split()[0]),
float(x.rstrip().split()[1])
] for x in fp.readlines()]
gps = np.array(gps)
gps = np.float32(gps)
import sklearn.metrics as skm
Distance = skm.pairwise_distances(gps, gps)
D = (Distance > opt.mt_neg)
self.DM = D
self.DV = Distance
else:
self.DM = None
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
### input A (label maps)
# dir_A = '_A' if self.opt.label_nc == 0 else '_label'
# self.dir_A = os.path.join(opt.dataroot, opt.phase + dir_A) # opt.phase can control the train set or test set
# self.A_paths = sorted(make_dataset(self.dir_A))
# dir_A = 'maps/train/' if self.opt.label_nc == 0 else '_label'
dir_A = 'images/train/' if self.opt.label_nc == 0 else '_label' # in our task, sourece domain A is the Image
self.dir_A = os.path.join(opt.dataroot, dir_A)
self.A_paths = sorted(make_dataset(self.dir_A))
### input B (real images)
# if opt.isTrain or opt.use_encoded_image:
# dir_B = '_B' if self.opt.label_nc == 0 else '_img'
# self.dir_B = os.path.join(opt.dataroot, opt.phase + dir_B) # opt.phase can control the train set or test set
# self.B_paths = sorted(make_dataset(self.dir_B))
if opt.isTrain or opt.use_encoded_image:
# dir_B = 'images/train/' if self.opt.label_nc == 0 else '_img'
dir_B = 'maps/train/' if self.opt.label_nc == 0 else '_img' # in our task, sourece domain B is the Saliency Map
self.dir_B = os.path.join(opt.dataroot, dir_B)
self.B_paths = sorted(make_dataset(self.dir_B))
dir_C = 'fixations_img/train/' if self.opt.label_nc == 0 else '_img' # in our task, sourece domain B is the Saliency Map
self.dir_C = os.path.join(opt.dataroot, dir_C)
self.C_paths = sorted(make_dataset(self.dir_C))
### instance maps
if not opt.no_instance:
self.dir_inst = os.path.join(opt.dataroot, opt.phase + '_inst')
self.inst_paths = sorted(make_dataset(self.dir_inst))
### load precomputed instance-wise encoded features
if opt.load_features:
self.dir_feat = os.path.join(opt.dataroot, opt.phase + '_feat')
print('----------- loading features from %s ----------' %
self.dir_feat)
self.feat_paths = sorted(make_dataset(self.dir_feat))
self.dataset_size = len(self.A_paths)
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_AB = os.path.join(opt.dataroot, opt.phase)
self.AB_paths = sorted(make_dataset(
self.dir_AB)) #,key=lambda x: int(os.path.splitext(x)[0]))
self.AB_paths.sort(key=lambda f: int(filter(str.isdigit, f)))
assert (opt.resize_or_crop == 'resize_and_crop')
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_AB = os.path.join(opt.dataroot, opt.phase)
self.dir_AB = self.dir_AB + '_temp'
self.temp_path = self.dir_AB[:-9] + 'template/template.jpg'
self.AB_paths = sorted(make_dataset(self.dir_AB))
self.size = len(self.AB_paths)
#self.labels={}
assert (opt.resize_or_crop == 'resize_and_crop')
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.fineSize = opt.fineSize
self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')
self.A_paths = make_dataset(self.dir_A)
self.A_paths = sorted(self.A_paths)
self.A_size = len(self.A_paths)
self.transform = get_transform(opt)
self.nintrm = opt.nintrm
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.dir_canonical = os.path.join(
"/home/robot/andrewk/RobotTeleop/pilot_scripts/canonical_no_shift_priv_info",
opt.phase + 'canonical')
self.dir_random = os.path.join(opt.dataroot, opt.phase + 'random')
#self.dir_real = os.path.join(opt.dataroot, opt.phase + 'real')
self.dir_seg = os.path.join(opt.dataroot, opt.phase + 'segmentation')
self.dir_depth = os.path.join(opt.dataroot, opt.phase + 'depth')
#self.real_states = np.load(os.path.join(opt.dataroot, opt.phase + 'real_states.npy'))
#self.dir_canonical_pi = '/home/robot/andrewk/RobotTeleop/pilot_scripts/priv_info' # np.load(os.path.join(opt.dataroot, opt.phase + '_canonical_pi.npy'))
self.dir_canonical_pi = '/home/robot/andrewk/RobotTeleop/pilot_scripts/canonical_no_shift_priv_info'
self.canonical_pi_paths = [
pth for pth in os.listdir(self.dir_canonical_pi) if '.npy' in pth
]
self.canonical_paths = sorted(
make_dataset(self.dir_canonical, opt.max_dataset_size))[::-1]
self.random_paths = sorted(
make_dataset(self.dir_random, opt.max_dataset_size))[::-1]
#self.real_paths = sorted(make_dataset(self.dir_#real, opt.max_dataset_size))
self.seg_paths = sorted(
make_dataset(self.dir_seg, opt.max_dataset_size))[::-1]
self.depth_paths = sorted(
make_dataset(self.dir_depth, opt.max_dataset_size))[::-1]
self.canonical_size = len(self.canonical_paths)
self.random_size = len(self.random_paths)
#self.real_size = len(self.real_paths)
self.seg_size = len(self.seg_paths)
self.depth_size = len(self.depth_paths)
self.index = 0
#assert self.canonical_size == self.random_size == self.seg_size == self.depth_size, 'Dataset sizes are not the same'
self.transform_rgb = get_transform(self.opt, grayscale=False)
self.transform_grayscale = get_transform(self.opt, grayscale=True)
def initialize(self,
opt,
dataroot=None,
image_dir=None,
label_dir=None,
record_txt=None,
transform=None,
is_aug=False):
assert not None in [dataroot,image_dir,label_dir],\
'dataroot:%s \nimage_dir:%s\nlabel_dir:%s'%(dataroot,image_dir,label_dir)
self.transform = transform
self.is_aug = is_aug
self.input_nc = opt.input_nc
image_folder = os.path.join(dataroot, image_dir)
label_folder = os.path.join(dataroot, label_dir)
self.image_paths = make_dataset(image_folder, record_txt)
self.label_paths = make_dataset(label_folder, record_txt)
def __init__(self, opt, augment):
self.opt = opt
self.root = opt.dataroot
self.transforms = augment
# query label (label maps)
dir_query_label = '_query_label'
self.dir_query_label = os.path.join(opt.dataroot,
opt.phase + dir_query_label)
self.query_label_paths = sorted(make_dataset(self.dir_query_label))
# ref label (label images)
dir_ref_label = '_ref_label'
self.dir_ref_label = os.path.join(opt.dataroot,
opt.phase + dir_ref_label)
self.ref_label_paths = sorted(make_dataset(self.dir_ref_label))
# query img (RGB maps)
dir_query_img = '_query_img'
self.dir_query_img = os.path.join(opt.dataroot,
opt.phase + dir_query_img)
self.query_img_paths = sorted(make_dataset(self.dir_query_img))
# ref img (RGB images)
dir_ref_img = '_ref_img'
self.dir_ref_img = os.path.join(opt.dataroot, opt.phase + dir_ref_img)
self.ref_img_paths = sorted(make_dataset(self.dir_ref_img))
if self.opt.shape_generation:
# densepose maps
dir_densepose = '_densepose'
self.dir_densepose = os.path.join(opt.dataroot,
opt.phase + dir_densepose)
self.densepose_paths = sorted(glob(self.dir_densepose + '/*'))
if self.opt.appearance_generation:
# generated segmentation from shape_generation (label maps)
dir_query_ref_label = '_query_ref_label'
self.dir_query_ref_label = os.path.join(
opt.dataroot, opt.phase + dir_query_ref_label)
self.query_ref_label_paths = sorted(
make_dataset(self.dir_query_ref_label))
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
p = os.path.join(opt.dataroot, 'images',
'train' if opt.isTrain else 'test')
self.A_paths = sorted(make_dataset(p, opt.max_dataset_size))
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
# get the images path
self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')
self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B')
# put all paths of images in a list
self.A_paths = make_dataset(self.dir_A)
self.B_paths = make_dataset(self.dir_B)
self.A_paths = sorted(self.A_paths)
self.B_paths = sorted(self.B_paths)
self.A_size = len(self.A_paths)
self.B_size = len(self.B_paths)
self.transform = get_transform(opt)
def initialize(self, opt):
if opt.nc == 1:
opt.gray = True
self.opt = opt
self.root = opt.dataroot
#self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')
#self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B')
self.dir_A = os.path.join(opt.dataroot, opt.phase)
self.dir_B = self.dir_A
self.A_paths = make_dataset(self.dir_A)
self.B_paths = make_dataset(self.dir_B)
self.A_paths = sorted(self.A_paths)
self.B_paths = sorted(self.B_paths)
self.A_size = len(self.A_paths)
self.B_size = len(self.B_paths)
self.transform = get_transform(opt)
self.transformFM = get_transformFMLoss(opt)
def __init__(self, opt):
super(UnalignedDataset, self).__init__()
self.opt = opt
self.transform = get_transform(opt)
datapath = os.path.join(opt.dataroot, opt.phase + '*')
self.dirs = sorted(glob.glob(datapath))
self.paths = [sorted(make_dataset(d)) for d in self.dirs]
self.sizes = [len(p) for p in self.paths]
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.A_paths = sorted(make_dataset(opt.dataroot, opt.max_dataset_size))
input_nc = self.opt.output_nc if self.opt.direction == 'BtoA' else self.opt.input_nc
self.transform = get_transform(opt, input_nc == 1)
def get_paths(self, opt):
"""THIS IS WHERE WE FIND THE IMAGES"""
# TODO modify this so we can put the images in normal places from the
# server...
root = opt.dataroot
phase = 'val'
label_dir = os.path.join(root, '%s_label' % phase)
label_paths = make_dataset(label_dir, recursive=False, read_cache=True)
return label_paths
def __init__(self, opt):
"""
初始化类，保存类中的选项，opt是BaseOptions的一个子类
"""
BaseDataset.__init__(self, opt)
self.dir_AB = os.path.join(opt.dataroot, opt.phase)
self.AB_paths = sorted(make_dataset(self.dir_AB, opt.max_dataset_size))
assert (self.opt.load_size >= self.opt.crop_size)
self.input_nc = self.opt.output_nc if self.opt.direction == 'BtoA' else self.opt.input_nc
self.output_nc = self.opt.input_nc if self.opt.direction == 'BtoA' else self.opt.output_nc
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot)
self.A_paths = make_dataset(self.dir_A)
self.A_paths = sorted(self.A_paths)
self.transform = get_transform(opt)
def initialize(self, opt):
self.opt = opt
self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')
self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B')
self.A_paths = make_dataset(self.dir_A)
self.B_paths = make_dataset(self.dir_B)
self.A_paths = sorted(self.A_paths)
self.B_paths = sorted(self.B_paths)
self.A_size = len(self.A_paths)
self.B_size = len(self.B_paths)
assert (self.A_size == self.B_size)
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot, 'rgb')
self.dir_depth_A = os.path.join(opt.dataroot, 'depth')
self.dir_label_A = os.path.join(opt.dataroot, 'seg_labels')
self.dir_edge_A = os.path.join(opt.dataroot,'edge_labels')
self.A_paths = make_dataset(self.dir_A)
self.A_depth_paths = make_dataset(self.dir_depth_A)
self.A_label_paths = make_dataset(self.dir_label_A)
self.A_edge_paths = make_dataset(self.dir_edge_A)
self.A_paths = sorted(self.A_paths)
self.A_depth_paths = sorted(self.A_depth_paths)
self.A_label_paths = sorted(self.A_label_paths)
self.A_edge_paths = sorted(self.A_edge_paths)
self.A_size = len(self.A_paths)
def __init__(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot, opt.phase, opt.folder, 'imgs')
self.A_paths = make_dataset(self.dir_A)
self.A_paths = sorted(self.A_paths)
self.A_size = len(self.A_paths)
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_val = os.path.join(opt.dataroot, 'val')
self.val_path = make_dataset(self.dir_val)
self.val_paths = sorted(self.val_path)
self.val_size = len(self.val_path)
self.A_path = make_dataset(os.path.join(opt.dataroot, opt.phase + 'A'))
self.A_path = sorted(self.A_path)
self.A_size = len(self.A_path)
self.B_path = make_dataset(os.path.join(opt.dataroot, opt.phase + 'B'))
self.B_path = sorted(self.B_path)
self.B_size = len(self.B_path)
self.transform = get_transform(opt)
self.target_transform = get_target_transform(opt)
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot)
self.A_paths = make_dataset(self.dir_A)
self.A_paths = sorted(self.A_paths)
self.transform = get_transform(opt)
def __init__(self,
root,
rgb_transforms=[],
rgb_torchvision_transforms=[],
gt_transforms=[]):
self.root = root
self.dir_rgb = os.path.join(root, 'RGB')
self.dir_gt = os.path.join(root, 'GT')
self.rgb_paths = make_dataset(self.dir_rgb)
self.gt_paths = make_dataset(self.dir_gt)
self.rgb_paths = sorted(self.rgb_paths)
self.gt_paths = sorted(self.gt_paths)
self.rgb_transforms = rgb_transforms
self.gt_transforms = gt_transforms
self.rgb_torchvision_transforms = rgb_torchvision_transforms
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A') # create a path '/path/to/data/trainA'
self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B') # create a path '/path/to/data/trainB'
self.A_paths = sorted(make_dataset(self.dir_A, opt.max_dataset_size)) # load images from '/path/to/data/trainA'
self.B_paths = sorted(make_dataset(self.dir_B, opt.max_dataset_size)) # load images from '/path/to/data/trainB'
self.A_size = len(self.A_paths) # get the size of dataset A
self.B_size = len(self.B_paths) # get the size of dataset B
btoA = self.opt.direction == 'BtoA'
input_nc = self.opt.output_nc if btoA else self.opt.input_nc # get the number of channels of input image
output_nc = self.opt.input_nc if btoA else self.opt.output_nc # get the number of channels of output image
self.transform_A = get_transform(opt, grayscale=(input_nc == 1)) # if nc == 1, we convert RGB to grayscale image
self.transform_B = get_transform(opt, grayscale=(output_nc == 1)) # if nc == 1, we convert RGB to grayscale image
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
### input A (label maps)
self.img_paths = sorted(make_dataset(self.root))
self.A_paths = [img_name for img_name in self.img_paths if '_a.jpg' in img_name]
self.B_paths = [img_name for img_name in self.img_paths if '_b.jpg' in img_name]
self.dataset_size = len(self.A_paths)
def initialize(self, dataroot, load_size=64):
self.root = dataroot
self.load_size = load_size
self.dir_p0 = os.path.join(self.root, 'p0')
self.p0_paths = make_dataset(self.dir_p0)
self.p0_paths = sorted(self.p0_paths)
self.dir_p1 = os.path.join(self.root, 'p1')
self.p1_paths = make_dataset(self.dir_p1)
self.p1_paths = sorted(self.p1_paths)
transform_list = []
transform_list.append(transforms.Scale(load_size))
transform_list += [transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),(0.5, 0.5, 0.5))]
self.transform = transforms.Compose(transform_list)
# judgement directory
self.dir_S = os.path.join(self.root, 'same')
self.same_paths = make_dataset(self.dir_S,mode='np')
self.same_paths = sorted(self.same_paths)
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_AB = os.path.join(opt.dataroot, opt.phase)
self.AB_paths = sorted(make_dataset(self.dir_AB))
assert(opt.resize_or_crop == 'resize_and_crop')
transform_list = [transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5))]
self.transform = transforms.Compose(transform_list)
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_AB = os.path.join(opt.dataroot, opt.phase)
self.AB_paths = sorted(make_dataset(self.dir_AB))
assert(opt.resize_or_crop == 'resize_and_crop')
