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])
