steffen/point_to_primitive
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

New Model running

Browse Source
This commit is contained in:
Si11ium
2020-06-23 14:37:34 +02:00
parent a19bd9cafd
commit 1033b26195
12 changed files with 173 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

107
datasets/shapenet.py
View File

@ -1,8 +1,7 @@
from pathlib import Path
from typing import Union
from warnings import warn
import numpy as np
from collections import defaultdict
import os
@ -13,7 +12,7 @@ import torch
from torch_geometric.data import InMemoryDataset
from torch_geometric.data import Data
from utils.project_settings import Classes, DataSplit
from utils.project_settings import Classes, DataSplit, ClusterTypes
def save_names(name_list, path):
@ -23,10 +22,23 @@ def save_names(name_list, path):
class CustomShapeNet(InMemoryDataset):
categories = {key: val for val, key in Classes().items()}
modes = {key: val for val, key in DataSplit().items()}
name = 'CustomShapeNet'
def download(self):
pass
@property
def categories(self):
return {key: val for val, key in self.classes.items()}
@property
def modes(self):
return {key: val for val, key in DataSplit().items()}
@property
def cluster_types(self):
return {key: val for val, key in ClusterTypes().items()}
@property
def raw_dir(self):
return self.root / 'raw'
@ -40,14 +52,21 @@ class CustomShapeNet(InMemoryDataset):
return self.root / 'processed'
def __init__(self, root_dir, collate_per_segment=True, mode='train', transform=None, pre_filter=None,
pre_transform=None, refresh=False, with_normals=False):
assert mode in self.modes.keys(), f'"mode" must be one of {self.modes.keys()}'
pre_transform=None, refresh=False, cluster_type: Union[str, None] = '',
poly_as_plane=False):
assert mode in self.modes.keys(), \
f'"mode" must be one of {self.modes.keys()}'
assert cluster_type in self.cluster_types.keys() or cluster_type is None, \
f'"cluster_type" must be one of {self.cluster_types.keys()} or None, but was: {cluster_type}'
# Set the Dataset Parameters
self.cluster_type = cluster_type if cluster_type else 'pc'
self.classes = Classes()
self.poly_as_plane = poly_as_plane
self.collate_per_segment = collate_per_segment
self.mode = mode
self.refresh = refresh
self.with_normals = with_normals
root_dir = Path(root_dir)
super(CustomShapeNet, self).__init__(root_dir, transform, pre_transform, pre_filter)
self.data, self.slices = self._load_dataset()
@ -72,7 +91,7 @@ class CustomShapeNet(InMemoryDataset):
@property
def num_classes(self):
return len(self.categories)
return len(self.categories) if self.poly_as_plane else (len(self.categories) - 2)
def _load_dataset(self):
data, slices = None, None
@ -101,22 +120,17 @@ class CustomShapeNet(InMemoryDataset):
return data, slices
def _pre_transform_and_filter(self, data):
# ToDo: ANy filter to apply? Then do it here.
if self.pre_filter is not None and not self.pre_filter(data):
data = self.pre_filter(data)
raise NotImplementedError
# ToDo: ANy transformation to apply? Then do it here.
if self.pre_transform is not None:
data = self.pre_transform(data)
raise NotImplementedError
return data
def process(self, delimiter=' '):
datasets = defaultdict(list)
path_to_clouds = self.raw_dir / self.mode
for pointcloud in tqdm(path_to_clouds.glob('*.xyz')):
if 'grid' not in pointcloud.name:
if self.cluster_type not in pointcloud.name:
continue
data = None
@ -129,15 +143,32 @@ class CustomShapeNet(InMemoryDataset):
vals = [float(x) if x not in ['-nan(ind)', 'nan(ind)'] else 0 for x in vals]
src[vals[-1]].append(vals)
# Switch from un-pickable Defaultdict to Standard Dict
src = dict(src)
# Transform the Dict[List] to Dict[torch.Tensor]
for key, values in src.items():
src[key] = torch.tensor(values, dtype=torch.double).squeeze()
# Screw the Sorting and make it a FullCloud rather than a seperated
if not self.collate_per_segment:
src = dict(
all=torch.cat(tuple(src.values()))
)
# Transform Box and Polytope to Plane if poly_as_plane is set
for key, tensor in src.items():
if tensor.ndim == 1:
if all([x == 0 for x in tensor]):
continue
tensor = tensor.unsqueeze(0)
if self.poly_as_plane:
tensor[:, -2][tensor[:, -2] == float(self.classes.Plane)] = 4.0
tensor[:, -2][tensor[:, -2] == float(self.classes.Box)] = 4.0
tensor[:, -2][tensor[:, -2] == float(self.classes.Polytope)] = 4.0
tensor[:, -2][tensor[:, -2] == self.classes.Torus] = 3.0
src[key] = tensor
for key, values in src.items():
try:
points = values[:, :-2]
@ -147,36 +178,35 @@ class CustomShapeNet(InMemoryDataset):
y_c = torch.as_tensor(values[:, -1], dtype=torch.long)
####################################
# This is where you define the keys
attr_dict = dict(y=y, y_c=y_c)
if self.with_normals:
pos = points[:, :6]
norm = None
attr_dict.update(pos=pos, norm=norm)
if not self.with_normals:
pos = points[:, :3]
norm = points[:, 3:6]
attr_dict.update(pos=pos, norm=norm)
attr_dict = dict(
y=y,
y_c=y_c,
pos=points[:, :3],
norm=points[:, 3:6]
)
####################################
if self.collate_per_segment:
data = Data(**attr_dict)
else:
if data is None:
data = defaultdict(list)
# points=points, norm=points[:, 3:]
for key, val in attr_dict.items():
data[key].append(val)
# data = Data(**data)
for attr_key, val in attr_dict.items():
data[attr_key].append(val)
# data = self._pre_transform_and_filter(data)
if self.collate_per_segment:
datasets[self.mode].append(data)
if not self.collate_per_segment:
# This is just to be sure, but should not be needed, since src[all] == all there is in this cloud
datasets[self.mode].append(Data(**{key: torch.cat(data[key]) for key in data.keys()}))
# This is just to be sure, but should not be needed, since src[all] == all
raise TypeError('FIX THIS')
# old Code
# datasets[self.mode].append(Data(**{key: torch.cat(data[key]) for key in data.keys()}))
if datasets[self.mode]:
os.makedirs(self.processed_dir, exist_ok=True)
torch.save(self.collate(datasets[self.mode]), self.processed_paths[0])
collated_dataset = self.collate(datasets[self.mode])
torch.save(collated_dataset, self.processed_paths[0])
def __repr__(self):
return f'{self.__class__.__name__}({len(self)})'
@ -190,17 +220,18 @@ class ShapeNetPartSegDataset(Dataset):
name = 'ShapeNetPartSegDataset'
def __init__(self, root_dir, npoints=1024, mode='train', **kwargs):
def __init__(self, root_dir, mode='train', **kwargs):
super(ShapeNetPartSegDataset, self).__init__()
self.mode = mode
kwargs.update(dict(root_dir=root_dir, mode=self.mode))
self.npoints = npoints
# self.npoints = npoints
self.dataset = CustomShapeNet(**kwargs)
def __getitem__(self, index):
data = self.dataset[index]
# Resample to fixed number of points
'''
try:
npoints = self.npoints if self.mode != DataSplit.predict else data.pos.shape[0]
choice = np.random.choice(data.pos.shape[0], npoints,
@ -209,16 +240,16 @@ class ShapeNetPartSegDataset(Dataset):
except ValueError:
choice = []
pos, norm, y = data.pos[choice, :], data.norm[choice], data.y[choice]
pos, norm, y = data.pos[choice, :], data.norm[choice], data.y[choice]
# y -= 1 if self.num_classes() in y else 0 # Map label from [1, C] to [0, C-1]
data = Data(**dict(pos=pos, # torch.Tensor (n, 3/6)
y=y, # torch.Tensor (n,)
norm=norm # torch.Tensor (n, 3/0)
y=y, # torch.Tensor (n,)
norm=norm # torch.Tensor (n, 3/0)
)
)
)
'''
return data
def __len__(self):