def process(self): progress = Progress('Processing', end=self.num_graphs, type='') train_data_list, test_data_list = [], [] for y, category in enumerate(self.categories): path = osp.join(self.raw_dir, category) for filename in glob.glob('{}/train/*.off'.format(path)): data = read_off(filename) data.y = torch.LongTensor([y]) train_data_list.append(data) progress.inc() for filename in glob.glob('{}/test/*.off'.format(path)): data = read_off(filename) data.y = torch.LongTensor([y]) test_data_list.append(data) progress.inc() train_dataset, train_slices = collate_to_set(train_data_list) torch.save((train_dataset, train_slices), self._processed_files[0]) test_dataset, test_slices = collate_to_set(test_data_list) torch.save((test_dataset, test_slices), self._processed_files[1]) progress.success()
def process(self): data_list = [] for ip_path in glob.glob('{}/*.off'.format(self.raw_paths[0])): ip_data = read_off(ip_path) shape_id = osp.basename(ip_path).rsplit('.', 1)[0] gt_path = osp.join(self.raw_paths[1], shape_id + '.off') gt_data = read_off(gt_path) data = Data() data.shape_id = torch.tensor( [int(shape_id.split('_')[1]), int(shape_id.split('_')[2])]) data.x = ip_data.pos data.gt_x = gt_data.pos data.face = ip_data.face data.gt_face = gt_data.face # if int(shape_id) < 10: data_list.append(data) if self.pre_filter is not None: data_list = [d for d in data_list if self.pre_filter(d)] if self.pre_transform is not None: data_list = [self.pre_transform(d) for d in data_list] val_split = 0.15 val_size = int(len(data_list) * val_split) torch.save(self.collate(data_list[:len(data_list) - val_size]), self.processed_paths[0]) torch.save(self.collate(data_list[len(data_list) - val_size:]), self.processed_paths[1])
def process(self): data_list = [] for off_path in glob.glob('{}/*.off'.format(self.raw_paths[0])): data = read_off(off_path) shape_id = osp.basename(off_path).rsplit('.', 1)[0] label_path = osp.join(self.raw_paths[1], shape_id + '.seg') data.y = read_txt_array(label_path) - 1 # start from 0 if self.classification is not None: if self.classification in data.y: data.y = torch.tensor([1]) else: data.y = torch.tensor([0]) data.shape_id = torch.tensor([int(shape_id)]) # if int(shape_id) < 10: data_list.append(data) if self.pre_filter is not None: data_list = [d for d in data_list if self.pre_filter(d)] if self.pre_transform is not None: data_list = [self.pre_transform(d) for d in data_list] val_split = 0.15 val_size = int(len(data_list) * val_split) torch.save(self.collate(data_list[:len(data_list) - val_size]), self.processed_paths[0]) torch.save(self.collate(data_list[len(data_list) - val_size:]), self.processed_paths[1])
def get_pcd(self, cat_num, model_num): path = os.path.join(self.raw_dir, self.categories[cat_num], self.dataset, self.get_model_name(cat_num, model_num)) data = read_off(path) data.y = torch.tensor([cat_num, model_num]) if self.pre_filter is not None: data = self.pre_filter(data) if self.pre_transform is not None: data = self.pre_transform(data) return data
def process_set(self, dataset): categories = glob.glob(osp.join(self.raw_dir, '*', '')) categories = sorted([x.split(os.sep)[-2] for x in categories]) data_list = [] for target, category in enumerate(categories): folder = osp.join(self.raw_dir, category, dataset) paths = glob.glob('{}/{}_*.off'.format(folder, category)) for path in paths: data = read_off(path) data.y = torch.tensor([target]) pos = data.pos face = data.face.contiguous() - 1 random_list = sorted(random.sample(range(int(face.size(1)) - 1), int(face.size(1) / 10))) face = face[:, random_list] face_set_a = set(face[0, :]) face_set_b = set(face[1, :]) face_set_c = set(face[2, :]) pos_list = sorted(list(face_set_a | face_set_b | face_set_c)) for i, _ in enumerate(pos_list): pos_list[i] = int(pos_list[i]) dict_key = {} for i, j in enumerate(pos_list): dict_key[int(j)] = i for i in range(face.size(1)): face[0][i] = dict_key[int(face[0][i])] face[1][i] = dict_key[int(face[1][i])] face[2][i] = dict_key[int(face[2][i])] pos = pos[pos_list] assert pos.size(1) == 3 and face.size(0) == 3 edge_index = torch.cat([face[:2], face[1:], face[::2]], dim=1) edge_index = to_undirected(edge_index, num_nodes=self.num) data.pos = pos data.face = face data.edge_index = edge_index data_list.append(data) if self.pre_filter is not None: data_list = [d for d in data_list if self.pre_filter(d)] if self.pre_transform is not None: data_list = [self.pre_transform(d) for d in data_list] return self.collate(data_list)
def process(self): ref_data = read_off( osp.join(self.raw_paths[0], 'null', '{}.off'.format(self.cat))) train_list = [] name = '{}_{}_*.off'.format(self.part, self.cat) paths = glob.glob(osp.join(self.raw_paths[0], self.part, name)) paths = [path[:-4] for path in paths] paths = sorted(paths, key=lambda e: (len(e), e)) for path in paths: data = read_off('{}.off'.format(path)) y = read_txt_array('{}.baryc_gt'.format(path)) data.y = y[:, 0].to(torch.long) - 1 data.y_baryc = y[:, 1:] train_list.append(data) test_list = [] name = '{}_{}_*.off'.format(self.part, self.cat) paths = glob.glob(osp.join(self.raw_paths[1], self.part, name)) paths = [path[:-4] for path in paths] paths = sorted(paths, key=lambda e: (len(e), e)) for path in paths: test_list.append(read_off('{}.off'.format(path))) if self.pre_filter is not None: train_list = [d for d in train_list if self.pre_filter(d)] test_list = [d for d in test_list if self.pre_filter(d)] if self.pre_transform is not None: ref_data = self.pre_transform(ref_data) train_list = [self.pre_transform(d) for d in train_list] test_list = [self.pre_transform(d) for d in test_list] torch.save(ref_data, self.processed_paths[0]) torch.save(self.collate(train_list), self.processed_paths[1]) torch.save(self.collate(test_list), self.processed_paths[2])
def process_set(self, dataset): categories = glob.glob(osp.join(self.raw_dir, '*', '')) categories = sorted([x.split('/')[-2] for x in categories]) data_list = [] for target, category in enumerate(categories): folder = osp.join(self.raw_dir, category, dataset) paths = glob.glob('{}/{}_*.off'.format(folder, category)) for path in paths: data = read_off(path) data.y = torch.tensor([target]) data_list.append(data) if self.pre_transform is not None: data_list = [self.pre_transform(d) for d in data_list] return self.collate(data_list)
def process_set(self, dataset): if self.cache_pcds: data_list = [] for target, category in enumerate(self.categories): folder = os.path.join(self.raw_dir, category, dataset) paths = sorted( glob.glob('{}/{}_*.off'.format(folder, category))) for path in paths: data = read_off(path) model_num = int(os.path.basename(path)[-8:-4]) data.y = torch.tensor([target, model_num]) data_list.append(data) if self.pre_filter is not None: data_list = [d for d in data_list if self.pre_filter(d)] if self.pre_transform is not None: data_list = [self.pre_transform(d) for d in data_list] return self.collate(data_list) else: pass
def process_set(self, dataset): categories = glob.glob(osp.join(self.raw_dir, '*', '')) categories = sorted([x.split(os.sep)[-2] for x in categories]) data_list = [] for target, category in enumerate(categories): folder = osp.join(self.raw_dir, category, dataset) paths = glob.glob('{}/*.off'.format(folder)) for path in paths: data = read_off(path) data.pos = data.pos - data.pos.mean(dim=0, keepdim=True) data.y = torch.tensor([target]) data_list.append(data) if self.pre_filter is not None: data_list = [d for d in data_list if self.pre_filter(d)] if self.pre_transform is not None: data_list = [self.pre_transform(d) for d in data_list] return self.collate(data_list)
def process_set(self, dataset): categories = glob.glob(osp.join(self.raw_dir, '*', '')) categories = sorted([x.split(os.sep)[-2] for x in categories]) poisson = Poisson(1024) normalize = normalization.Normalize() data_list = [] for target, category in enumerate(categories): folder = osp.join(self.raw_dir, category, dataset) paths = glob.glob('{}/{}_*.off'.format(folder, category)) for path in paths: data = read_off(path) data.y = torch.tensor([target]) data = normalize(data) data = poisson(data) data_list.append(data) if self.pre_filter is not None: data_list = [d for d in data_list if self.pre_filter(d)] if self.pre_transform is not None: data_list = [self.pre_transform(d) for d in data_list] return self.collate(data_list)