def main():
args = parser.parse_args()
print('->model load: {}'.format(args.pretrained_posenet))
weights_pose = torch.load(args.pretrained_posenet)
pose_net = models.PoseNet().to(device)
pose_net.load_state_dict(weights_pose['state_dict'], strict=False)
pose_net.eval()
if args.type =='gray':
image_dir = Path(args.dataset_dir + args.sequence + "/image_1/")#gray 01, color 23
else:
image_dir = Path(args.dataset_dir + args.sequence + "/image_2/")#gray 01, color 23
output_dir = Path(args.output_dir)
print('-> out dif {}'.format(output_dir))
output_dir.makedirs_p()
test_files = sum([image_dir.files('*.{}'.format(ext))
for ext in args.img_exts], [])
test_files.sort()
print('{} files to test'.format(len(test_files)))
global_pose = np.identity(4)
poses = [global_pose[0:3, :].reshape(1, 12)]
n = len(test_files)
tensor_img1 = load_tensor_image(test_files[0], args)
for iter in tqdm(range(n - 1)):
tensor_img2 = load_tensor_image(test_files[iter+1], args)
pose = pose_net(tensor_img1, tensor_img2)
pose_mat = pose_vec2mat(pose).squeeze(0).cpu().numpy()#1,6-->3x4
pose_mat = np.vstack([pose_mat, np.array([0, 0, 0, 1])])#4X4
global_pose = global_pose @ np.linalg.inv(pose_mat)
pose = global_pose[0:3, :].reshape(1, 12)
poses.append(pose)
# update
tensor_img1 = tensor_img2
poses = np.concatenate(poses, axis=0)
if args.scale_factor:
poses[:,3]*=args.scale_factor#x-axis
poses[:,11]*=args.scale_factor#z-axis
filename = Path(args.output_dir + args.sequence + ".txt")
np.savetxt(filename, poses, delimiter=' ', fmt='%1.8e')
def main():
args = parser.parse_args()
weights_pose = torch.load(args.pretrained_posenet)
pose_net = models.PoseNet().to(device)
pose_net.load_state_dict(weights_pose['state_dict'], strict=False)
pose_net.eval()
sequences = os.listdir(args.dataset_dir)
for seq in sequences:
if '.txt' not in seq:
args.sequence = seq
image_dir = Path(args.dataset_dir + args.sequence + "/")
output_dir = Path(args.output_dir)
output_dir.makedirs_p()
test_files = sum(
[image_dir.files('*.{}'.format(ext)) for ext in args.img_exts],
[])
test_files.sort()
print('{} files to test'.format(len(test_files)))
global_pose = np.identity(4)
poses = [global_pose[0:3, :].reshape(1, 12)]
n = len(test_files)
tensor_img1 = load_tensor_image(test_files[0], args)
for iter in tqdm(range(n - 1)):
tensor_img2 = load_tensor_image(test_files[iter + 1], args)
pose = pose_net(tensor_img1, tensor_img2)
pose_mat = pose_vec2mat(pose).squeeze(0).cpu().numpy()
pose_mat = np.vstack([pose_mat, np.array([0, 0, 0, 1])])
global_pose = global_pose @ np.linalg.inv(pose_mat)
poses.append(global_pose[0:3, :].reshape(1, 12))
# update
tensor_img1 = tensor_img2
poses = np.concatenate(poses, axis=0)
filename = Path(args.output_dir + args.sequence + ".txt")
np.savetxt(filename, poses, delimiter=' ', fmt='%1.8e')
def main():
args = parser.parse_args()
weights = torch.load(args.pretrained_posenet)
pose_net = models.PoseNet().to(device)
pose_net.load_state_dict(weights['state_dict'], strict=False)
pose_net.eval()
seq_length = 5
dataset_dir = Path(args.dataset_dir)
framework = test_framework(dataset_dir, args.sequences, seq_length)
print('{} snippets to test'.format(len(framework)))
errors = np.zeros((len(framework), 2), np.float32)
if args.output_dir is not None:
output_dir = Path(args.output_dir)
output_dir.makedirs_p()
predictions_array = np.zeros((len(framework), seq_length, 3, 4))
for j, sample in enumerate(tqdm(framework)):
imgs = sample['imgs']
h, w, _ = imgs[0].shape
if (not args.no_resize) and (h != args.img_height
or w != args.img_width):
imgs = [
imresize(img,
(args.img_height, args.img_width)).astype(np.float32)
for img in imgs
]
imgs = [np.transpose(img, (2, 0, 1)) for img in imgs]
tensor_imgs = []
for i, img in enumerate(imgs):
img = ((torch.from_numpy(img).unsqueeze(0) / 255 - 0.5) /
0.5).to(device)
tensor_imgs.append(img)
global_pose = np.identity(4)
poses = []
poses.append(global_pose[0:3, :])
for iter in range(seq_length - 1):
pose = pose_net(tensor_imgs[iter], tensor_imgs[iter + 1])
pose_mat = pose_vec2mat(pose).squeeze(0).cpu().numpy()
pose_mat = np.vstack([pose_mat, np.array([0, 0, 0, 1])])
global_pose = global_pose @ np.linalg.inv(pose_mat)
poses.append(global_pose[0:3, :])
final_poses = np.stack(poses, axis=0)
if args.output_dir is not None:
predictions_array[j] = final_poses
ATE, RE = compute_pose_error(sample['poses'], final_poses)
errors[j] = ATE, RE
mean_errors = errors.mean(0)
std_errors = errors.std(0)
error_names = ['ATE', 'RE']
print('')
print("Results")
print("\t {:>10}, {:>10}".format(*error_names))
print("mean \t {:10.4f}, {:10.4f}".format(*mean_errors))
print("std \t {:10.4f}, {:10.4f}".format(*std_errors))
if args.output_dir is not None:
np.save(output_dir / 'predictions.npy', predictions_array)
def prepare_environment():
env = {}
args = parser.parse_args()
if args.dataset_format == 'KITTI':
from datasets.shifted_sequence_folders import ShiftedSequenceFolder
elif args.dataset_format == 'StillBox':
from datasets.shifted_sequence_folders import StillBox as ShiftedSequenceFolder
elif args.dataset_format == 'TUM':
from datasets.shifted_sequence_folders import TUM as ShiftedSequenceFolder
save_path = save_path_formatter(args, parser)
args.save_path = 'checkpoints' / save_path
print('=> will save everything to {}'.format(args.save_path))
args.save_path.makedirs_p()
torch.manual_seed(args.seed)
args.test_batch_size = 4 * args.batch_size
if args.evaluate:
args.epochs = 0
env['training_writer'] = SummaryWriter(args.save_path)
output_writers = []
if args.log_output:
for i in range(3):
output_writers.append(
SummaryWriter(args.save_path / 'valid' / str(i)))
env['output_writers'] = output_writers
# Data loading code
normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
train_transform = custom_transforms.Compose([
# custom_transforms.RandomHorizontalFlip(),
custom_transforms.ArrayToTensor(),
normalize
])
valid_transform = custom_transforms.Compose(
[custom_transforms.ArrayToTensor(), normalize])
print("=> fetching scenes in '{}'".format(args.data))
train_set = ShiftedSequenceFolder(args.data,
transform=train_transform,
seed=args.seed,
train=True,
with_depth_gt=False,
with_pose_gt=args.supervise_pose,
sequence_length=args.sequence_length)
val_set = ShiftedSequenceFolder(args.data,
transform=valid_transform,
seed=args.seed,
train=False,
sequence_length=args.sequence_length,
with_depth_gt=args.with_gt,
with_pose_gt=args.with_gt)
print('{} samples found in {} train scenes'.format(len(train_set),
len(train_set.scenes)))
print('{} samples found in {} valid scenes'.format(len(val_set),
len(val_set.scenes)))
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=4 * args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
env['train_set'] = train_set
env['val_set'] = val_set
env['train_loader'] = train_loader
env['val_loader'] = val_loader
if args.epoch_size == 0:
args.epoch_size = len(train_loader)
# create model
print("=> creating model")
pose_net = models.PoseNet(seq_length=args.sequence_length,
batch_norm=args.bn in ['pose',
'both']).to(device)
if args.pretrained_pose:
print("=> using pre-trained weights for pose net")
weights = torch.load(args.pretrained_pose)
pose_net.load_state_dict(weights['state_dict'], strict=False)
depth_net = models.DepthNet(depth_activation="elu",
batch_norm=args.bn in ['depth',
'both']).to(device)
if args.pretrained_depth:
print("=> using pre-trained DepthNet model")
data = torch.load(args.pretrained_depth)
depth_net.load_state_dict(data['state_dict'])
cudnn.benchmark = True
depth_net = torch.nn.DataParallel(depth_net)
pose_net = torch.nn.DataParallel(pose_net)
env['depth_net'] = depth_net
env['pose_net'] = pose_net
print('=> setting adam solver')
optim_params = [{
'params': depth_net.parameters(),
'lr': args.lr
}, {
'params': pose_net.parameters(),
'lr': args.lr
}]
# parameters = chain(depth_net.parameters(), pose_exp_net.parameters())
optimizer = torch.optim.Adam(optim_params,
betas=(args.momentum, args.beta),
weight_decay=args.weight_decay)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
args.lr_decay_frequency,
gamma=0.5)
env['optimizer'] = optimizer
env['scheduler'] = scheduler
with open(args.save_path / args.log_summary, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(['train_loss', 'validation_loss'])
with open(args.save_path / args.log_full, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(
['train_loss', 'photo_loss', 'explainability_loss', 'smooth_loss'])
logger = TermLogger(n_epochs=args.epochs,
train_size=min(len(train_loader), args.epoch_size),
valid_size=len(val_loader))
logger.epoch_bar.start()
env['logger'] = logger
env['args'] = args
return env
def main():
global best_error, n_iter, device
args = parser.parse_args()
if args.dataset_format == 'stacked':
from datasets.stacked_sequence_folders import SequenceFolder
elif args.dataset_format == 'sequential':
from datasets.sequence_folders import SequenceFolder
timestamp = datetime.datetime.now().strftime("%m-%d-%H:%M")
args.save_path = 'checkpoints' / Path(args.name) / timestamp
print('=> will save everything to {}'.format(args.save_path))
args.save_path.makedirs_p()
torch.manual_seed(args.seed)
training_writer = SummaryWriter(args.save_path)
# Data loading
normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
train_transform = custom_transforms.Compose([
custom_transforms.RandomHorizontalFlip(),
custom_transforms.RandomScaleCrop(),
custom_transforms.ArrayToTensor(), normalize
])
valid_transform = custom_transforms.Compose(
[custom_transforms.ArrayToTensor(), normalize])
print("=> fetching scenes in '{}'".format(args.data))
train_set = SequenceFolder(args.data,
transform=train_transform,
seed=args.seed,
train=True,
sequence_length=args.sequence_length)
# if no Groundtruth is avalaible, Validation set is the same type as training set to measure photometric loss from warping
if args.with_gt:
from datasets.validation_folders import ValidationSet
val_set = ValidationSet(args.data, transform=valid_transform)
else:
val_set = SequenceFolder(
args.data,
transform=valid_transform,
seed=args.seed,
train=False,
sequence_length=args.sequence_length,
)
print('{} samples found in {} train scenes'.format(len(train_set),
len(train_set.scenes)))
print('{} samples found in {} valid scenes'.format(len(val_set),
len(val_set.scenes)))
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.epoch_size == 0:
args.epoch_size = len(train_loader)
# create model
print("=> creating model")
disp_net = getattr(models, args.dispnet)().to(device)
pose_net = models.PoseNet().to(device)
if args.pretrained_pose:
print("=> using pre-trained weights for PoseNet")
weights = torch.load(args.pretrained_pose)
pose_net.load_state_dict(weights['state_dict'], strict=False)
else:
pose_net.init_weights()
if args.pretrained_disp:
print("=> using pre-trained weights for DispNet")
weights = torch.load(args.pretrained_disp)
disp_net.load_state_dict(weights['state_dict'])
else:
disp_net.init_weights()
cudnn.benchmark = True
disp_net = torch.nn.DataParallel(disp_net)
pose_net = torch.nn.DataParallel(pose_net)
print('=> setting adam solver')
optim_params = [{
'params': disp_net.parameters(),
'lr': args.lr
}, {
'params': pose_net.parameters(),
'lr': args.lr
}]
optimizer = torch.optim.Adam(optim_params,
betas=(args.momentum, args.beta),
weight_decay=args.weight_decay)
with open(args.save_path / args.log_summary, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(['train_loss', 'validation_loss'])
with open(args.save_path / args.log_full, 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(
['photo_loss', 'smooth_loss', 'geometry_loss', 'train_loss'])
logger = TermLogger(n_epochs=args.epochs,
train_size=min(len(train_loader), args.epoch_size),
valid_size=len(val_loader))
logger.epoch_bar.start()
if args.pretrained_disp:
logger.reset_valid_bar()
if args.with_gt:
errors, error_names = validate_with_gt(args, val_loader, disp_net,
0, logger)
else:
errors, error_names = validate_without_gt(args, val_loader,
disp_net, pose_net, 0,
logger)
for error, name in zip(errors, error_names):
training_writer.add_scalar(name, error, 0)
error_string = ', '.join(
'{} : {:.3f}'.format(name, error)
for name, error in zip(error_names[2:9], errors[2:9]))
logger.valid_writer.write(' * Avg {}'.format(error_string))
for epoch in range(args.epochs):
logger.epoch_bar.update(epoch)
# train for one epoch
logger.reset_train_bar()
train_loss = train(args, train_loader, disp_net, pose_net, optimizer,
args.epoch_size, logger, training_writer)
logger.train_writer.write(' * Avg Loss : {:.3f}'.format(train_loss))
# evaluate on validation set
logger.reset_valid_bar()
if args.with_gt:
errors, error_names = validate_with_gt(args, val_loader, disp_net,
epoch, logger)
else:
errors, error_names = validate_without_gt(args, val_loader,
disp_net, pose_net,
epoch, logger)
error_string = ', '.join('{} : {:.3f}'.format(name, error)
for name, error in zip(error_names, errors))
logger.valid_writer.write(' * Avg {}'.format(error_string))
for error, name in zip(errors, error_names):
training_writer.add_scalar(name, error, epoch)
# Up to you to chose the most relevant error to measure your model's performance, careful some measures are to maximize (such as a1,a2,a3)
decisive_error = errors[1]
if best_error < 0:
best_error = decisive_error
# remember lowest error and save checkpoint
is_best = decisive_error < best_error
best_error = min(best_error, decisive_error)
save_checkpoint(args.save_path, {
'epoch': epoch + 1,
'state_dict': disp_net.module.state_dict()
}, {
'epoch': epoch + 1,
'state_dict': pose_net.module.state_dict()
}, is_best)
with open(args.save_path / args.log_summary, 'a') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow([train_loss, decisive_error])
logger.epoch_bar.finish()
def load_model(self, path):
self.pose_net = models.PoseNet().to(device)
weights_pose = torch.load(path)
self.pose_net.load_state_dict(weights_pose['state_dict'], strict=False)
self.pose_net.eval()