Data loader for NeRF
(args)
| 495 | return train_dl, val_dl, test_dl, hwf, i_split, bounds.min(), bounds.max() |
| 496 | |
| 497 | def load_7Scenes_dataloader_NeRF(args): |
| 498 | ''' Data loader for NeRF ''' |
| 499 | |
| 500 | data_dir, scene = osp.split(args.datadir) # ../data/7Scenes, chess |
| 501 | dataset_folder, dataset = osp.split(data_dir) # ../data, 7Scenes |
| 502 | |
| 503 | data_transform = transforms.Compose([ |
| 504 | transforms.ToTensor()]) |
| 505 | target_transform = transforms.Lambda(lambda x: torch.Tensor(x)) |
| 506 | |
| 507 | data_dir = osp.join('..', 'data', 'deepslam_data', dataset) # ../data/deepslam_data/7Scenes |
| 508 | |
| 509 | ret_idx = False # return frame index |
| 510 | fix_idx = False # return frame index=0 in training |
| 511 | ret_hist = False |
| 512 | |
| 513 | if 'NeRFH' in args: |
| 514 | ret_idx = True |
| 515 | if args.fix_index: |
| 516 | fix_idx = True |
| 517 | |
| 518 | # encode hist experiment |
| 519 | if args.encode_hist: |
| 520 | ret_idx = False |
| 521 | fix_idx = False |
| 522 | ret_hist = True |
| 523 | |
| 524 | kwargs = dict(scene=scene, data_path=data_dir, |
| 525 | transform=data_transform, target_transform=target_transform, |
| 526 | df=args.df, ret_idx=ret_idx, fix_idx=fix_idx, ret_hist=ret_hist, hist_bin=args.hist_bin) |
| 527 | |
| 528 | train_set = SevenScenes(train=True, trainskip=args.trainskip, **kwargs) |
| 529 | val_set = SevenScenes(train=False, testskip=args.testskip, **kwargs) |
| 530 | |
| 531 | i_train = train_set.gt_idx |
| 532 | i_val = val_set.gt_idx |
| 533 | i_test = val_set.gt_idx |
| 534 | |
| 535 | # use a pose average stats computed earlier to unify posenet and nerf training |
| 536 | if args.save_pose_avg_stats or args.load_pose_avg_stats: |
| 537 | pose_avg_stats_file = osp.join(args.datadir, 'pose_avg_stats.txt') |
| 538 | train_set, val_set, bounds = fix_coord(args, train_set, val_set, pose_avg_stats_file) |
| 539 | else: |
| 540 | train_set, val_set, bounds = fix_coord(args, train_set, val_set) |
| 541 | |
| 542 | render_poses = None |
| 543 | render_img = None |
| 544 | |
| 545 | train_shuffle=True |
| 546 | if args.render_video_train or args.render_test: # or args.method3: |
| 547 | train_shuffle=False |
| 548 | train_dl = DataLoader(train_set, batch_size=1, shuffle=train_shuffle) |
| 549 | val_dl = DataLoader(val_set, batch_size=1, shuffle=False) |
| 550 | |
| 551 | hwf = [train_set.H, train_set.W, train_set.focal] |
| 552 | |
| 553 | i_split = [i_train, i_val, i_test] |
| 554 |
no test coverage detected