def test(model_path): test_args = TestOptions().parse() test_args.thread = 0 test_args.batchsize = 1 merge_cfg_from_file(test_args) data_loader = CustomerDataLoader(test_args) test_datasize = len(data_loader) logger.info('{:>15}: {:<30}'.format('test_data_size', test_datasize)) # load model model = MetricDepthModel() model.eval() test_args.load_ckpt = model_path # load checkpoint if test_args.load_ckpt: load_ckpt(test_args, model) model.cuda() # model = torch.nn.DataParallel(model) # test smoothed_absRel = SmoothedValue(test_datasize) smoothed_rms = SmoothedValue(test_datasize) smoothed_logRms = SmoothedValue(test_datasize) smoothed_squaRel = SmoothedValue(test_datasize) smoothed_silog = SmoothedValue(test_datasize) smoothed_silog2 = SmoothedValue(test_datasize) smoothed_log10 = SmoothedValue(test_datasize) smoothed_delta1 = SmoothedValue(test_datasize) smoothed_delta2 = SmoothedValue(test_datasize) smoothed_delta3 = SmoothedValue(test_datasize) smoothed_whdr = SmoothedValue(test_datasize) smoothed_criteria = { 'err_absRel': smoothed_absRel, 'err_squaRel': smoothed_squaRel, 'err_rms': smoothed_rms, 'err_silog': smoothed_silog, 'err_logRms': smoothed_logRms, 'err_silog2': smoothed_silog2, 'err_delta1': smoothed_delta1, 'err_delta2': smoothed_delta2, 'err_delta3': smoothed_delta3, 'err_log10': smoothed_log10, 'err_whdr': smoothed_whdr } for i, data in enumerate(data_loader): out = model.inference(data) pred_depth = torch.squeeze(out['b_fake']) img_path = data['A_paths'] invalid_side = data['invalid_side'][0] pred_depth = pred_depth[invalid_side[0]:pred_depth.size(0) - invalid_side[1], :] pred_depth = pred_depth / data['ratio'].cuda() # scale the depth pred_depth = resize_image(pred_depth, torch.squeeze(data['B_raw']).shape) smoothed_criteria = evaluate_err(pred_depth, data['B_raw'], smoothed_criteria, mask=(45, 471, 41, 601), scale=10.) # save images model_name = test_args.load_ckpt.split('/')[-1].split('.')[0] image_dir = os.path.join(cfg.ROOT_DIR, './evaluation', cfg.MODEL.ENCODER, model_name) if not os.path.exists(image_dir): os.makedirs(image_dir) img_name = img_path[0].split('/')[-1] #plt.imsave(os.path.join(image_dir, 'd_' + img_name), pred_depth, cmap='rainbow') #cv2.imwrite(os.path.join(image_dir, 'rgb_' + img_name), data['A_raw'].numpy().squeeze()) # print('processing (%04d)-th image... %s' % (i, img_path)) # print("###############absREL ERROR: %f", smoothed_criteria['err_absRel'].GetGlobalAverageValue()) # print("###############silog ERROR: %f", np.sqrt(smoothed_criteria['err_silog2'].GetGlobalAverageValue() - ( # smoothed_criteria['err_silog'].GetGlobalAverageValue()) ** 2)) # print("###############log10 ERROR: %f", smoothed_criteria['err_log10'].GetGlobalAverageValue()) # print("###############RMS ERROR: %f", np.sqrt(smoothed_criteria['err_rms'].GetGlobalAverageValue())) # print("###############delta_1 ERROR: %f", smoothed_criteria['err_delta1'].GetGlobalAverageValue()) # print("###############delta_2 ERROR: %f", smoothed_criteria['err_delta2'].GetGlobalAverageValue()) # print("###############delta_3 ERROR: %f", smoothed_criteria['err_delta3'].GetGlobalAverageValue()) # print("###############squaRel ERROR: %f", smoothed_criteria['err_squaRel'].GetGlobalAverageValue()) # print("###############logRms ERROR: %f", np.sqrt(smoothed_criteria['err_logRms'].GetGlobalAverageValue())) f.write("tested model:" + model_path) f.write('\n') f.write("###############absREL ERROR:" + str(smoothed_criteria['err_absRel'].GetGlobalAverageValue())) f.write('\n') f.write("###############silog ERROR:" + str( np.sqrt(smoothed_criteria['err_silog2'].GetGlobalAverageValue() - (smoothed_criteria['err_silog'].GetGlobalAverageValue())**2))) f.write('\n') f.write("###############log10 ERROR:" + str(smoothed_criteria['err_log10'].GetGlobalAverageValue())) f.write('\n') f.write("###############RMS ERROR:" + str(np.sqrt(smoothed_criteria['err_rms'].GetGlobalAverageValue()))) f.write('\n') f.write("###############delta_1 ERROR:" + str(smoothed_criteria['err_delta1'].GetGlobalAverageValue())) f.write('\n') f.write("###############delta_2 ERROR:" + str(smoothed_criteria['err_delta2'].GetGlobalAverageValue())) f.write('\n') f.write("###############delta_3 ERROR:" + str(smoothed_criteria['err_delta3'].GetGlobalAverageValue())) f.write('\n') f.write("###############squaRel ERROR:" + str(smoothed_criteria['err_squaRel'].GetGlobalAverageValue())) f.write('\n') f.write( "###############logRms ERROR:" + str(np.sqrt(smoothed_criteria['err_logRms'].GetGlobalAverageValue()))) f.write('\n') f.write( '-----------------------------------------------------------------------------' ) f.write('\n')
loss_func = ModelLoss() val_err = [{'abs_rel': 0, 'silog': 0}] ignore_step = -1 # Lerning strategy lr_optim_lambda = lambda iter: (1.0 - iter / (float(total_iters)))**0.9 scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer.optimizer, lr_lambda=lr_optim_lambda) #scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer.optimizer, milestones=train_args.epoch, gamma=0.1) # load checkpoint if train_args.load_ckpt: load_ckpt(train_args, model, optimizer.optimizer, scheduler, val_err) ignore_step = train_args.start_step - train_args.start_epoch * math.ceil( train_datasize / train_args.batchsize) if gpu_num != -1: model = torch.nn.DataParallel(model) try: for epoch in range(train_args.start_epoch, cfg.TRAIN.EPOCH[-1]): # training train(train_dataloader, model, epoch, loss_func, optimizer, scheduler, training_stats, val_dataloader, val_err, ignore_step) ignore_step = -1 except (RuntimeError, KeyboardInterrupt): logger.info('Save ckpt on exception ...')
test_args = TestOptions().parse() test_args.thread = 1 test_args.batchsize = 1 merge_cfg_from_file(test_args) data_loader = CustomerDataLoader(test_args) test_datasize = len(data_loader) logger.info('{:>15}: {:<30}'.format('test_data_size', test_datasize)) # load model model = MetricDepthModel() model.eval() # load checkpoint if test_args.load_ckpt: load_ckpt(test_args, model) model.cuda() model = torch.nn.DataParallel(model) # test smoothed_absRel = SmoothedValue(test_datasize) smoothed_rms = SmoothedValue(test_datasize) smoothed_logRms = SmoothedValue(test_datasize) smoothed_squaRel = SmoothedValue(test_datasize) smoothed_silog = SmoothedValue(test_datasize) smoothed_silog2 = SmoothedValue(test_datasize) smoothed_log10 = SmoothedValue(test_datasize) smoothed_delta1 = SmoothedValue(test_datasize) smoothed_delta2 = SmoothedValue(test_datasize) smoothed_delta3 = SmoothedValue(test_datasize) smoothed_whdr = SmoothedValue(test_datasize)
def main(): test_args = TestOptions().parse() test_args.thread = 1 # test code only supports thread = 1 test_args.batchsize = 1 # test code only supports batchSize = 1 data_loader = CustomerDataLoader(test_args) test_datasize = len(data_loader) logger.info('{:>15}: {:<30}'.format('test_data_size', test_datasize)) # load model model = DepthNormal() # evaluate mode model.eval() # load checkpoint if test_args.load_ckpt: load_ckpt(test_args, model) model.cuda() model = torch.nn.DataParallel(model) for i, data in enumerate(data_loader): out = model.module.inference(data) pred_depth = np.squeeze(out['b_fake']) * 80. # [h, w] pred_conf = np.squeeze(out['b_fake_conf']) # [c, h, w] # the image size has been padded to the size (385, 1243) pred_depth_crop = pred_depth[data['pad_raw'][0][0]:, data['pad_raw'][0][2]:] pred_conf_crop = pred_conf[:, data['pad_raw'][0][0]:, data['pad_raw'][0][2]:] sample_th = 0.15 sample_mask = get_sample_mask(pred_conf_crop.cpu().numpy(), threshold=sample_th) # [h, w] ####################################################################################### # add by users img_name = data['A_paths'][0].split('/')[-1][:-4] calib_name = img_name + '.txt' calib_dir = os.path.join(calib_fold, calib_name) camera_para = np.genfromtxt(calib_dir, delimiter=' ', skip_footer=3, dtype=None) P3_0 = camera_para[3] P2_0 = camera_para[2] P3_2 = P3_0 P3_2[4] -= P2_0[4] R0_rect = np.genfromtxt(calib_dir, delimiter=' ', skip_header=4, skip_footer=2) Tr_velo_to_cam0 = np.genfromtxt(calib_dir, delimiter=' ', skip_header=5, skip_footer=1) pcd_cam2 = reconstruct_3D(pred_depth_crop.cpu().numpy(), P3_2[3], P3_2[7], P3_2[1], P3_2[6]) # Transfer points in cam2 coordinate to cam0 coordinate pcd_cam0 = pcd_cam2 - np.array([[[P2_0[4] / P2_0[1]]], [[P2_0[8] / P2_0[1]]], [[P2_0[12] / P2_0[1]]]]) # Transfer points in cam0 coordinate to velo coordinate pcd_velo = transfer_points_in_cam0_to_velo(pcd_cam0, R0_rect, Tr_velo_to_cam0) rgb = data['A_raw'][0].cpu().numpy() save_ply(pcd_velo, rgb, os.path.join(pcd_folder, img_name) + '_sample.ply', sample_mask=sample_mask) #save_ply(pcd_cam2, rgb, os.path.join(pcd_folder, img_name) + '.ply') print('saved', img_name)
else: img = img.astype(np.float32) img = torch.from_numpy(img) return img if __name__ == '__main__': args = parse_args() # create depth model depth_model = RelDepthModel() depth_model.eval() # load checkpoint load_ckpt(args, depth_model, None, None) depth_model.cuda() image_dir = os.path.dirname(os.path.dirname(__file__)) + '/test_images/' imgs_list = os.listdir(image_dir) imgs_list.sort() imgs_path = [ os.path.join(image_dir, i) for i in imgs_list if i != 'outputs' ] image_dir_out = image_dir + '/outputs' os.makedirs(image_dir_out, exist_ok=True) for i, v in enumerate(imgs_path): print('processing (%04d)-th image... %s' % (i, v)) rgb = cv2.imread(v) rgb_c = rgb[:, :, ::-1].copy()
def main_worker(local_rank: int, ngpus_per_node: int, train_args, val_args): train_args.global_rank = train_args.node_rank * ngpus_per_node + local_rank train_args.local_rank = local_rank val_args.global_rank = train_args.global_rank val_args.local_rank = local_rank merge_cfg_from_file(train_args) global logger # Set logger log_output_dir = cfg.TRAIN.LOG_DIR if log_output_dir: try: os.makedirs(log_output_dir) except OSError as e: if e.errno != errno.EEXIST: raise logger = setup_distributed_logger("lib", log_output_dir, local_rank, cfg.TRAIN.RUN_NAME + '.txt') tblogger = None if train_args.use_tfboard and local_rank == 0: from tensorboardX import SummaryWriter tblogger = SummaryWriter(cfg.TRAIN.LOG_DIR) # init if train_args.distributed: torch.cuda.set_device(local_rank) dist.init_process_group(backend='nccl', init_method=train_args.dist_url, world_size=train_args.world_size, rank=train_args.global_rank) # load model model = RelDepthModel() if train_args.world_size > 1: assert is_pytorch_1_1_0_or_later(), \ "SyncBatchNorm is only available in pytorch >= 1.1.0" model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) logger.info('Using SyncBN!') if train_args.distributed: model = torch.nn.parallel.DistributedDataParallel( model.cuda(), device_ids=[local_rank], output_device=local_rank) else: model = torch.nn.DataParallel(model.cuda()) val_err = [{'abs_rel': 0, 'whdr': 0}] # Print configs and logs print_options(train_args, logger) # training and validation dataloader train_dataloader, train_sample_size = MultipleDataLoaderDistributed( train_args) val_dataloader, val_sample_size = MultipleDataLoaderDistributed(val_args) cfg.TRAIN.LR_SCHEDULER_MULTISTEPS = np.array( train_args.lr_scheduler_multiepochs) * math.ceil( np.sum(train_sample_size) / (train_args.world_size * train_args.batchsize)) # Optimizer optimizer = ModelOptimizer(model) # lr_optim_lambda = lambda iter: (1.0 - iter / (float(total_iters))) ** 0.9 # scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer.optimizer, lr_lambda=lr_optim_lambda) # lr_scheduler_step = 15000 # scheduler = torch.optim.lr_scheduler.StepLR(optimizer.optimizer, step_size=lr_scheduler_step, gamma=0.9) scheduler = make_lr_scheduler(cfg=cfg, optimizer=optimizer.optimizer) if train_args.load_ckpt: load_ckpt(train_args, model, optimizer.optimizer, scheduler, val_err) # obtain the current sample ratio sample_ratio = increase_sample_ratio_steps( train_args.start_step, base_ratio=train_args.sample_start_ratio, step_size=train_args.sample_ratio_steps) # reconstruct the train_dataloader with the new sample_ratio train_dataloader, train_sample_size = MultipleDataLoaderDistributed( train_args, sample_ratio=sample_ratio) total_iters = math.ceil( np.sum(train_sample_size) / (train_args.world_size * train_args.batchsize)) * train_args.epoch cfg.TRAIN.MAX_ITER = total_iters cfg.TRAIN.GPU_NUM = train_args.world_size print_configs(cfg) save_to_disk = main_process(train_args.distributed, local_rank) do_train(train_dataloader, val_dataloader, train_args, model, save_to_disk, scheduler, optimizer, val_err, logger, tblogger)
def __init__(self): self.device = "cuda" if torch.cuda.is_available() else "cpu" # Setup AdaBins model self.adabins_nyu_infer_helper = InferenceHelper(dataset='nyu', device=self.device) self.adabins_kitti_infer_helper = InferenceHelper(dataset='kitti', device=self.device) # Setup DiverseDepth model class DiverseDepthArgs: def __init__(self): self.resume = False self.cfg_file = "lib/configs/resnext50_32x4d_diversedepth_regression_vircam" self.load_ckpt = "pretrained/DiverseDepth.pth" diverse_depth_args = DiverseDepthArgs() merge_cfg_from_file(diverse_depth_args) self.diverse_depth_model = RelDepthModel() self.diverse_depth_model.eval() # load checkpoint load_ckpt(diverse_depth_args, self.diverse_depth_model) # TODO: update this - see how `device` argument should be processsed if self.device == "cuda": self.diverse_depth_model.cuda() self.diverse_depth_model = torch.nn.DataParallel( self.diverse_depth_model) # Setup MiDaS model self.midas_model_path = "./pretrained/MiDaS_f6b98070.pt" midas_model_type = "large" # load network if midas_model_type == "large": self.midas_model = MidasNet(self.midas_model_path, non_negative=True) self.midas_net_w, self.midas_net_h = 384, 384 elif midas_model_type == "small": self.midas_model = MidasNet_small(self.midas_model_path, features=64, backbone="efficientnet_lite3", exportable=True, non_negative=True, blocks={'expand': True}) self.midas_net_w, self.midas_net_h = 256, 256 self.midas_transform = Compose([ Resize( self.midas_net_w, self.midas_net_h, resize_target=None, keep_aspect_ratio=True, ensure_multiple_of=32, resize_method="upper_bound", image_interpolation_method=cv2.INTER_CUBIC, ), NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), PrepareForNet(), ]) self.midas_model.eval() self.midas_optimize = True if self.midas_optimize == True: rand_example = torch.rand(1, 3, self.midas_net_h, self.midas_net_w) self.midas_model(rand_example) traced_script_module = torch.jit.trace(self.midas_model, rand_example) self.midas_model = traced_script_module if self.device == "cuda": self.midas_model = self.midas_model.to( memory_format=torch.channels_last) self.midas_model = self.midas_model.half() self.midas_model.to(torch.device(self.device)) # Setup SGDepth model self.sgdepth_model = InferenceEngine.SgDepthInference() # Setup monodepth2 model self.monodepth2_model_path = "pretrained/monodepth2_mono+stereo_640x192" monodepth2_device = torch.device(self.device) encoder_path = os.path.join(self.monodepth2_model_path, "encoder.pth") depth_decoder_path = os.path.join(self.monodepth2_model_path, "depth.pth") # LOADING PRETRAINED MODEL print(" Loading Monodepth2 pretrained encoder") self.monodepth2_encoder = networks.ResnetEncoder(18, False) loaded_dict_enc = torch.load(encoder_path, map_location=monodepth2_device) # extract the height and width of image that this model was trained with self.feed_height = loaded_dict_enc['height'] self.feed_width = loaded_dict_enc['width'] filtered_dict_enc = { k: v for k, v in loaded_dict_enc.items() if k in self.monodepth2_encoder.state_dict() } self.monodepth2_encoder.load_state_dict(filtered_dict_enc) self.monodepth2_encoder.to(monodepth2_device) self.monodepth2_encoder.eval() print(" Loading pretrained decoder") self.monodepth2_depth_decoder = networks.DepthDecoder( num_ch_enc=self.monodepth2_encoder.num_ch_enc, scales=range(4)) loaded_dict = torch.load(depth_decoder_path, map_location=monodepth2_device) self.monodepth2_depth_decoder.load_state_dict(loaded_dict) self.monodepth2_depth_decoder.to(monodepth2_device) self.monodepth2_depth_decoder.eval()
def train(local_rank, distributed, train_args, logger, tblogger=None): # load model model = RelDepthModel() device = torch.device(cfg.MODEL.DEVICE) model.to(device) if cfg.MODEL.USE_SYNCBN: assert is_pytorch_1_1_0_or_later(), \ "SyncBatchNorm is only available in pytorch >= 1.1.0" model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) # Optimizer optimizer = ModelOptimizer(model) #lr_optim_lambda = lambda iter: (1.0 - iter / (float(total_iters))) ** 0.9 #scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer.optimizer, lr_lambda=lr_optim_lambda) lr_scheduler_step = 15000 scheduler = torch.optim.lr_scheduler.StepLR(optimizer.optimizer, step_size=lr_scheduler_step, gamma=0.9) if distributed: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[local_rank], output_device=local_rank) val_err = [{'abs_rel': 0, 'whdr': 0}] # training and validation dataloader val_dataloader = MultipleDataLoaderDistributed(val_args) if train_args.load_ckpt: load_ckpt(train_args, model, optimizer.optimizer, scheduler, val_err) # obtain the current sample ratio sample_ratio = increase_sample_ratio_steps( train_args.start_step, base_ratio=train_args.sample_start_ratio, step_size=train_args.sample_ratio_steps) # reconstruct the train_dataloader with the new sample_ratio train_dataloader = MultipleDataLoaderDistributed( train_args, sample_ratio=sample_ratio) if not train_args.resume: scheduler.__setattr__('step_size', lr_scheduler_step) else: train_dataloader = MultipleDataLoaderDistributed(train_args) train_datasize = len(train_dataloader) val_datasize = len(val_dataloader) merge_cfg_from_file(train_args) # total iterations total_iters = math.ceil( train_datasize / train_args.batchsize) * train_args.epoch cfg.TRAIN.MAX_ITER = total_iters cfg.TRAIN.GPU_NUM = gpu_num # Print configs and logs if get_rank() == 0: train_opt.print_options(train_args) val_opt.print_options(val_args) print_configs(cfg) logger.info('{:>15}: {:<30}'.format('GPU_num', gpu_num)) logger.info('{:>15}: {:<30}'.format('train_data_size', train_datasize)) logger.info('{:>15}: {:<30}'.format('val_data_size', val_datasize)) logger.info('{:>15}: {:<30}'.format('total_iterations', total_iters)) save_to_disk = get_rank() == 0 do_train(train_dataloader, val_dataloader, train_args, model, save_to_disk, scheduler, optimizer, val_err, logger, tblogger)