def test_read_off(): data = read_off(osp.join('test', 'io', 'example1.off')) assert len(data) == 2 assert data.pos.tolist() == [[0, 0, 0], [0, 1, 0], [1, 0, 0], [1, 1, 0]] assert data.face.tolist() == [[0, 1], [1, 2], [2, 3]] data = read_off(osp.join('test', 'io', 'example2.off')) assert len(data) == 2 assert data.pos.tolist() == [[0, 0, 0], [0, 1, 0], [1, 0, 0], [1, 1, 0]] assert data.face.tolist() == [[0, 0], [1, 2], [2, 3]]
def test_write_off(): pos = torch.tensor([[0, 0, 0], [0, 1, 0], [1, 0, 0], [1, 1, 0]]) face = torch.tensor([[0, 1], [1, 2], [2, 3]]) name = str(random.randrange(sys.maxsize)) path = osp.join('/', 'tmp', f'{name}.off') write_off(Data(pos=pos, face=face), path) data = read_off(path) os.unlink(path) assert data.pos.tolist() == pos.tolist() assert data.face.tolist() == face.tolist()
def process(self): ref_data = read_off( osp.join(self.raw_paths[0], 'null', f'{self.cat}.off')) train_list = [] name = f'{self.part}_{self.cat}_*.off' 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(f'{path}.off') y = read_txt_array(f'{path}.baryc_gt') data.y = y[:, 0].to(torch.long) - 1 data.y_baryc = y[:, 1:] train_list.append(data) test_list = [] name = f'{self.part}_{self.cat}_*.off' 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(f'{path}.off')) 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 = self.raw_file_names # 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]) 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: str): 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(f'{folder}/*.off') 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(self): path = osp.join(self.root, model_dir, 'mesh{0:03d}.off') txt_path = osp.join(self.root, data_dir, 'mesh{0:03d}.txt') train_list = [] test_list = [] for i in range(0, 72): if i == 51: continue data = read_off(path.format(i)) data.y = torch.tensor([i % 10], dtype=torch.long) if self.pre_filter is not None and not self.pre_filter(data): continue if self.pre_transform is not None: data = self.pre_transform(data) data.x = torch.tensor(np.loadtxt(txt_path.format(i)), dtype=torch.float32) Winn = sio.loadmat(txt_path.format(i)[:-3] + 'mat') data.V = torch.tensor(Winn['V'], dtype=torch.float32) data.A = torch.tensor(Winn['A'], dtype=torch.float32) data.D = torch.tensor(Winn['D'], dtype=torch.float32) data.clk = torch.tensor(Winn['clk'], dtype=torch.float32) data.name = int(path.format(i).split('/')[-1].split('.')[0][4:7]) data.flag = 0 data.map = torch.tensor(np.loadtxt(path.format(i)[:-3] + 'vts') - 1, dtype=torch.long) if i < 62 and i != 51: train_list.append(data) elif i != 51: test_list.append(data) torch.save(self.collate(train_list), self.processed_paths[0]) torch.save(self.collate(test_list), self.processed_paths[1])