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-08-01 14:16:50 +02:00
parent 47a76dc978
commit a9bf053794
3 changed files with 77 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset/shapenet.py
+59 -60
View File
@@ -20,6 +20,7 @@ class CustomShapeNet(InMemoryDataset):
headers=True, **kwargs): headers=True, **kwargs):
self.has_headers = headers self.has_headers = headers
self.collate_per_element = collate_per_segment self.collate_per_element = collate_per_segment
self.train = train
super(CustomShapeNet, self).__init__(root, transform, pre_transform, pre_filter) 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) self.data, self.slices = torch.load(path)
@@ -70,73 +71,71 @@ class CustomShapeNet(InMemoryDataset):
return data return data
def process(self, delimiter=' '): def process(self, delimiter=' '):
# idx = self.categories[self.category]
# paths = [osp.join(path, idx) for path in self.raw_paths]
datasets = defaultdict(list) datasets = defaultdict(list)
for idx, setting in enumerate(self.raw_file_names): idx, data_folder = (0, self.raw_file_names[0]) if self.train else (1, self.raw_file_names[1])
path_to_clouds = os.path.join(self.raw_dir, setting) path_to_clouds = os.path.join(self.raw_dir, data_folder)
if '.headers' in os.listdir(path_to_clouds): if '.headers' in os.listdir(path_to_clouds):
self.has_headers = True self.has_headers = True
elif 'no.headers' in os.listdir(path_to_clouds): elif 'no.headers' in os.listdir(path_to_clouds):
self.has_headers = False self.has_headers = False
else: else:
pass pass
for pointcloud in tqdm(os.scandir(path_to_clouds)): for pointcloud in tqdm(os.scandir(path_to_clouds)):
if not os.path.isdir(pointcloud):
continue
data, paths = None, list()
for ext in ['dat', 'xyz']:
paths.extend(glob.glob(os.path.join(pointcloud.path, f'*.{ext}')))
if not os.path.isdir(pointcloud): for element in paths:
continue if all([x not in os.path.split(element)[-1] for x in ['pc.dat', 'pc.xyz']]):
data, paths = None, list() # Assign training data to the data container
for ext in ['dat', 'xyz']: # Following the original logic;
paths.extend(glob.glob(os.path.join(pointcloud.path, f'*.{ext}'))) # y should be the label;
for element in paths: # pos should be the six dimensional vector describing: !its pos not points!!
if all([x not in os.path.split(element)[-1] for x in ['pc.dat', 'pc.xyz']]): # x,y,z,x_rot,y_rot,z_rot
# 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
# Get the y - Label # Get the y - Label
y_raw = next(i for i, v in enumerate(self.categories.keys()) if v.lower() in element.lower()) 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] # y_raw = os.path.splitext(element)[0].split('_')[-2]
with open(element,'r') as f: with open(element,'r') as f:
if self.has_headers: if self.has_headers:
headers = f.__next__() headers = f.__next__()
# Check if there are no useable nodes in this file, header says 0. # Check if there are no useable nodes in this file, header says 0.
if not int(headers.rstrip().split(delimiter)[0]): if not int(headers.rstrip().split(delimiter)[0]):
continue continue
# Iterate over all rows # Iterate over all rows
src = [[float(x) if x not in ['-nan(ind)', 'nan(ind)'] else 0 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 != ''] for x in line.rstrip().split(delimiter)[None:None]] for line in f if line != '']
points = torch.tensor(src, dtype=None).squeeze() points = torch.tensor(src, dtype=None).squeeze()
if not len(points.shape) > 1: if not len(points.shape) > 1:
continue continue
# pos = points[:, :3] # pos = points[:, :3]
# norm = points[:, 3:] # norm = points[:, 3:]
y_all = [y_raw] * points.shape[0] y_all = [y_raw] * points.shape[0]
y = torch.as_tensor(y_all, dtype=torch.int) y = torch.as_tensor(y_all, dtype=torch.int)
# points = torch.as_tensor(points, dtype=torch.float) # points = torch.as_tensor(points, dtype=torch.float)
# norm = torch.as_tensor(norm, dtype=torch.float) # norm = torch.as_tensor(norm, dtype=torch.float)
if self.collate_per_element: if self.collate_per_element:
data = Data(y=y, pos=points[:, :3]) data = Data(y=y, pos=points[:, :3])
else: else:
if not data: if not data:
data = defaultdict(list) data = defaultdict(list)
for key, val in dict(y=y, pos= points[:, :3]).items(): for key, val in dict(y=y, pos= points[:, :3]).items():
data[key].append(val) data[key].append(val)
# , points=points, norm=points[:3], ) # , points=points, norm=points[:3], )
data = self._transform_and_filter(data) data = self._transform_and_filter(data)
if self.collate_per_element: if self.collate_per_element:
datasets[setting].append(data) datasets[data_folder].append(data)
if not self.collate_per_element: if not self.collate_per_element:
datasets[setting].append(Data(**{key: torch.cat(data[key]) for key in data.keys()})) datasets[data_folder].append(Data(**{key: torch.cat(data[key]) for key in data.keys()}))
if datasets[data_folder]:
os.makedirs(self.processed_dir, exist_ok=True) os.makedirs(self.processed_dir, exist_ok=True)
torch.save(self.collate(datasets[setting]), self.processed_paths[idx]) torch.save(self.collate(datasets[data_folder]), self.processed_paths[idx])
def __repr__(self): def __repr__(self):
return f'{self.__class__.__name__}({len(self)})' return f'{self.__class__.__name__}({len(self)})'
@@ -291,7 +290,7 @@ class PredictNetPartSegDataset(Dataset):
def __init__(self, root_dir, transform=None, npoints=2048, headers=True): def __init__(self, root_dir, transform=None, npoints=2048, headers=True):
super(PredictNetPartSegDataset, self).__init__() super(PredictNetPartSegDataset, self).__init__()
self.npoints = npoints self.npoints = npoints
self.dataset = PredictionShapeNet(root=root_dir, train=False, transform=transform, headers=headers) self.dataset = ShapeNetPartSegDataset(root=root_dir, train=False, transform=transform, headers=headers)
def __getitem__(self, index): def __getitem__(self, index):
data = self.dataset[index] data = self.dataset[index]
vis/show_seg_res.py
+17 -14
View File
@@ -5,7 +5,7 @@ import sys
import os import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../') # add project root directory sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../') # add project root directory
from dataset.shapenet import PredictNetPartSegDataset from dataset.shapenet import ShapeNetPartSegDataset
from model.pointnet2_part_seg import PointNet2PartSegmentNet from model.pointnet2_part_seg import PointNet2PartSegmentNet
import torch_geometric.transforms as GT import torch_geometric.transforms as GT
import torch import torch
@@ -16,8 +16,8 @@ import argparse
## ##
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default='data', help='dataset path') 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('--npoints', type=int, default=2048, help='resample points number')
parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_8.pth', help='model path') parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_249.pth', help='model path')
parser.add_argument('--sample_idx', type=int, default=0, help='select a sample to segment and view result') parser.add_argument('--sample_idx', type=int, default=0, help='select a sample to segment and view result')
opt = parser.parse_args() opt = parser.parse_args()
print(opt) print(opt)
@@ -28,8 +28,10 @@ if __name__ == '__main__':
print('Construct dataset ..') print('Construct dataset ..')
test_transform = GT.Compose([GT.NormalizeScale(),]) test_transform = GT.Compose([GT.NormalizeScale(),])
test_dataset = PredictNetPartSegDataset( test_dataset = ShapeNetPartSegDataset(
root_dir=opt.dataset, root_dir=opt.dataset,
collate_per_segment=False,
train=False,
transform=test_transform, transform=test_transform,
npoints=opt.npoints npoints=opt.npoints
) )
@@ -121,16 +123,17 @@ if __name__ == '__main__':
print('mIoU: ', compute_mIoU(pred_labels, gt_labels)) print('mIoU: ', compute_mIoU(pred_labels, gt_labels))
# View result # View result
if True:
print('View gt labels ..')
view_points_labels(points, gt_labels)
# print('View gt labels ..') if True:
# view_points_labels(points, gt_labels) print('View diff labels ..')
print(diff_labels)
view_points_labels(points, diff_labels)
print('View diff labels ..') if True:
print(diff_labels) print('View pred labels ..')
view_points_labels(points, diff_labels) print(pred_labels)
view_points_labels(points, pred_labels)
# print('View pred labels ..')
# print(pred_labels)
# view_points_labels(points, pred_labels)
vis/view.py
+1 -1
View File
@@ -48,7 +48,7 @@ def label2color(labels):
minl, maxl = np.min(labels), np.max(labels) minl, maxl = np.min(labels), np.max(labels)
for l in range(minl, maxl + 1): for l in range(minl, maxl + 1):
colors[labels==l, :] = mini_color_table(l) colors[labels == l, :] = mini_color_table(l)
return colors return colors