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

File based header detection, collate_per_PC training.

Browse Source
This commit is contained in:
Si11ium 2019-07-31 12:55:47 +02:00
parent 0b9d03a25d
commit 47a76dc978
6 changed files with 218 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -127,3 +127,4 @@ dmypy.json
/data/
/checkpoint/
/shapenet/
/vis/

1
.idea/pointnet2-pytorch.iml generated
View File

@ -4,6 +4,7 @@
<content url="file://$MODULE_DIR$">
<excludeFolder url="file://$MODULE_DIR$/data" />
<excludeFolder url="file://$MODULE_DIR$/net" />
<excludeFolder url="file://$MODULE_DIR$/vis" />
</content>
<orderEntry type="jdk" jdkName="Python 3.7 (torch)" jdkType="Python SDK" />
<orderEntry type="sourceFolder" forTests="false" />

197
dataset/shapenet.py
View File

@ -16,11 +16,12 @@ 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,
def __init__(self, root, collate_per_segment=True, train=True, transform=None, pre_filter=None, pre_transform=None,
headers=True, **kwargs):
self.has_headers = headers
self.collate_per_element = collate_per_segment
super(CustomShapeNet, self).__init__(root, transform, pre_transform, pre_filter)
path = self.processed_paths[0] if train else self.processed_paths[1]
path = self.processed_paths[0] if train else self.processed_paths[-1]
self.data, self.slices = torch.load(path)
print("Initialized")
@ -57,16 +58,37 @@ class CustomShapeNet(InMemoryDataset):
continue
return data, slices
def _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=' '):
# idx = self.categories[self.category]
# paths = [osp.join(path, idx) for path in self.raw_paths]
datasets = defaultdict(list)
for idx, setting in enumerate(self.raw_file_names):
for pointcloud in tqdm(os.scandir(os.path.join(self.raw_dir, setting))):
path_to_clouds = os.path.join(self.raw_dir, setting)
if '.headers' in os.listdir(path_to_clouds):
self.has_headers = True
elif 'no.headers' in os.listdir(path_to_clouds):
self.has_headers = False
else:
pass
for pointcloud in tqdm(os.scandir(path_to_clouds)):
if not os.path.isdir(pointcloud):
continue
paths = list()
data, paths = None, list()
for ext in ['dat', 'xyz']:
paths.extend(glob.glob(os.path.join(pointcloud.path, f'*.{ext}')))
for element in paths:
@ -99,17 +121,19 @@ class CustomShapeNet(InMemoryDataset):
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)
data = Data(y=y, pos=points[:, :3])
if self.collate_per_element:
data = Data(y=y, pos=points[:, :3])
else:
if not data:
data = defaultdict(list)
for key, val in dict(y=y, pos= points[:, :3]).items():
data[key].append(val)
# , points=points, norm=points[:3], )
# 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
datasets[setting].append(data)
data = self._transform_and_filter(data)
if self.collate_per_element:
datasets[setting].append(data)
if not self.collate_per_element:
datasets[setting].append(Data(**{key: torch.cat(data[key]) for key in data.keys()}))
os.makedirs(self.processed_dir, exist_ok=True)
torch.save(self.collate(datasets[setting]), self.processed_paths[idx])
@ -123,10 +147,151 @@ 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, headers=True):
def __init__(self, root_dir, collate_per_segment=True, 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, headers=headers)
self.dataset = CustomShapeNet(root=root_dir, collate_per_segment=collate_per_segment,
train=train, transform=transform, headers=headers)
def __getitem__(self, index):
data = self.dataset[index]
points, labels = data.pos, data.y
# Resample to fixed number of points
try:
choice = np.random.choice(points.shape[0], self.npoints, replace=True)
except ValueError:
choice = []
points, labels = points[choice, :], labels[choice]
labels -= 1 if self.num_classes() in labels else 0 # Map label from [1, C] to [0, C-1]
sample = {
'points': points, # torch.Tensor (n, 3)
'labels': labels # torch.Tensor (n,)
}
return sample
def __len__(self):
return len(self.dataset)
def num_classes(self):
return self.dataset.num_classes
class PredictionShapeNet(InMemoryDataset):
categories = {key: val for val, key in enumerate(['Box', 'Cone', 'Cylinder', 'Sphere'])}
def __init__(self, root, transform=None, pre_filter=None, pre_transform=None,
headers=True, **kwargs):
self.has_headers = headers
super(PredictionShapeNet, self).__init__(root, transform, pre_transform, pre_filter)
path = self.processed_paths[0]
self.data, self.slices = torch.load(path)
print("Initialized")
@property
def raw_file_names(self):
# Maybe add more data like validation sets
return ['predict']
@property
def processed_file_names(self):
return [f'{x}.pt' for x in self.raw_file_names]
def download(self):
dir_count = len([name for name in os.listdir(self.raw_dir) if os.path.isdir(os.path.join(self.raw_dir, name))])
print(f'{dir_count} folders have been found....')
if dir_count:
return dir_count
raise IOError("No raw pointclouds have been found.")
@property
def num_classes(self):
return len(self.categories)
def _load_dataset(self):
data, slices = None, None
while True:
try:
filepath = os.path.join(self.root, self.processed_dir, f'{"train" if self.train else "test"}.pt')
data, slices = torch.load(filepath)
print('Dataset Loaded')
break
except FileNotFoundError:
self.process()
continue
return data, slices
def process(self, delimiter=' '):
datasets = defaultdict(list)
for idx, setting in enumerate(self.raw_file_names):
path_to_clouds = os.path.join(self.raw_dir, setting)
if '.headers' in os.listdir(path_to_clouds):
self.has_headers = True
elif 'no.headers' in os.listdir(path_to_clouds):
self.has_headers = False
else:
pass
for pointcloud in tqdm(os.scandir(path_to_clouds)):
if not os.path.isdir(pointcloud):
continue
for extention in ['dat', 'xyz']:
file = os.path.join(pointcloud.path, f'pc.{extention}')
if not os.path.exists(file):
continue
with open(file, 'r') as f:
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
for x in line.rstrip().split(delimiter)[None:None]] for line in f if line != '']
points = torch.tensor(src, dtype=None).squeeze()
if not len(points.shape) > 1:
continue
# pos = points[:, :3]
# norm = points[:, 3:]
y_fake_all = [-1] * points.shape[0]
y = torch.as_tensor(y_fake_all, dtype=torch.int)
# points = torch.as_tensor(points, dtype=torch.float)
# norm = torch.as_tensor(norm, dtype=torch.float)
data = Data(y=y, pos=points[:, :3])
# , points=points, norm=points[:3], )
# 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
datasets[setting].append(data)
os.makedirs(self.processed_dir, exist_ok=True)
torch.save(self.collate(datasets[setting]), self.processed_paths[idx])
def __repr__(self):
return f'{self.__class__.__name__}({len(self)})'
class PredictNetPartSegDataset(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, transform=None, npoints=2048, headers=True):
super(PredictNetPartSegDataset, self).__init__()
self.npoints = npoints
self.dataset = PredictionShapeNet(root=root_dir, train=False, transform=transform, headers=headers)
def __getitem__(self, index):
data = self.dataset[index]

