Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,766 Bytes
4893ce0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
"""
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):
# load data
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
|