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", 5, args.numdepth, interval_scale=args.interval_scale) TestImgLoader = DataLoader(test_dataset, args.batch_size, shuffle=False, num_workers=4, drop_last=False) # model model = MVSNet(refine=False, upsample=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) outputs = model(sample_cuda["imgs"], sample_cuda["proj_matrices"], sample_cuda["depth_values"]) 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', '.npy')) confidence_filename = os.path.join( args.outdir, filename.format('confidence', '.npy')) os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True) os.makedirs(confidence_filename.rsplit('/', 1)[0], exist_ok=True) # save depth maps np.save(depth_filename, depth_est) # save_pfm(depth_filename, depth_est) # save confidence maps np.save(confidence_filename, photometric_confidence)
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()
# p: probability volume [B, D, H, W] # depth_values: discrete depth values [B, D] def depth_regression(p, depth_values): depth_values = depth_values.view(*depth_values.shape, 1, 1) depth = torch.sum(p * depth_values, 1) return depth if __name__ == "__main__": # some testing code, just IGNORE it from datasets import find_dataset_def from torch.utils.data import DataLoader import numpy as np import cv2 MVSDataset = find_dataset_def("dtu_yao") dataset = MVSDataset("/home/xyguo/dataset/dtu_mvs/processed/mvs_training/dtu/", '../lists/dtu/train.txt', 'train', 3, 256) dataloader = DataLoader(dataset, batch_size=2) item = next(iter(dataloader)) imgs = item["imgs"][:, :, :, ::4, ::4].cuda() proj_matrices = item["proj_matrices"].cuda() mask = item["mask"].cuda() depth = item["depth"].cuda() depth_values = item["depth_values"].cuda() imgs = torch.unbind(imgs, 1) proj_matrices = torch.unbind(proj_matrices, 1) ref_img, src_imgs = imgs[0], imgs[1:] ref_proj, src_projs = proj_matrices[0], proj_matrices[1:]
if args.mode == "train": if not os.path.isdir(args.logdir): os.mkdir(args.logdir) current_time_str = str(datetime.datetime.now().strftime('%Y%m%d_%H%M%S')) print("current time", current_time_str) print("creating new summary file") logger = SummaryWriter(args.logdir) print("argv:", sys.argv[1:]) print_args(args) # dataset, dataloader MVSDataset = find_dataset_def(args.dataset) if args.dataset == 'dtu_yao': train_dataset = MVSDataset(args.trainpath, args.trainlist, "train", 5, robust_train=True) test_dataset = MVSDataset(args.valpath, args.vallist, "val", 5, robust_train=False) TrainImgLoader = DataLoader(train_dataset, args.batch_size, shuffle=True,
if not os.path.exists(os.path.join(cfg.TEST.PATH, 'SyncedPoses.txt')): logger.info("First run on this captured data, start the pre-processing...") process_data(cfg.TEST.PATH) else: logger.info("Found SyncedPoses.txt, skipping data pre-processing...") logger.info("Running NeuralRecon...") transform = [transforms.ResizeImage((640, 480)), transforms.ToTensor(), transforms.RandomTransformSpace( cfg.MODEL.N_VOX, cfg.MODEL.VOXEL_SIZE, random_rotation=False, random_translation=False, paddingXY=0, paddingZ=0, max_epoch=cfg.TRAIN.EPOCHS), transforms.IntrinsicsPoseToProjection(cfg.TEST.N_VIEWS, 4)] transforms = transforms.Compose(transform) ARKitDataset = find_dataset_def(cfg.DATASET) test_dataset = ARKitDataset(cfg.TEST.PATH, "test", transforms, cfg.TEST.N_VIEWS, len(cfg.MODEL.THRESHOLDS) - 1) data_loader = DataLoader(test_dataset, cfg.BATCH_SIZE, shuffle=False, num_workers=cfg.TEST.N_WORKERS, drop_last=False) # model logger.info("Initializing the model on GPU...") model = NeuralRecon(cfg).cuda().eval() model = torch.nn.DataParallel(model, device_ids=[0]) # use the latest checkpoint file saved_models = [fn for fn in os.listdir(cfg.LOGDIR) if fn.endswith(".ckpt")] saved_models = sorted(saved_models, key=lambda x: int(x.split('_')[-1].split('.')[0])) loadckpt = os.path.join(cfg.LOGDIR, saved_models[-1]) logger.info("Resuming from " + str(loadckpt)) state_dict = torch.load(loadckpt) model.load_state_dict(state_dict['model'], strict=False)
transforms.ResizeImage((640, 480)), transforms.ToTensor(), transforms.RandomTransformSpace(cfg.MODEL.N_VOX, cfg.MODEL.VOXEL_SIZE, random_rotation, random_translation, paddingXY, paddingZ, max_epoch=cfg.TRAIN.EPOCHS), transforms.IntrinsicsPoseToProjection(n_views, 4), ] transforms = transforms.Compose(transform) # dataset, dataloader MVSDataset = find_dataset_def(cfg.DATASET) train_dataset = MVSDataset(cfg.TRAIN.PATH, "train", transforms, cfg.TRAIN.N_VIEWS, len(cfg.MODEL.THRESHOLDS) - 1) test_dataset = MVSDataset(cfg.TEST.PATH, "test", transforms, cfg.TEST.N_VIEWS, len(cfg.MODEL.THRESHOLDS) - 1) if cfg.DISTRIBUTED: train_sampler = DistributedSampler(train_dataset, shuffle=False) TrainImgLoader = torch.utils.data.DataLoader( train_dataset, batch_size=cfg.BATCH_SIZE, sampler=train_sampler, num_workers=cfg.TRAIN.N_WORKERS, pin_memory=True, drop_last=True)
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 save_depth(): # dataset, dataloader MVSDataset = find_dataset_def(args.dataset) test_dataset = MVSDataset(args.testpath, args.testlist, "test", args.numviews, 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)) # from collections import OrderedDict # new_state_dict = OrderedDict() state_dict = torch.load(args.loadckpt) # for k, v in state_dict["model"].items(): # name = k.replace("module.", "") # remove module. # new_state_dict[name] = v # model.load_state_dict(new_state_dict) #, strict=False) model.load_state_dict(state_dict['model']) #, strict=False) model.eval() with torch.no_grad(): for batch_idx, sample in enumerate(TestImgLoader): orig_images = torch.unbind(sample["imgs"], 1) ext_matrices = sample["ext_matrices"] int_matrices = sample["int_matrices"] ref_img = orig_images[0].squeeze().numpy() n_img = orig_images[1].squeeze().numpy() ref_ext = ext_matrices[0].squeeze().numpy() ref_int = int_matrices[0].squeeze().numpy() n_ext = ext_matrices[1].squeeze().numpy() n_int = int_matrices[1].squeeze().numpy() sample_cuda = tocuda(sample) outputs = model(sample_cuda["imgs"], sample_cuda["proj_matrices"], sample_cuda["depth_values"]) outputs = tensor2numpy(outputs) del sample_cuda ref_depth = cv2.resize(outputs["depth"].squeeze(), (ref_img.shape[2], ref_img.shape[1])) x = np.linspace(0, ref_img.shape[2], ref_img.shape[2], axis=-1) y = np.linspace(0, ref_img.shape[1], ref_img.shape[1], axis=-1) xv, yv = np.meshgrid(x, y) W = PixelCoordToWorldCoord(ref_int, ref_ext[:3, :3], ref_ext[:3, 3:4], xv, yv, ref_depth) xp, yp = WorldCoordTopixelCoord(n_int, n_ext[:3, :3], n_ext[:3, 3:4], W) # xp, yp = WorldCoordTopixelCoord(ref_int, ref_ext[:3,:3], ref_ext[:3, 3:4], W) projected_img = GetImageAtPixelCoordinates(ref_img, xp, yp) fig = plt.figure() plt.subplot(1, 3, 1) plt.imshow(np.swapaxes(ref_img, 0, 2).swapaxes(0, 1)) plt.subplot(1, 3, 2) plt.imshow(np.swapaxes(n_img, 0, 2).swapaxes(0, 1)) plt.subplot(1, 3, 3) plt.imshow(np.swapaxes(projected_img, 0, 2).swapaxes(0, 1)) plt.show() plt.tight_layout() print('Iter {}/{}'.format(batch_idx, len(TestImgLoader))) filenames = sample["filename"]