Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,578 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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
"""
ScanNet20 / ScanNet200 / ScanNet Data Efficient Dataset
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 collections.abc import Sequence
from pointcept.utils.logger import get_root_logger
from pointcept.utils.cache import shared_dict
from .builder import DATASETS
from .defaults import DefaultDataset
from .transform import Compose, TRANSFORMS
from .preprocessing.scannet.meta_data.scannet200_constants import (
VALID_CLASS_IDS_20,
VALID_CLASS_IDS_200,
)
@DATASETS.register_module()
class ScanNetDataset(DefaultDataset):
VALID_ASSETS = [
"coord",
"color",
"normal",
"segment20",
"instance",
]
class2id = np.array(VALID_CLASS_IDS_20)
def __init__(
self,
lr_file=None,
la_file=None,
**kwargs,
):
self.lr = np.loadtxt(lr_file, dtype=str) if lr_file is not None else None
self.la = torch.load(la_file) if la_file is not None else None
super().__init__(**kwargs)
def get_data_list(self):
if self.lr is None:
data_list = super().get_data_list()
else:
data_list = [
os.path.join(self.data_root, "train", name) for name in self.lr
]
return data_list
def get_data(self, idx):
data_path = self.data_list[idx % len(self.data_list)]
name = self.get_data_name(idx)
if self.cache:
cache_name = f"pointcept-{name}"
return shared_dict(cache_name)
data_dict = {}
assets = os.listdir(data_path)
for asset in assets:
if not asset.endswith(".npy"):
continue
if asset[:-4] not in self.VALID_ASSETS:
continue
data_dict[asset[:-4]] = np.load(os.path.join(data_path, asset))
data_dict["name"] = name
data_dict["coord"] = data_dict["coord"].astype(np.float32)
data_dict["color"] = data_dict["color"].astype(np.float32)
data_dict["normal"] = data_dict["normal"].astype(np.float32)
if "segment20" in data_dict.keys():
data_dict["segment"] = (
data_dict.pop("segment20").reshape([-1]).astype(np.int32)
)
elif "segment200" in data_dict.keys():
data_dict["segment"] = (
data_dict.pop("segment200").reshape([-1]).astype(np.int32)
)
else:
data_dict["segment"] = (
np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
)
if "instance" in data_dict.keys():
data_dict["instance"] = (
data_dict.pop("instance").reshape([-1]).astype(np.int32)
)
else:
data_dict["instance"] = (
np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
)
if self.la:
sampled_index = self.la[self.get_data_name(idx)]
mask = np.ones_like(data_dict["segment"], dtype=bool)
mask[sampled_index] = False
data_dict["segment"][mask] = self.ignore_index
data_dict["sampled_index"] = sampled_index
return data_dict
@DATASETS.register_module()
class ScanNet200Dataset(ScanNetDataset):
VALID_ASSETS = [
"coord",
"color",
"normal",
"segment200",
"instance",
]
class2id = np.array(VALID_CLASS_IDS_200)
|