Data loader for Pose Regression Network
(args)
| 420 | return |
| 421 | |
| 422 | def load_7Scenes_dataloader(args): |
| 423 | ''' Data loader for Pose Regression Network ''' |
| 424 | if args.pose_only: # if train posenet is true |
| 425 | pass |
| 426 | else: |
| 427 | raise Exception('load_7Scenes_dataloader() currently only support PoseNet Training, not NeRF training') |
| 428 | data_dir, scene = osp.split(args.datadir) # ../data/7Scenes, chess |
| 429 | dataset_folder, dataset = osp.split(data_dir) # ../data, 7Scenes |
| 430 | |
| 431 | # transformer |
| 432 | data_transform = transforms.Compose([ |
| 433 | # transforms.Resize(240), |
| 434 | transforms.ToTensor(), |
| 435 | ]) |
| 436 | target_transform = transforms.Lambda(lambda x: torch.Tensor(x)) |
| 437 | |
| 438 | data_dir = osp.join(dataset_folder, 'deepslam_data', dataset) # ../data/deepslam_data/7Scenes |
| 439 | |
| 440 | ret_idx = False # return frame index |
| 441 | fix_idx = False # return frame index=0 in training |
| 442 | ret_hist = False |
| 443 | |
| 444 | if 'NeRFH' in args: |
| 445 | if args.NeRFH == True: |
| 446 | ret_idx = True |
| 447 | if args.fix_index: |
| 448 | fix_idx = True |
| 449 | |
| 450 | # encode hist experiment |
| 451 | if args.encode_hist: |
| 452 | ret_idx = False |
| 453 | fix_idx = False |
| 454 | ret_hist = True |
| 455 | |
| 456 | # kwargs = dict(scene=scene, data_path=data_dir, transform=data_transform, target_transform=target_transform) |
| 457 | kwargs = dict(scene=scene, data_path=data_dir, |
| 458 | transform=data_transform, target_transform=target_transform, |
| 459 | df=args.df, ret_idx=ret_idx, fix_idx=fix_idx, |
| 460 | ret_hist=ret_hist, hist_bin=args.hist_bin) |
| 461 | # pdb.set_trace() |
| 462 | if args.finetune_unlabel: # direct-pn + unlabel |
| 463 | train_set = SevenScenes(train=False, testskip=args.trainskip, **kwargs) |
| 464 | val_set = SevenScenes(train=False, testskip=args.testskip, **kwargs) |
| 465 | |
| 466 | # if not args.eval: |
| 467 | # # remove overlap data in val_set that was already in train_set, |
| 468 | # train_set, val_set = remove_overlap_data(train_set, val_set) |
| 469 | else: |
| 470 | train_set = SevenScenes(train=True, trainskip=args.trainskip, **kwargs) |
| 471 | val_set = SevenScenes(train=False, testskip=args.testskip, **kwargs) |
| 472 | L = len(train_set) |
| 473 | |
| 474 | i_train = train_set.gt_idx |
| 475 | i_val = val_set.gt_idx |
| 476 | i_test = val_set.gt_idx |
| 477 | # print('TRAIN views are', i_train) |
| 478 | # print('TEST views are', i_test) |
| 479 | # print('VAL views are', i_val) |
no test coverage detected