| """ |
| ScanNet Pair Dataset (Frame-level contrastive view) |
| |
| Refer PointContrast |
| |
| Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) |
| Please cite our work if the code is helpful to you. |
| """ |
|
|
| import os |
| import glob |
| import numpy as np |
| import torch |
| from copy import deepcopy |
| from torch.utils.data import Dataset |
|
|
| from pointcept.utils.logger import get_root_logger |
| from .builder import DATASETS |
| from .transform import Compose, TRANSFORMS |
|
|
|
|
| @DATASETS.register_module() |
| class ScanNetPairDataset(Dataset): |
| def __init__( |
| self, |
| data_root="data/scannet_pair", |
| overlap_threshold=0.3, |
| view1_transform=None, |
| view2_transform=None, |
| loop=1, |
| **kwargs |
| ): |
| super(ScanNetPairDataset, self).__init__() |
| self.data_root = data_root |
| self.overlap_threshold = overlap_threshold |
| self.view1_transform = Compose(view1_transform) |
| self.view2_transform = Compose(view2_transform) |
| self.loop = loop |
| self.data_list = self.get_data_list() |
| logger = get_root_logger() |
| logger.info("Totally {} x {} samples.".format(len(self.data_list), self.loop)) |
|
|
| def get_data_list(self): |
| data_list = [] |
| overlap_list = glob.glob( |
| os.path.join(self.data_root, "*", "pcd", "overlap.txt") |
| ) |
| for overlap_file in overlap_list: |
| with open(overlap_file) as f: |
| overlap = f.readlines() |
| overlap = [pair.strip().split() for pair in overlap] |
| data_list.extend( |
| [ |
| pair[:2] |
| for pair in overlap |
| if float(pair[2]) > self.overlap_threshold |
| ] |
| ) |
| return data_list |
|
|
| def get_data(self, idx): |
| pair = self.data_list[idx % len(self.data_list)] |
| view1_dict = torch.load(self.data_root + pair[0]) |
| view2_dict = torch.load(self.data_root + pair[1]) |
| return view1_dict, view2_dict |
|
|
| def get_data_name(self, idx): |
| return os.path.basename(self.data_list[idx % len(self.data_list)]).split(".")[0] |
|
|
| def prepare_train_data(self, idx): |
| |
| view1_dict, view2_dict = self.get_data(idx) |
| view1_dict = self.view1_transform(view1_dict) |
| view2_dict = self.view2_transform(view2_dict) |
| data_dict = dict() |
| for key, value in view1_dict.items(): |
| data_dict["view1_" + key] = value |
| for key, value in view2_dict.items(): |
| data_dict["view2_" + key] = value |
| return data_dict |
|
|
| def prepare_test_data(self, idx): |
| raise NotImplementedError |
|
|
| def __getitem__(self, idx): |
| return self.prepare_train_data(idx) |
|
|
| def __len__(self): |
| return len(self.data_list) * self.loop |
|
|
