dmlc--dgl
Tong He
558673e139
* commit patch * commit patch * pointnet basic * fix data * reorg * reorg * temp status * remove validate set * add partseg data and model * partseg miou * clean up * fix loss * network definition match paper * fix * fix miou * update data format * fix * fix * working pointnet ssg cls * avoid some pytorch bug * fix script * update hyperparams * add msg module * try different dataset * update new dataset info * quick fix to subgraph * fix speed * update training * update * fix bs * update docstring * update * update * remove parallel reduction in fps * switch to kernel fps, training is 30% faster Co-authored-by: Ubuntu <ubuntu@ip-172-31-20-181.us-west-2.compute.internal>
56 行
1.8 KiB
Python
56 行
1.8 KiB
Python
import numpy as np
|
|
from torch.utils.data import Dataset
|
|
|
|
class ModelNet(object):
|
|
def __init__(self, path, num_points):
|
|
import h5py
|
|
self.f = h5py.File(path)
|
|
self.num_points = num_points
|
|
|
|
self.n_train = self.f['train/data'].shape[0]
|
|
self.n_valid = int(self.n_train / 5)
|
|
self.n_train -= self.n_valid
|
|
self.n_test = self.f['test/data'].shape[0]
|
|
|
|
def train(self):
|
|
return ModelNetDataset(self, 'train')
|
|
|
|
def valid(self):
|
|
return ModelNetDataset(self, 'valid')
|
|
|
|
def test(self):
|
|
return ModelNetDataset(self, 'test')
|
|
|
|
class ModelNetDataset(Dataset):
|
|
def __init__(self, modelnet, mode):
|
|
super(ModelNetDataset, self).__init__()
|
|
self.num_points = modelnet.num_points
|
|
self.mode = mode
|
|
|
|
if mode == 'train':
|
|
self.data = modelnet.f['train/data'][:modelnet.n_train]
|
|
self.label = modelnet.f['train/label'][:modelnet.n_train]
|
|
elif mode == 'valid':
|
|
self.data = modelnet.f['train/data'][modelnet.n_train:]
|
|
self.label = modelnet.f['train/label'][modelnet.n_train:]
|
|
elif mode == 'test':
|
|
self.data = modelnet.f['test/data'].value
|
|
self.label = modelnet.f['test/label'].value
|
|
|
|
def translate(self, x, scale=(2/3, 3/2), shift=(-0.2, 0.2)):
|
|
xyz1 = np.random.uniform(low=scale[0], high=scale[1], size=[3])
|
|
xyz2 = np.random.uniform(low=shift[0], high=shift[1], size=[3])
|
|
x = np.add(np.multiply(x, xyz1), xyz2).astype('float32')
|
|
return x
|
|
|
|
def __len__(self):
|
|
return self.data.shape[0]
|
|
|
|
def __getitem__(self, i):
|
|
x = self.data[i][:self.num_points]
|
|
y = self.label[i]
|
|
if self.mode == 'train':
|
|
x = self.translate(x)
|
|
np.random.shuffle(x)
|
|
return x, y
|