7
main.py
View File

@ -37,6 +37,7 @@ parser.add_argument('--batch_size', type=int, default=8, help='input batch size'
parser.add_argument('--test_per_batches', type=int, default=1000, help='run a test batch per training batches number')
parser.add_argument('--num_workers', type=int, default=4, help='number of data loading workers')
parser.add_argument('--headers', type=strtobool, default=True, help='if raw files come with headers')
parser.add_argument('--collate_per_segment', type=strtobool, default=True, help='whether to look at pointclouds or sub')
opt = parser.parse_args()
@ -68,10 +69,12 @@ if __name__ == '__main__':
train_transform = GT.Compose([GT.NormalizeScale(), RotTransform, TransTransform])
test_transform = GT.Compose([GT.NormalizeScale(), ])
dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, train=True, transform=train_transform, npoints=opt.npoints, headers=opt.headers)
dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, collate_per_segment=opt.collate_per_segment,
train=True, transform=train_transform, npoints=opt.npoints, headers=opt.headers)
dataLoader = DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers)
test_dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, train=False, transform=test_transform, npoints=opt.npoints, headers=opt.headers)
test_dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, collate_per_segment=opt.collate_per_segment,
train=False, transform=test_transform, npoints=opt.npoints, headers=opt.headers)
test_dataLoader = DataLoader(test_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers)
num_classes = dataset.num_classes()

