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])