def get_test_dataloader(config, train=False):
data_dir = osp.join('..', 'data', config.dataset)
stats_filename = osp.join(data_dir, config.scene, 'stats.txt')
stats = np.loadtxt(stats_filename)
# transformer
data_transform = transforms.Compose([
transforms.Resize(config.image_size),
transforms.ToTensor(),
transforms.Normalize(mean=stats[0], std=np.sqrt(stats[1]))
])
target_transform = transforms.Lambda(lambda x: torch.from_numpy(x).float())
kwargs = dict(scene=config.scene,
data_path=config.dataset_path,
train=train,
transform=data_transform,
target_transform=target_transform,
seed=config.seed,
data_dir=config.preprocessed_data_path,
config=config)
if config.dataset == '7Scenes':
dataset = SevenScenes(**kwargs)
elif config.dataset == 'RobotCar':
dataset = RobotCar(**kwargs)
elif config.dataset == 'SenseTime':
dataset = SenseTime(**kwargs)
elif config.dataset == 'Cambridge' or config.dataset == 'NewCambridge':
dataset = Cambridge(**kwargs)
loader = DataLoader(dataset,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.num_workers,
pin_memory=True)
return loader
def get_posenet_train_dataloader(config):
data_transform, target_transform = get_train_transforms(config)
kwargs = dict(
scene=config.scene,
data_path=config.dataset_path,
transform=data_transform,
target_transform=target_transform,
seed=config.seed,
data_dir=config.preprocessed_data_path,
config=config
)
if config.dataset == '7Scenes':
train_data = SevenScenes(train=True, **kwargs)
valid_data = SevenScenes(train=False, **kwargs)
elif config.dataset == 'RobotCar':
train_data = RobotCar(train=True, **kwargs)
valid_data = RobotCar(train=False, **kwargs)
elif config.dataset == 'SenseTime':
train_data = SenseTime(train=True, **kwargs)
valid_data = SenseTime(train=False, **kwargs)
elif config.dataset == 'Cambridge' or config.dataset == 'NewCambridge':
train_data = Cambridge(train=True, **kwargs)
valid_data = Cambridge(train=False, **kwargs)
else:
raise NotImplementedError
dataloader_kwargs = dict(
batch_size=config.batch_size,
shuffle=config.shuffle,
num_workers=config.num_workers,
pin_memory=True,
collate_fn=safe_collate
)
train_dataloader = torch.utils.data.DataLoader(
train_data,
**dataloader_kwargs
)
if config.do_val:
val_dataloader = torch.utils.data.DataLoader(
valid_data,
**dataloader_kwargs
)
else:
val_dataloader = None
return train_dataloader, val_dataloader
def ResizeImages(train=True):
data_transform = transforms.Compose([
transforms.Resize(image_size),
transforms.Lambda(lambda x: np.asarray(x))
])
dset = Cambridge(scene=args.scene,
data_path=args.data_dir,
train=train,
transform=data_transform,
config=config,
data_dir=config.data_dir)
loader = DataLoader(dset,
batch_size=batch_size,
num_workers=num_workers,
collate_fn=safe_collate,
shuffle=False)
base_dir = osp.join(args.data_dir, args.scene)
if train:
split_filename = osp.join(base_dir, 'dataset_train.txt')
else:
split_filename = osp.join(base_dir, 'dataset_test.txt')
with open(split_filename, 'r') as f:
data = f.readlines()
im_fns = [
osp.join(base_dir,
l.split(' ')[0]) for l in data[3:]
if not l.startswith('#')
]
assert len(dset) == len(im_fns)
for batch_idx, (imgs, _) in enumerate(loader):
for idx, im in enumerate(imgs):
im_fn = im_fns[batch_idx * batch_size + idx]
im = Image.fromarray(im.numpy())
try:
im.save(im_fn)
except IOError:
print('IOError while saving {:s}'.format(im_fn))
if batch_idx % 50 == 0:
print('Processed {:d}/{:d}'.format(batch_idx * batch_size,
len(dset)))
ResizeImages(args.train)
# In[6]:
# get statistical information from training set only
data_transform = transforms.Compose(
[transforms.Resize(256), transforms.ToTensor()])
kwargs = dict(scene=args.scene,
data_path=args.data_dir,
train=True,
transform=data_transform,
config=config,
data_dir=config.data_dir)
dset = Cambridge(**kwargs)
loader = DataLoader(dset,
batch_size=batch_size,
num_workers=num_workers,
collate_fn=safe_collate)
acc = np.zeros((3, image_size[0], image_size[1]))
sq_acc = np.zeros((3, image_size[0], image_size[1]))
for batch_idx, (imgs, _) in enumerate(loader):
imgs = imgs.numpy()
acc += np.sum(imgs, axis=0)
sq_acc += np.sum(imgs**2, axis=0)
if batch_idx % 50 == 0:
print('Accumulated {:d}/{:d}'.format(batch_idx * batch_size,
len(dset)))