53
vis/show_seg_res.py
View File

@ -1,11 +1,11 @@
# Warning: import open3d may lead crash, try to to import open3d first here
from view import view_points_labels
from vis.view import view_points_labels
import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../') # add project root directory
from dataset.shapenet import ShapeNetPartSegDataset
from dataset.shapenet import PredictNetPartSegDataset
from model.pointnet2_part_seg import PointNet2PartSegmentNet
import torch_geometric.transforms as GT
import torch
@ -17,9 +17,8 @@ import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default='data', help='dataset path')
parser.add_argument('--npoints', type=int, default=50, help='resample points number')
parser.add_argument('--model', type=str, default='./checkpoint/seg_model_Airplane_24.pth', help='model path')
parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_8.pth', help='model path')
parser.add_argument('--sample_idx', type=int, default=0, help='select a sample to segment and view result')
opt = parser.parse_args()
print(opt)
@ -29,9 +28,8 @@ if __name__ == '__main__':
print('Construct dataset ..')
test_transform = GT.Compose([GT.NormalizeScale(),])
test_dataset = ShapeNetPartSegDataset(
test_dataset = PredictNetPartSegDataset(
root_dir=opt.dataset,
train=False,
transform=test_transform,
npoints=opt.npoints
)
@ -49,7 +47,7 @@ if __name__ == '__main__':
# net = PointNetPartSegmentNet(num_classes)
net = PointNet2PartSegmentNet(num_classes)
net.load_state_dict(torch.load(opt.model))
net.load_state_dict(torch.load(opt.model, map_location=device.type))
net = net.to(device, dtype)
net.eval()
@ -104,30 +102,35 @@ if __name__ == '__main__':
# Get one sample and eval
sample = test_dataset[opt.sample_idx]
#sample = test_dataset[opt.sample_idx]
r_idx = np.random.randint(0, len(test_dataset), 20)
for idx in r_idx:
sample = test_dataset[int(idx)]
print('Eval test sample ..')
pred_label, gt_label = eval_sample(net, sample)
print('Eval done ..')
print('Eval test sample ..')
pred_label, gt_label = eval_sample(net, sample)
print('Eval done ..')
# Get sample result
print('Compute mIoU ..')
points = sample['points'].numpy()
pred_labels = pred_label.numpy()
gt_labels = gt_label.numpy()
diff_labels = label_diff(pred_labels, gt_labels)
# Get sample result
print('Compute mIoU ..')
points = sample['points'].numpy()
pred_labels = pred_label.numpy()
gt_labels = gt_label.numpy()
diff_labels = label_diff(pred_labels, gt_labels)
print('mIoU: ', compute_mIoU(pred_labels, gt_labels))
print('mIoU: ', compute_mIoU(pred_labels, gt_labels))
# View result
# View result
# print('View gt labels ..')
# view_points_labels(points, gt_labels)
# print('View gt labels ..')
# view_points_labels(points, gt_labels)
# print('View diff labels ..')
# view_points_labels(points, diff_labels)
print('View diff labels ..')
print(diff_labels)
view_points_labels(points, diff_labels)
print('View pred labels ..')
view_points_labels(points, pred_labels)
# print('View pred labels ..')
# print(pred_labels)
# view_points_labels(points, pred_labels)

3
vis/view.py
View File

@ -28,7 +28,7 @@ def view_points(points, colors=None):
'''
cloud = o3d.PointCloud()
cloud.points = o3d.Vector3dVector(points)
# frame = o3d.create_mesh_coordinate_frame(-1, -1, -1)
if colors is not None:
if isinstance(colors, np.ndarray):
cloud.colors = o3d.Vector3dVector(colors)
@ -37,6 +37,7 @@ def view_points(points, colors=None):
o3d.draw_geometries([cloud])
def label2color(labels):
'''
labels: np.ndarray with shape (n, )