def main():
args = parser.parse_args()
output_dir = Path(args.output_dir)
normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
valid_transform = custom_transforms.Compose(
[custom_transforms.ArrayToTensor(), normalize])
val_set = SequenceFolder(args.data,
transform=valid_transform,
seed=args.seed,
sequence_length=args.sequence_length)
print('{} samples found in {} valid scenes'.format(len(val_set),
len(val_set.scenes)))
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
dpsnet = PSNet(args.nlabel, args.mindepth).cuda()
weights = torch.load(args.pretrained_dps)
dpsnet.load_state_dict(weights['state_dict'])
dpsnet.eval()
output_dir = Path(args.output_dir)
if not os.path.isdir(output_dir):
os.mkdir(output_dir)
errors = np.zeros((2, 8, int(len(val_loader) / args.print_freq) + 1),
np.float32)
with torch.no_grad():
for ii, (tgt_img, ref_imgs, ref_poses, intrinsics, intrinsics_inv,
tgt_depth, scale_) in enumerate(val_loader):
if ii % args.print_freq == 0:
i = int(ii / args.print_freq)
tgt_img_var = Variable(tgt_img.cuda())
ref_imgs_var = [Variable(img.cuda()) for img in ref_imgs]
ref_poses_var = [Variable(pose.cuda()) for pose in ref_poses]
intrinsics_var = Variable(intrinsics.cuda())
intrinsics_inv_var = Variable(intrinsics_inv.cuda())
tgt_depth_var = Variable(tgt_depth.cuda())
scale = scale_.numpy()[0]
# compute output
pose = torch.cat(ref_poses_var, 1)
start = time.time()
output_depth = dpsnet(tgt_img_var, ref_imgs_var, pose,
intrinsics_var, intrinsics_inv_var)
elps = time.time() - start
mask = (tgt_depth <= args.maxdepth) & (
tgt_depth >= args.mindepth) & (tgt_depth == tgt_depth)
tgt_disp = args.mindepth * args.nlabel / tgt_depth
output_disp = args.mindepth * args.nlabel / output_depth
output_disp_ = torch.squeeze(output_disp.data.cpu(), 1)
output_depth_ = torch.squeeze(output_depth.data.cpu(), 1)
errors[0, :,
i] = compute_errors_test(tgt_depth[mask] / scale,
output_depth_[mask] / scale)
errors[1, :,
i] = compute_errors_test(tgt_disp[mask] / scale,
output_disp_[mask] / scale)
print('Elapsed Time {} Abs Error {:.4f}'.format(
elps, errors[0, 0, i]))
if args.output_print:
output_disp_n = (output_disp_).numpy()[0]
np.save(output_dir / '{:04d}{}'.format(i, '.npy'),
output_disp_n)
disp = (255 * tensor2array(torch.from_numpy(output_disp_n),
max_value=args.nlabel,
colormap='bone')).astype(
np.uint8)
imsave(output_dir / '{:04d}_disp{}'.format(i, '.png'),
disp)
mean_errors = errors.mean(2)
error_names = [
'abs_rel', 'abs_diff', 'sq_rel', 'rms', 'log_rms', 'a1', 'a2', 'a3'
]
print("{}".format(args.output_dir))
print("Depth Results : ")
print("{:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}".
format(*error_names))
print(
"{:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}"
.format(*mean_errors[0]))
print("Disparity Results : ")
print("{:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}".
format(*error_names))
print(
"{:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}"
.format(*mean_errors[1]))
np.savetxt(output_dir / 'errors.csv',
mean_errors,
fmt='%1.4f',
delimiter=',')
def main():
args = parser.parse_args()
output_dir = Path(args.output_dir)
normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
valid_transform = custom_transforms.Compose(
[custom_transforms.ArrayToTensor(), normalize])
val_set = ScannetDataset(args.data,
args.testlist,
mode='test',
n_frames=args.seq_len,
r=args.seq_len * 2,
transform=valid_transform)
print('{} samples found in {} valid scenes'.format(len(val_set),
len(val_set.scenes)))
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
dpsnet = PSNet(args.nlabel, args.mindepth).cuda()
weights = torch.load(args.pretrained_dps)
dpsnet.load_state_dict(weights['state_dict'])
dpsnet.eval()
output_dir = Path(args.output_dir)
if not os.path.isdir(output_dir):
os.mkdir(output_dir)
with torch.no_grad():
for ii, (tgt_img, ref_imgs, ref_poses, intrinsics, intrinsics_inv,
tgt_depth, scene, tgt_filename) in enumerate(val_loader):
tgt_img_var = Variable(tgt_img.cuda())
ref_imgs_var = [Variable(img.cuda()) for img in ref_imgs]
ref_poses_var = [Variable(pose.cuda()) for pose in ref_poses]
intrinsics_var = Variable(intrinsics.cuda())
intrinsics_inv_var = Variable(intrinsics_inv.cuda())
tgt_depth_var = Variable(tgt_depth.cuda())
# compute output
pose = torch.cat(ref_poses_var, 1)
start = time.time()
output_depth = dpsnet(tgt_img_var, ref_imgs_var, pose,
intrinsics_var, intrinsics_inv_var)
elps = time.time() - start
tgt_disp = args.mindepth * args.nlabel / tgt_depth
output_disp = args.mindepth * args.nlabel / output_depth
mask = (tgt_depth <= args.maxdepth) & (
tgt_depth >= args.mindepth) & (tgt_depth == tgt_depth)
output_disp_ = torch.squeeze(output_disp.data.cpu(), 1)
output_depth_ = torch.squeeze(output_depth.data.cpu(), 1)
for idx in range(tgt_img_var.shape[0]):
scene_name = scene[idx]
rgb_basename = tgt_filename[idx]
_, img_ext = os.path.splitext(rgb_basename)
pred_depth_dir = os.path.join(args.output_dir, scene_name,
"pred_depth")
if not os.path.exists(pred_depth_dir):
os.makedirs(pred_depth_dir)
pred_depth = output_depth[idx]
pred_depth = np.float16(pred_depth.squeeze(1).cpu().numpy())
pred_depth_filepath = os.path.join(
pred_depth_dir,
rgb_basename.replace("color" + img_ext, "pred_depth.npy"))
np.save(pred_depth_filepath, pred_depth)
pred_depth_color = colorize_depth(
pred_depth.squeeze(1),
max_depth=5.0).permute(0, 2, 3, 1).squeeze().cpu().numpy()
pred_depth_color_filepath = os.path.join(
pred_depth_dir,
rgb_basename.replace("color" + img_ext, "pred_depth.jpg"))
cv2.imwrite(
pred_depth_color_filepath,
cv2.cvtColor(np.uint8(pred_depth_color),
cv2.COLOR_RGB2BGR))
def main():
global n_iter
args = parser.parse_args()
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)
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)))
# 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.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,
ttype=args.ttype,
add_geo=args.geo,
depth_source=args.depth_init,
pose_source='%s_poses.txt' %
args.pose_init if args.pose_init else 'poses.txt',
scale=False)
val_set = SequenceFolder(args.data,
transform=valid_transform,
seed=args.seed,
ttype=args.ttype2,
add_geo=args.geo,
depth_source=args.depth_init,
pose_source='%s_poses.txt' %
args.pose_init if args.pose_init else 'poses.txt',
scale=False)
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)
if args.epoch_size == 0:
args.epoch_size = len(train_loader)
# create model
print("=> creating model")
depth_net = PSNet(args.nlabel,
args.mindepth,
add_geo_cost=args.geo,
depth_augment=False).cuda()
if args.pretrained_dps:
# for param in depth_net.feature_extraction.parameters():
# param.requires_grad = False
print("=> using pre-trained weights for DPSNet")
model_dict = depth_net.state_dict()
weights = torch.load(args.pretrained_dps)['state_dict']
pretrained_dict = {k: v for k, v in weights.items() if k in model_dict}
model_dict.update(pretrained_dict)
depth_net.load_state_dict(model_dict)
else:
depth_net.init_weights()
cudnn.benchmark = True
depth_net = torch.nn.DataParallel(depth_net)
print('=> setting adam solver')
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
depth_net.parameters()),
args.lr,
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(['train_loss'])
for epoch in range(args.epochs):
adjust_learning_rate(args, optimizer, epoch)
# train for one epoch
train_loss = train(args, train_loader, depth_net, optimizer,
args.epoch_size, training_writer)
save_checkpoint(args.save_path, {
'epoch': epoch + 1,
'state_dict': depth_net.module.state_dict()
}, epoch)
with open(args.save_path / args.log_summary, 'a') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow([train_loss])
def main():
args = parser.parse_args()
normalize = custom_transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
valid_transform = custom_transforms.Compose(
[custom_transforms.ArrayToTensor(), normalize])
val_set = SequenceFolder(args.data,
transform=valid_transform,
seed=args.seed,
ttype=args.ttype,
dataset='',
sequence_length=args.sequence_length,
add_geo=args.geo,
depth_source=args.depth_init,
pose_source='%s_poses.txt' %
args.pose_init if args.pose_init else 'poses.txt',
scale=False,
size=0,
req_gt=True,
get_path=True)
print('{} samples found in {} valid scenes'.format(len(val_set),
len(val_set.scenes)))
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
depth_net = PSNet(args.nlabel, args.mindepth, add_geo_cost=args.geo).cuda()
weights = torch.load(args.pretrained_dps)
depth_net.load_state_dict(weights['state_dict'])
depth_net.eval()
output_dir = Path(args.output_dir)
if not os.path.isdir(output_dir):
os.mkdir(output_dir)
errors = np.zeros((2, 13, int(np.ceil(len(val_loader) / args.print_freq))),
np.float32)
with torch.no_grad():
for ii, (tgt_img, ref_imgs, ref_poses, poses_gt, intrinsics,
intrinsics_inv, tgt_depth, ref_depths,
tgt_path) in enumerate(val_loader):
if ii % args.print_freq == 0:
i = int(ii / args.print_freq)
tgt_img_var = Variable(tgt_img.cuda())
ref_imgs_var = [Variable(img.cuda()) for img in ref_imgs]
ref_poses_var = [Variable(pose.cuda()) for pose in ref_poses]
poses_gt_var = [
Variable(pose_gt.cuda()) for pose_gt in poses_gt
]
ref_depths_var = [Variable(dep.cuda()) for dep in ref_depths]
intrinsics_var = Variable(intrinsics.cuda())
intrinsics_inv_var = Variable(intrinsics_inv.cuda())
# compute output
pose = torch.cat(ref_poses_var, 1)
poses_gt = torch.cat(poses_gt_var, 1)
rel_pose = poses_gt.squeeze().data.cpu().numpy()
scale = float(np.sqrt(rel_pose[:3, 3].dot(rel_pose[:3, 3])))
start = time.time()
output_depth = depth_net(tgt_img_var, ref_imgs_var, pose,
intrinsics_var, intrinsics_inv_var,
ref_depths_var)
elps = time.time() - start
tgt_disp = args.mindepth * args.nlabel / tgt_depth
output_disp = args.mindepth * args.nlabel / output_depth
mask = (tgt_depth <= args.maxdepth) & (tgt_depth >= 0.5) & (
tgt_depth == tgt_depth)
tgt_depth = tgt_depth / scale
output_depth_scaled = output_depth / scale
output_disp_ = torch.squeeze(output_disp.data.cpu(), 1)
output_depth_ = torch.squeeze(output_depth_scaled.data.cpu(),
1)
if args.save:
output_depth_n = torch.squeeze(output_depth.data.cpu(),
1).numpy()[0]
save_path = tgt_path[0][:-4] + "_" + args.save + ".npy"
if not os.path.exists(save_path):
np.save(save_path, output_depth_n)
errors[0, :10,
i] = compute_errors_test(tgt_depth[mask],
output_depth_[mask])
errors[1, :10,
i] = compute_errors_test(tgt_disp[mask],
output_disp_[mask])
print('iter{}, Elapsed Time {} Abs Error {:.10f}'.format(
i, elps, errors[0, 0, i]))
if args.output_print:
output_disp_n = (output_disp_).numpy()[0]
np.save(output_dir / '{:08d}{}'.format(i, '.npy'),
output_disp_n)
disp = (255 * tensor2array(torch.from_numpy(output_disp_n),
max_value=args.nlabel,
colormap='bone')).astype(
np.uint8)
disp = disp.transpose(1, 2, 0)
imsave(output_dir / '{:08d}_disp{}'.format(i, '.png'),
disp)
mean_errors = errors.mean(2)
error_names = [
'abs_rel', 'abs_diff', 'sq_rel', 'rms', 'log_rms', 'a1', 'a2', 'a3',
'L1-inv', "sc-inv", 'ra', 'rd', 'ta'
]
print("{}".format(args.output_dir))
print("Depth & angle Results : ")
print(
"{:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}"
.format(*error_names))
print(
"{:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, "
"{:10.4f}, {:10.4f}, {:10.4f}".format(*mean_errors[0]))
np.savetxt(output_dir / 'errors.csv',
mean_errors,
fmt='%1.4f',
delimiter=',')
def main():
global n_iter
args = parser.parse_args()
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)
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)))
# 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.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,
ttype=args.ttype)
val_set = SequenceFolder(args.data,
transform=valid_transform,
seed=args.seed,
ttype=args.ttype2)
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")
dpsnet = PSNet(args.nlabel, args.mindepth).cuda()
if args.pretrained_dps:
print("=> using pre-trained weights for DPSNet")
weights = torch.load(args.pretrained_dps)
dpsnet.load_state_dict(weights['state_dict'])
else:
dpsnet.init_weights()
cudnn.benchmark = True
dpsnet = torch.nn.DataParallel(dpsnet)
print('=> setting adam solver')
parameters = chain(dpsnet.parameters())
optimizer = torch.optim.Adam(parameters,
args.lr,
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(['train_loss'])
for epoch in range(args.epochs):
adjust_learning_rate(args, optimizer, epoch)
# train for one epoch
train_loss = train(args, train_loader, dpsnet, optimizer,
args.epoch_size, training_writer)
errors, error_names = validate_with_gt(args, val_loader, dpsnet, epoch,
output_writers)
error_string = ', '.join('{} : {:.3f}'.format(name, error)
for name, error in zip(error_names, errors))
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[0]
save_checkpoint(args.save_path, {
'epoch': epoch + 1,
'state_dict': dpsnet.module.state_dict()
}, epoch)
with open(args.save_path / args.log_summary, 'a') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow([train_loss, decisive_error])