File size: 3,253 Bytes
b30c1d8 |
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 |
import os
import shlex
import subprocess
import h5py
import numpy as np
import torch
import torch.utils.data as data
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
def _get_data_files(list_filename):
with open(list_filename) as f:
return [line.rstrip() for line in f]
def _load_data_file(name):
f = h5py.File(name, "r")
data = f["data"][:]
label = f["label"][:]
return data, label
class Indoor3DSemSeg(data.Dataset):
def __init__(self, num_points, train=True, download=True, data_precent=1.0):
super().__init__()
self.data_precent = data_precent
self.folder = "indoor3d_sem_seg_hdf5_data"
self.data_dir = os.path.join(BASE_DIR, self.folder)
self.url = (
"https://shapenet.cs.stanford.edu/media/indoor3d_sem_seg_hdf5_data.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, self.num_points = train, num_points
all_files = _get_data_files(os.path.join(self.data_dir, "all_files.txt"))
room_filelist = _get_data_files(
os.path.join(self.data_dir, "room_filelist.txt")
)
data_batchlist, label_batchlist = [], []
for f in all_files:
data, label = _load_data_file(os.path.join(BASE_DIR, f))
data_batchlist.append(data)
label_batchlist.append(label)
data_batches = np.concatenate(data_batchlist, 0)
labels_batches = np.concatenate(label_batchlist, 0)
test_area = "Area_5"
train_idxs, test_idxs = [], []
for i, room_name in enumerate(room_filelist):
if test_area in room_name:
test_idxs.append(i)
else:
train_idxs.append(i)
if self.train:
self.points = data_batches[train_idxs, ...]
self.labels = labels_batches[train_idxs, ...]
else:
self.points = data_batches[test_idxs, ...]
self.labels = labels_batches[test_idxs, ...]
def __getitem__(self, idx):
pt_idxs = np.arange(0, self.num_points)
np.random.shuffle(pt_idxs)
current_points = torch.from_numpy(self.points[idx, pt_idxs].copy()).float()
current_labels = torch.from_numpy(self.labels[idx, pt_idxs].copy()).long()
return current_points, current_labels
def __len__(self):
return int(self.points.shape[0] * self.data_precent)
def set_num_points(self, pts):
self.num_points = pts
def randomize(self):
pass
if __name__ == "__main__":
dset = Indoor3DSemSeg(16, "./", train=True)
print(dset[0])
print(len(dset))
dloader = torch.utils.data.DataLoader(dset, batch_size=32, shuffle=True)
for i, data in enumerate(dloader, 0):
inputs, labels = data
if i == len(dloader) - 1:
print(inputs.size())
|