forked from sudhir5595/Mesh_Neural_Network
-
Notifications
You must be signed in to change notification settings - Fork 0
/
dataset.py
71 lines (53 loc) · 1.62 KB
/
dataset.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
import numpy as np
import torch
import os
from torch.utils.data import Dataset
from sklearn.preprocessing import LabelEncoder
from input import input_stl
class MeshData(Dataset):
def __init__(self,root_dir):
self.classes = set()
self.X = []
self.Y = []
self.final_dataset = []
self.classes_codec = LabelEncoder()
lst_of_classes = os.listdir(root_dir)
lst_of_classes.sort()
self.classes_codec.fit(lst_of_classes)
for x in lst_of_classes:
path = root_dir + '/' + x + '/train'
lst_of_objects = os.listdir(path)
for y in lst_of_objects:
file_path = path + '/' + y
dict_para = input_stl(file_path)
if dict_para == False:
print('Input file',y,'has problems')
continue
else:
neigh = dict_para["neigh_index"]
corner = dict_para["corners"]
center = dict_para["centroids"]
normal = dict_para["normals"]
self.X.append(np.concatenate((center,corner,normal,neigh),axis=1))
self.Y.append((self.one_hot_encode(self.classes_codec,[x])))
self.final_dataset = [self.X,self.Y]
def one_hot_encode(self,codec,values):
value_idxs = codec.transform(values)
val,idx = torch.max(torch.eye(len(codec.classes_))[value_idxs],1)
return torch.LongTensor(idx)
def __len__(self):
return len(final_dataset)
def __getitem__(self,idx):
return torch.from_numpy(self.X[idx]),self.Y[idx]
#DataObject = MeshData('/home/prathmesh/Desktop/SoC-2020/stl/b')
#print(DataObject.X.type())
#print(len(DataObject.X))
#x,y = DataObject[0]
#print(x.shape)
#print(y)
#x,y = DataObject[1]
#print(x.shape)
#print(y)
#print(DataObject[1].shape)
#print(y.shape)
#print(DataObject[1].shape)