def save_depth(): # dataset, dataloader MVSDataset = find_dataset_def(args.dataset) # test_dataset = MVSDataset(args.testpath, args.testlist, "test", 5, args.numdepth, args.interval_scale) test_dataset = MVSDataset(args.testpath) TestImgLoader = DataLoader(test_dataset, args.batch_size, shuffle=False, num_workers=4, drop_last=False) # model model = EdgeFlow() model = nn.DataParallel(model) model.cuda() # load checkpoint file specified by args.loadckpt print("loading model {}".format(args.loadckpt)) state_dict = torch.load(args.loadckpt) model.load_state_dict(state_dict['model']) model.eval() with torch.no_grad(): for batch_idx, sample in enumerate(TestImgLoader): sample_cuda = tocuda(sample) # print(sample_cuda) depth, confidence, _ = model(sample_cuda["imgs"], sample_cuda["proj_matrices"], sample_cuda["depth_values"]) depth = tensor2numpy(depth) confidence = tensor2numpy(confidence) del sample_cuda torch.cuda.empty_cache() print('Iter {}/{}'.format(batch_idx, len(TestImgLoader))) filenames = sample["filename"] # save depth maps and confidence maps for filename, depth_est, photometric_confidence in zip( filenames, depth["{}x".format(args.stage)], confidence["{}x".format(args.stage)]): depth_filename = os.path.join( args.outdir, filename.format('depth_est_{}x'.format(args.stage), '.pfm')) confidence_filename = os.path.join( args.outdir, filename.format('confidence_{}x'.format(args.stage), '.pfm')) # print(photometric_confidence.shape) # mask = sample_cuda["mask"] os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True) os.makedirs(confidence_filename.rsplit('/', 1)[0], exist_ok=True) # save depth maps save_pfm(depth_filename, depth_est) # save confidence maps save_pfm(confidence_filename, photometric_confidence)
def save_depth(): # dataset, dataloader MVSDataset = find_dataset_def(args.dataset) test_dataset = MVSDataset(args.testpath, args.n_views) TestImgLoader = DataLoader(test_dataset, args.batch_size, shuffle=False, num_workers=4, drop_last=False) # model model = PatchmatchNet( patchmatch_interval_scale=args.patchmatch_interval_scale, propagation_range=args.patchmatch_range, patchmatch_iteration=args.patchmatch_iteration, patchmatch_num_sample=args.patchmatch_num_sample, propagate_neighbors=args.propagate_neighbors, evaluate_neighbors=args.evaluate_neighbors) model = nn.DataParallel(model) model.cuda() # load checkpoint file specified by args.loadckpt print("loading model {}".format(args.loadckpt)) state_dict = torch.load(args.loadckpt) model.load_state_dict(state_dict['model']) model.eval() with torch.no_grad(): for batch_idx, sample in enumerate(TestImgLoader): start_time = time.time() sample_cuda = tocuda(sample) outputs = model(sample_cuda["imgs"], sample_cuda["proj_matrices"], sample_cuda["depth_min"], sample_cuda["depth_max"]) outputs = tensor2numpy(outputs) del sample_cuda print('Iter {}/{}, time = {:.3f}'.format(batch_idx, len(TestImgLoader), time.time() - start_time)) filenames = sample["filename"] # save depth maps and confidence maps for filename, depth_est, photometric_confidence in zip( filenames, outputs["refined_depth"]['stage_0'], outputs["photometric_confidence"]): depth_filename = os.path.join( args.outdir, filename.format('depth_est', '.pfm')) confidence_filename = os.path.join( args.outdir, filename.format('confidence', '.pfm')) os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True) os.makedirs(confidence_filename.rsplit('/', 1)[0], exist_ok=True) # save depth maps depth_est = np.squeeze(depth_est, 0) save_pfm(depth_filename, depth_est) # save confidence maps save_pfm(confidence_filename, photometric_confidence)
def save_depth(): # dataset, dataloader MVSDataset = find_dataset_def(args.dataset) test_dataset = MVSDataset(args.testpath, args.testlist, "test", 3, args.numdepth, args.interval_scale) TestImgLoader = DataLoader(test_dataset, args.batch_size, shuffle=False, num_workers=4, drop_last=False) # model model = MVSNet(refine=False) model = nn.DataParallel(model) model.cuda() # load checkpoint file specified by args.loadckpt print("loading model {}".format(args.loadckpt)) state_dict = torch.load(args.loadckpt) model.load_state_dict(state_dict['model']) model.eval() with torch.no_grad(): for batch_idx, sample in enumerate(TestImgLoader): sample_cuda = tocuda(sample) print("begin forward") outputs = model(sample_cuda["imgs"], sample_cuda["proj_matrices"], sample_cuda["depth_values"]) print("forward over") outputs = tensor2numpy(outputs) del sample_cuda print('Iter {}/{}'.format(batch_idx, len(TestImgLoader))) filenames = sample["filename"] # save depth maps and confidence maps for filename, depth_est, photometric_confidence in zip( filenames, outputs["depth"], outputs["photometric_confidence"]): depth_filename = os.path.join( args.outdir, filename.format('depth_est', '.pfm')) confidence_filename = os.path.join( args.outdir, filename.format('confidence', '.pfm')) os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True) os.makedirs(confidence_filename.rsplit('/', 1)[0], exist_ok=True) # save depth maps print("save .pfm") save_pfm(depth_filename, depth_est) # save confidence maps save_pfm(confidence_filename, photometric_confidence)
def probability_filter(dense_folder, prob_threshold): image_folder = os.path.join(dense_folder, 'images') # convert cameras image_names = os.listdir(image_folder) for image_name in image_names: image_prefix = os.path.splitext(image_name)[0] init_depth_map_path = os.path.join(dense_folder, "depth_est", image_prefix + '.pfm') prob_map_path = os.path.join(dense_folder, "confidence", image_prefix + '.pfm') out_depth_map_path = os.path.join(dense_folder, "depth_est", image_prefix + '_prob_filtered.pfm') depth_map, _ = read_pfm(init_depth_map_path) prob_map, _ = read_pfm(prob_map_path) depth_map[prob_map < prob_threshold] = 0 save_pfm(out_depth_map_path, depth_map)
def save_scene_depth(testlist): # dataset, dataloader MVSDataset = find_dataset_def(args.dataset) test_dataset = MVSDataset(args.testpath, testlist, "test", args.num_view, args.numdepth, Interval_Scale, max_h=args.max_h, max_w=args.max_w, fix_res=args.fix_res) TestImgLoader = DataLoader(test_dataset, args.batch_size, shuffle=False, num_workers=4, drop_last=False) # model model = DDR_Net( refine=False, ndepths=[int(nd) for nd in args.ndepths.split(",") if nd], depth_interals_ratio=[ float(d_i) for d_i in args.depth_inter_r.split(",") if d_i ], share_cr=args.share_cr, cr_base_chs=[int(ch) for ch in args.cr_base_chs.split(",") if ch], grad_method=args.grad_method) # load checkpoint file specified by args.loadckpt print("loading model {}".format(args.loadckpt)) state_dict = torch.load(args.loadckpt, map_location=torch.device("cpu")) model.load_state_dict(state_dict['model'], strict=True) model = nn.DataParallel(model) model.cuda() model.eval() with torch.no_grad(): for batch_idx, sample in enumerate(TestImgLoader): sample_cuda = tocuda(sample) start_time = time.time() outputs, refine_depth_map = model(sample_cuda["imgs"], sample_cuda["proj_matrices"], sample_cuda["depth_values"]) end_time = time.time() #outputs = tensor2numpy(outputs) depth = outputs["depth"] #feature=outputs photometric_confidence = outputs["photometric_confidence"] depth = tensor2numpy(depth) photometric_confidence = tensor2numpy(photometric_confidence) del sample_cuda filenames = sample["filename"] cams = sample["proj_matrices"]["stage{}".format(num_stage)].numpy() imgs = sample["imgs"].numpy() print('Iter {}/{}, Time:{} Res:{}'.format(batch_idx, len(TestImgLoader), end_time - start_time, imgs[0].shape)) # save depth maps and confidence maps for filename, cam, img, depth_est, photometric_confidence in zip(filenames, cams, imgs, \ depth,photometric_confidence ): img = img[0] #ref view cam = cam[0] #ref cam depth_filename = os.path.join( args.outdir, filename.format('depth_est', '.pfm')) confidence_filename = os.path.join( args.outdir, filename.format('confidence', '.pfm')) cam_filename = os.path.join( args.outdir, filename.format('cams', '_cam.txt')) img_filename = os.path.join(args.outdir, filename.format('images', '.jpg')) ply_filename = os.path.join( args.outdir, filename.format('ply_local', '.ply')) feature_filename = os.path.join( args.outdir, filename.format('feature', '.jpg')) os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True) os.makedirs(confidence_filename.rsplit('/', 1)[0], exist_ok=True) os.makedirs(cam_filename.rsplit('/', 1)[0], exist_ok=True) os.makedirs(img_filename.rsplit('/', 1)[0], exist_ok=True) os.makedirs(ply_filename.rsplit('/', 1)[0], exist_ok=True) #save depth maps save_pfm(depth_filename, depth_est) #save confidence maps save_pfm(confidence_filename, photometric_confidence) #save cams, img write_cam(cam_filename, cam) img = np.clip(np.transpose(img, (1, 2, 0)) * 255, 0, 255).astype(np.uint8) img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) cv2.imwrite(img_filename, img_bgr) # vis # print(photometric_confidence.mean(), photometric_confidence.min(), photometric_confidence.max()) # import matplotlib.pyplot as plt # plt.subplot(1, 3, 1) # plt.imshow(img) # plt.subplot(1, 3, 2) # plt.imshow((depth_est - depth_est.min())/(depth_est.max() - depth_est.min())) # plt.subplot(1, 3, 3) # plt.imshow(photometric_confidence) # plt.show() if num_stage == 1: downsample_img = cv2.resize( img, (int(img.shape[1] * 0.25), int(img.shape[0] * 0.25))) elif num_stage == 2: downsample_img = cv2.resize( img, (int(img.shape[1] * 0.5), int(img.shape[0] * 0.5))) elif num_stage == 3: downsample_img = img if batch_idx % args.save_freq == 0: generate_pointcloud(downsample_img, depth_est, ply_filename, cam[1, :3, :3]) torch.cuda.empty_cache() gc.collect()
def save_depth(): # dataset, dataloader MVSDataset = find_dataset_def(args.dataset) test_dataset = MVSDataset(args.testpath, args.testlist, "test", 5, args.numdepth, args.interval_scale, args.inverse_depth, args.pyramid) TestImgLoader = DataLoader(test_dataset, args.batch_size, shuffle=False, num_workers=4, drop_last=False) # model if args.model == 'mvsnet': print('use MVSNet') model = MVSNet(refine=args.refine, fea_net=args.fea_net, cost_net=args.cost_net, refine_net=args.refine_net, origin_size=args.origin_size, cost_aggregation=args.cost_aggregation, dp_ratio=args.dp_ratio) else: print('input pre-defined model') model = nn.DataParallel(model) model.cuda() # load checkpoint file specified by args.loadckpt print("loading model {}".format(args.loadckpt)) state_dict = torch.load(args.loadckpt) model.load_state_dict(state_dict['model']) model.eval() count = -1 total_time = 0 with torch.no_grad(): for batch_idx, sample in enumerate(TestImgLoader): count += 1 print('process', sample['filename']) sample_cuda = tocuda(sample) print(sample_cuda["imgs"].shape) time_s = time.time() outputs = model(sample_cuda["imgs"], sample_cuda["proj_matrices"], sample_cuda["depth_values"]) one_time = time.time() - time_s total_time += one_time print('one forward: ', one_time) if count % 50 == 0: print('avg time:', total_time / 50) total_time = 0 if 'High' in args.fea_net and 'Coarse2Fine' in args.cost_net: tmp_outputs = {} for key, value in outputs.items(): tmp_outputs[key] = value[0] outputs = tmp_outputs outputs = tensor2numpy(outputs) del sample_cuda print('Iter {}/{}'.format(batch_idx, len(TestImgLoader))) filenames = sample["filename"] # save depth maps and confidence maps for filename, depth_est, photometric_confidence in zip( filenames, outputs["depth"], outputs["photometric_confidence"]): depth_filename = os.path.join( save_dir, filename.format('depth_est_{}'.format(args.pyramid), '.pfm')) confidence_filename = os.path.join( save_dir, filename.format('confidence_{}'.format(args.pyramid), '.pfm')) os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True) os.makedirs(confidence_filename.rsplit('/', 1)[0], exist_ok=True) # save depth maps print(depth_est.shape) save_pfm(depth_filename, depth_est.squeeze()) # save confidence maps save_pfm(confidence_filename, photometric_confidence.squeeze())
def multi_scale(depth_folder_down4, depth_folder_down8, depth_folder_down16, final_folder, test_list): depth_folder = depth_folder_down4 conf_folder = depth_folder output_folder = depth_folder depth2_folder = depth_folder_down8 conf2_folder = depth2_folder mask2_folder = depth2_folder depth3_folder = depth_folder_down16 conf3_folder = depth3_folder mask3_folder = depth3_folder if not os.path.exists(final_folder): os.mkdir(final_folder) file = open(test_list) scans = file.readlines() scans = [line.rstrip() for line in scans] sum = 0 ct = 0 shape_all = 0 for scan in scans: ct = ct + 1 depth_t = os.path.join(depth_folder, scan) depth2_t = os.path.join(depth2_folder, scan) depth3_t = os.path.join(depth3_folder, scan) conf_t = os.path.join(conf_folder, scan) conf2_t = os.path.join(conf2_folder, scan) conf3_t = os.path.join(conf3_folder, scan) mask2_folder_t = os.path.join(mask2_folder, scan) mask2_folder_t = os.path.join(mask2_folder_t, 'mask') mask3_folder_t = os.path.join(mask3_folder, scan) mask3_folder_t = os.path.join(mask3_folder_t, 'mask') output_t = os.path.join(output_folder, 'output') final_t = os.path.join(final_folder, scan) if not os.path.exists(final_t): os.mkdir(final_t) if not os.path.exists(output_t): os.mkdir(output_t) output_t = os.path.join(output_t, scan) if not os.path.exists(output_t): os.mkdir(output_t) sum_t = 0 shape = 0 for i in range(49): print('process depth ' + str(i)) depth_s = os.path.join(depth_t, 'depth_est') final_d = os.path.join(final_t, 'depth_est') if not os.path.exists(final_d): os.mkdir(final_d) final_c = os.path.join(final_t, 'confidence') if not os.path.exists(final_c): os.mkdir(final_c) depth_s = os.path.join(depth_s, '%08d.pfm' % i) depth2_s = os.path.join(depth2_t, 'depth_est') depth2_s = os.path.join(depth2_s, '%08d.pfm' % i) mask2_s = os.path.join(mask2_folder_t, '%08d_final.png' % i) depth3_s = os.path.join(depth3_t, 'depth_est') depth3_s = os.path.join(depth3_s, '%08d.pfm' % i) mask3_s = os.path.join(mask3_folder_t, '%08d_final.png' % i) conf_s = os.path.join(conf_t, 'confidence') conf_s = os.path.join(conf_s, '%08d.pfm' % i) final_d_s = os.path.join(final_d, '%08d.pfm' % i) final_c_s = os.path.join(final_c, '%08d.pfm' % i) conf2_s = os.path.join(conf2_t, 'confidence') conf2_s = os.path.join(conf2_s, '%08d.pfm' % i) conf3_s = os.path.join(conf3_t, 'confidence') conf3_s = os.path.join(conf3_s, '%08d.pfm' % i) output_s = os.path.join(output_t, 'error_map_%04d.png' % i) mask2_img_t = cv2.imread(mask2_s) mask2_img_t = cv2.cvtColor(mask2_img_t, cv2.COLOR_BGR2GRAY) mask2_arr = np.array(mask2_img_t) mask2_t = mask2_arr > 0 mask3_img_t = cv2.imread(mask3_s) mask3_img_t = cv2.cvtColor(mask3_img_t, cv2.COLOR_BGR2GRAY) mask3_arr = np.array(mask3_img_t) mask3_arr = cv2.pyrUp(mask3_arr) depth = read_pfm(depth_s)[0] depth2 = read_pfm(depth2_s)[0] depth3 = read_pfm(depth3_s)[0] depth3 = cv2.pyrUp(depth3) confidence = read_pfm(conf_s)[0] confidence2 = read_pfm(conf2_s)[0] confidence3 = read_pfm(conf3_s)[0] confidence3 = cv2.pyrUp(confidence3) depth2 = cv2.pyrUp(depth2) zero_depth = np.zeros([8, 400]) depth2 = np.r_[depth2, zero_depth] confidence2 = cv2.pyrUp(confidence2) zero_conf = np.zeros([8, 400]) confidence2 = np.r_[confidence2, zero_conf] mask2_arr = cv2.pyrUp(mask2_arr) zero_depth = np.zeros([8, 400]) mask2_arr = np.r_[mask2_arr, zero_depth] mask3 = (0.9 < confidence2) & (confidence < 0.5) & (mask2_arr > 0) depth[mask3] = depth2[mask3] confidence[mask3] = confidence2[mask3] save_pfm(final_d_s, depth) save_pfm(final_c_s, confidence)