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 18:16:17 +02:00
parent a9bf053794
commit b7d127e840
2 changed files with 7 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
dataset/shapenet.py
View File

@ -287,10 +287,11 @@ 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):
def __init__(self, root_dir, train=False, transform=None, npoints=2048, headers=True, collate_per_segment=False):
super(PredictNetPartSegDataset, self).__init__()
self.npoints = npoints
self.dataset = ShapeNetPartSegDataset(root=root_dir, train=False, transform=transform, headers=headers)
self.dataset = PredictionShapeNet(root=root_dir, train=train, transform=transform,
headers=headers, collate_per_segment=collate_per_segment)
def __getitem__(self, index):
data = self.dataset[index]

8
vis/show_seg_res.py
View File

@ -5,7 +5,7 @@ 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, ShapeNetPartSegDataset
from model.pointnet2_part_seg import PointNet2PartSegmentNet
import torch_geometric.transforms as GT
import torch
@ -28,7 +28,7 @@ if __name__ == '__main__':
print('Construct dataset ..')
test_transform = GT.Compose([GT.NormalizeScale(),])
test_dataset = ShapeNetPartSegDataset(
test_dataset = PredictNetPartSegDataset(
root_dir=opt.dataset,
collate_per_segment=False,
train=False,
@ -128,12 +128,12 @@ if __name__ == '__main__':
print('View gt labels ..')
view_points_labels(points, gt_labels)
if True:
if False:
print('View diff labels ..')
print(diff_labels)
view_points_labels(points, diff_labels)
if True:
if False:
print('View pred labels ..')
print(pred_labels)
view_points_labels(points, pred_labels)