steffen/Labeling-Primitives-With-Point2Net
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

New File Types, automatic detection and header parameters

Browse Source
This commit is contained in:
Si11ium
2019-07-30 15:11:11 +02:00
parent 4e38de9a5b
commit 0b9d03a25d
2 changed files with 26 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
dataset/shapenet.py
View File

@ -16,7 +16,9 @@ class CustomShapeNet(InMemoryDataset):
categories = {key: val for val, key in enumerate(['Box', 'Cone', 'Cylinder', 'Sphere'])}
def __init__(self, root, train=True, transform=None, pre_filter=None, pre_transform=None, **kwargs):
def __init__(self, root, train=True, transform=None, pre_filter=None, pre_transform=None,
headers=True, **kwargs):
self.has_headers = headers
super(CustomShapeNet, self).__init__(root, transform, pre_transform, pre_filter)
path = self.processed_paths[0] if train else self.processed_paths[1]
self.data, self.slices = torch.load(path)
@ -64,20 +66,26 @@ class CustomShapeNet(InMemoryDataset):
for pointcloud in tqdm(os.scandir(os.path.join(self.raw_dir, setting))):
if not os.path.isdir(pointcloud):
continue
for element in glob.glob(os.path.join(pointcloud.path, '*.dat')):
if os.path.split(element)[-1] not in ['pc.dat']:
paths = list()
for ext in ['dat', 'xyz']:
paths.extend(glob.glob(os.path.join(pointcloud.path, f'*.{ext}')))
for element in paths:
if all([x not in os.path.split(element)[-1] for x in ['pc.dat', 'pc.xyz']]):
# Assign training data to the data container
# Following the original logic;
# y should be the label;
# pos should be the six dimensional vector describing: !its pos not points!!
# x,y,z,x_rot,y_rot,z_rot
y_raw = os.path.splitext(element)[0].split('_')[-2]
# Get the y - Label
y_raw = next(i for i, v in enumerate(self.categories.keys()) if v.lower() in element.lower())
# y_raw = os.path.splitext(element)[0].split('_')[-2]
with open(element,'r') as f:
headers = f.__next__()
# Check if there are no useable nodes in this file, header says 0.
if not int(headers.rstrip().split(delimiter)[0]):
continue
# Get the y - Label
if self.has_headers:
headers = f.__next__()
# Check if there are no useable nodes in this file, header says 0.
if not int(headers.rstrip().split(delimiter)[0]):
continue
# Iterate over all rows
src = [[float(x) if x not in ['-nan(ind)', 'nan(ind)'] else 0
@ -87,7 +95,7 @@ class CustomShapeNet(InMemoryDataset):
continue
# pos = points[:, :3]
# norm = points[:, 3:]
y_all = [self.categories[y_raw]] * points.shape[0]
y_all = [y_raw] * points.shape[0]
y = torch.as_tensor(y_all, dtype=torch.int)
# points = torch.as_tensor(points, dtype=torch.float)
# norm = torch.as_tensor(norm, dtype=torch.float)
@ -115,10 +123,10 @@ class ShapeNetPartSegDataset(Dataset):
Resample raw point cloud to fixed number of points.
Map raw label from range [1, N] to [0, N-1].
"""
def __init__(self, root_dir, train=True, transform=None, npoints=1024):
def __init__(self, root_dir, train=True, transform=None, npoints=1024, headers=True):
super(ShapeNetPartSegDataset, self).__init__()
self.npoints = npoints
self.dataset = CustomShapeNet(root=root_dir, train=train, transform=transform)
self.dataset = CustomShapeNet(root=root_dir, train=train, transform=transform, headers=headers)
def __getitem__(self, index):
data = self.dataset[index]