| import os |
| import os.path as osp |
| import shlex |
| import shutil |
| import subprocess |
|
|
| import lmdb |
| import msgpack_numpy |
| import numpy as np |
| import torch |
| import torch.utils.data as data |
| import tqdm |
|
|
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) |
|
|
|
|
| def pc_normalize(pc): |
| l = pc.shape[0] |
| centroid = np.mean(pc, axis=0) |
| pc = pc - centroid |
| m = np.max(np.sqrt(np.sum(pc ** 2, axis=1))) |
| pc = pc / m |
| return pc |
|
|
|
|
| class ModelNet40Cls(data.Dataset): |
| def __init__(self, num_points, transforms=None, train=True, download=True): |
| super().__init__() |
|
|
| self.transforms = transforms |
|
|
| self.set_num_points(num_points) |
| self._cache = os.path.join(BASE_DIR, "modelnet40_normal_resampled_cache") |
|
|
| if not osp.exists(self._cache): |
| self.folder = "modelnet40_normal_resampled" |
| self.data_dir = os.path.join(BASE_DIR, self.folder) |
| self.url = ( |
| "https://shapenet.cs.stanford.edu/media/modelnet40_normal_resampled.zip" |
| ) |
|
|
| if download and not os.path.exists(self.data_dir): |
| zipfile = os.path.join(BASE_DIR, os.path.basename(self.url)) |
| subprocess.check_call( |
| shlex.split("curl {} -o {}".format(self.url, zipfile)) |
| ) |
|
|
| subprocess.check_call( |
| shlex.split("unzip {} -d {}".format(zipfile, BASE_DIR)) |
| ) |
|
|
| subprocess.check_call(shlex.split("rm {}".format(zipfile))) |
|
|
| self.train = train |
| self.set_num_points(num_points) |
|
|
| self.catfile = os.path.join(self.data_dir, "modelnet40_shape_names.txt") |
| self.cat = [line.rstrip() for line in open(self.catfile)] |
| self.classes = dict(zip(self.cat, range(len(self.cat)))) |
|
|
| os.makedirs(self._cache) |
|
|
| print("Converted to LMDB for faster dataloading while training") |
| for split in ["train", "test"]: |
| if split == "train": |
| shape_ids = [ |
| line.rstrip() |
| for line in open( |
| os.path.join(self.data_dir, "modelnet40_train.txt") |
| ) |
| ] |
| else: |
| shape_ids = [ |
| line.rstrip() |
| for line in open( |
| os.path.join(self.data_dir, "modelnet40_test.txt") |
| ) |
| ] |
|
|
| shape_names = ["_".join(x.split("_")[0:-1]) for x in shape_ids] |
| |
| self.datapath = [ |
| ( |
| shape_names[i], |
| os.path.join(self.data_dir, shape_names[i], shape_ids[i]) |
| + ".txt", |
| ) |
| for i in range(len(shape_ids)) |
| ] |
|
|
| with lmdb.open( |
| osp.join(self._cache, split), map_size=1 << 36 |
| ) as lmdb_env, lmdb_env.begin(write=True) as txn: |
| for i in tqdm.trange(len(self.datapath)): |
| fn = self.datapath[i] |
| point_set = np.loadtxt(fn[1], delimiter=",").astype(np.float32) |
| cls = self.classes[self.datapath[i][0]] |
| cls = int(cls) |
|
|
| txn.put( |
| str(i).encode(), |
| msgpack_numpy.packb( |
| dict(pc=point_set, lbl=cls), use_bin_type=True |
| ), |
| ) |
|
|
| shutil.rmtree(self.data_dir) |
|
|
| self._lmdb_file = osp.join(self._cache, "train" if train else "test") |
| with lmdb.open(self._lmdb_file, map_size=1 << 36) as lmdb_env: |
| self._len = lmdb_env.stat()["entries"] |
|
|
| self._lmdb_env = None |
|
|
| def __getitem__(self, idx): |
| if self._lmdb_env is None: |
| self._lmdb_env = lmdb.open( |
| self._lmdb_file, map_size=1 << 36, readonly=True, lock=False |
| ) |
|
|
| with self._lmdb_env.begin(buffers=True) as txn: |
| ele = msgpack_numpy.unpackb(txn.get(str(idx).encode()), raw=False) |
|
|
| point_set = ele["pc"] |
|
|
| pt_idxs = np.arange(0, self.num_points) |
| np.random.shuffle(pt_idxs) |
|
|
| point_set = point_set[pt_idxs, :] |
| point_set[:, 0:3] = pc_normalize(point_set[:, 0:3]) |
|
|
| if self.transforms is not None: |
| point_set = self.transforms(point_set) |
|
|
| return point_set, ele["lbl"] |
|
|
| def __len__(self): |
| return self._len |
|
|
| def set_num_points(self, pts): |
| self.num_points = min(int(1e4), pts) |
|
|
|
|
| if __name__ == "__main__": |
| from torchvision import transforms |
| import data_utils as d_utils |
|
|
| transforms = transforms.Compose( |
| [ |
| d_utils.PointcloudToTensor(), |
| d_utils.PointcloudRotate(axis=np.array([1, 0, 0])), |
| d_utils.PointcloudScale(), |
| d_utils.PointcloudTranslate(), |
| d_utils.PointcloudJitter(), |
| ] |
| ) |
| dset = ModelNet40Cls(16, train=True, transforms=transforms) |
| print(dset[0][0]) |
| print(dset[0][1]) |
| print(len(dset)) |
| dloader = torch.utils.data.DataLoader(dset, batch_size=32, shuffle=True) |
|
|
