| 214 | |
| 215 | class ModelNetDataset(Dataset): |
| 216 | def __init__(self, data_root, num_category, num_points, split='train'): |
| 217 | self.root = data_root |
| 218 | self.npoints = num_points |
| 219 | self.uniform = True |
| 220 | self.use_normals = True |
| 221 | self.num_category = num_category |
| 222 | |
| 223 | if self.num_category == 10: |
| 224 | self.catfile = os.path.join(self.root, 'modelnet10_shape_names.txt') |
| 225 | else: |
| 226 | self.catfile = os.path.join(self.root, 'modelnet40_shape_names.txt') |
| 227 | |
| 228 | self.cat = [line.rstrip() for line in open(self.catfile)] |
| 229 | self.classes = dict(zip(self.cat, range(len(self.cat)))) |
| 230 | |
| 231 | shape_ids = {} |
| 232 | if self.num_category == 10: |
| 233 | shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet10_train.txt'))] |
| 234 | shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet10_test.txt'))] |
| 235 | else: |
| 236 | shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_train.txt'))] |
| 237 | shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_test.txt'))] |
| 238 | |
| 239 | assert (split == 'train' or split == 'test') |
| 240 | shape_names = ['_'.join(x.split('_')[0:-1]) for x in shape_ids[split]] |
| 241 | self.datapath = [(shape_names[i], os.path.join(self.root, shape_names[i], shape_ids[split][i]) + '.txt') for i |
| 242 | in range(len(shape_ids[split]))] |
| 243 | print('The size of %s data is %d' % (split, len(self.datapath))) |
| 244 | |
| 245 | if self.uniform: |
| 246 | self.data_path = os.path.join(data_root, |
| 247 | 'modelnet%d_%s_%dpts_fps.dat' % (self.num_category, split, self.npoints)) |
| 248 | else: |
| 249 | self.data_path = os.path.join(data_root, |
| 250 | 'modelnet%d_%s_%dpts.dat' % (self.num_category, split, self.npoints)) |
| 251 | |
| 252 | print('Load processed data from %s...' % self.data_path) |
| 253 | with open(self.data_path, 'rb') as f: |
| 254 | self.list_of_points, self.list_of_labels = pickle.load(f) |
| 255 | |
| 256 | def __len__(self): |
| 257 | return len(self.datapath) |