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

dataset modification

Browse Source
This commit is contained in:
Si11ium
2020-06-19 19:00:07 +02:00
parent b3c67bab40
commit a19bd9cafd
4 changed files with 36 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
datasets/shapenet.py
View File

@ -8,7 +8,6 @@ from collections import defaultdict
import os
from torch.utils.data import Dataset
from tqdm import tqdm
import glob
import torch
from torch_geometric.data import InMemoryDataset
@ -45,7 +44,9 @@ class CustomShapeNet(InMemoryDataset):
assert mode in self.modes.keys(), f'"mode" must be one of {self.modes.keys()}'
# Set the Dataset Parameters
self.collate_per_segment, self.mode, self.refresh = collate_per_segment, mode, refresh
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)
@ -57,15 +58,15 @@ class CustomShapeNet(InMemoryDataset):
return [f'{self.mode}.pt']
def check_and_resolve_cloud_count(self):
if self.raw_dir.exists():
dir_count = len([name for name in os.listdir(self.raw_dir) if os.path.isdir(os.path.join(self.raw_dir, name))])
if (self.raw_dir / self.mode).exists():
file_count = len([cloud for cloud in (self.raw_dir / self.mode).iterdir() if cloud.is_file()])
if dir_count:
print(f'{dir_count} folders have been found....')
return dir_count
if file_count:
print(f'{file_count} files have been found....')
return file_count
else:
warn(ResourceWarning("No raw pointclouds have been found. Was this intentional?"))
return dir_count
return file_count
warn(ResourceWarning("The raw data folder does not exist. Was this intentional?"))
return -1
@ -99,7 +100,7 @@ class CustomShapeNet(InMemoryDataset):
continue
return data, slices
def _transform_and_filter(self, data):
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)
@ -133,7 +134,9 @@ class CustomShapeNet(InMemoryDataset):
src[key] = torch.tensor(values, dtype=torch.double).squeeze()
if not self.collate_per_segment:
src = dict(all=torch.stack([x for x in src.values()]))
src = dict(
all=torch.cat(tuple(src.values()))
)
for key, values in src.items():
try:
@ -157,17 +160,18 @@ class CustomShapeNet(InMemoryDataset):
if self.collate_per_segment:
data = Data(**attr_dict)
else:
if not data:
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)
data = self._transform_and_filter(data)
# data = self._pre_transform_and_filter(data)
if self.collate_per_segment:
datasets[self.mode].append(data)
if not self.collate_per_segment:
# Todo: What is this?
# 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()}))
if datasets[self.mode]:
@ -198,8 +202,10 @@ class ShapeNetPartSegDataset(Dataset):
# Resample to fixed number of points
try:
npoints = self.npoints if self.mode != 'predict' else data.pos.shape[0]
choice = np.random.choice(data.pos.shape[0], npoints, replace=False if self.mode == 'predict' else True)
npoints = self.npoints if self.mode != DataSplit.predict else data.pos.shape[0]
choice = np.random.choice(data.pos.shape[0], npoints,
replace=False if self.mode == DataSplit.predict else True
)
except ValueError:
choice = []