Data loader for NeRF
(args)
| 418 | return train_dl, val_dl, test_dl, hwf, i_split, bounds.min(), bounds.max() |
| 419 | |
| 420 | def load_Cambridge_dataloader_NeRF(args): |
| 421 | ''' Data loader for NeRF ''' |
| 422 | |
| 423 | data_dir, scene = osp.split(args.datadir) # ../data/7Scenes, chess |
| 424 | dataset_folder, dataset = osp.split(data_dir) # ../data, 7Scenes |
| 425 | |
| 426 | data_transform = transforms.Compose([ |
| 427 | transforms.ToTensor()]) |
| 428 | target_transform = transforms.Lambda(lambda x: torch.Tensor(x)) |
| 429 | |
| 430 | ret_idx = False # return frame index |
| 431 | fix_idx = False # return frame index=0 in training |
| 432 | ret_hist = False |
| 433 | |
| 434 | if 'NeRFH' in args: |
| 435 | ret_idx = True |
| 436 | if args.fix_index: |
| 437 | fix_idx = True |
| 438 | |
| 439 | # encode hist experiment |
| 440 | if args.encode_hist: |
| 441 | ret_idx = False |
| 442 | fix_idx = False |
| 443 | ret_hist = True |
| 444 | |
| 445 | kwargs = dict(scene=scene, data_path=data_dir, |
| 446 | transform=data_transform, target_transform=target_transform, |
| 447 | df=args.df, ret_idx=ret_idx, fix_idx=fix_idx, ret_hist=ret_hist, hist_bin=args.hist_bin) |
| 448 | |
| 449 | train_set = Cambridge2(train=True, trainskip=args.trainskip, **kwargs) |
| 450 | val_set = Cambridge2(train=False, testskip=args.testskip, **kwargs) |
| 451 | |
| 452 | i_train = train_set.gt_idx |
| 453 | i_val = val_set.gt_idx |
| 454 | i_test = val_set.gt_idx |
| 455 | |
| 456 | # use a pose average stats computed earlier to unify posenet and nerf training |
| 457 | if args.save_pose_avg_stats or args.load_pose_avg_stats: |
| 458 | pose_avg_stats_file = osp.join(args.datadir, 'pose_avg_stats.txt') |
| 459 | train_set, val_set, bounds = fix_coord(args, train_set, val_set, pose_avg_stats_file) |
| 460 | else: |
| 461 | train_set, val_set, bounds = fix_coord(args, train_set, val_set) |
| 462 | |
| 463 | render_poses = None |
| 464 | render_img = None |
| 465 | |
| 466 | train_shuffle=True |
| 467 | if args.render_video_train or args.render_test or args.dataset_type == 'Cambridge2': |
| 468 | train_shuffle=False |
| 469 | train_dl = DataLoader(train_set, batch_size=1, shuffle=train_shuffle) # default |
| 470 | # train_dl = DataLoader(train_set, batch_size=1, shuffle=False) # for debug only |
| 471 | val_dl = DataLoader(val_set, batch_size=1, shuffle=False) |
| 472 | |
| 473 | hwf = [train_set.H, train_set.W, train_set.focal] |
| 474 | |
| 475 | i_split = [i_train, i_val, i_test] |
| 476 | |
| 477 | return train_dl, val_dl, hwf, i_split, bounds, render_poses, render_img |
no test coverage detected