File based header detection, collate_per_PC training.

vis/show_seg_res.py

@@ -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 PredictNetPartSegDataset
+from dataset.shapenet import ShapeNetPartSegDataset
 from model.pointnet2_part_seg import PointNet2PartSegmentNet
 import torch_geometric.transforms as GT
 import torch
@@ -16,8 +16,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_custom_8.pth', help='model path')
+parser.add_argument('--npoints', type=int, default=2048, help='resample points number')
+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')
 opt = parser.parse_args()
 print(opt)
@@ -28,8 +28,10 @@ if __name__ == '__main__':
     print('Construct dataset ..')
     test_transform = GT.Compose([GT.NormalizeScale(),])
 
-    test_dataset = PredictNetPartSegDataset(
+    test_dataset = ShapeNetPartSegDataset(
         root_dir=opt.dataset,
+        collate_per_segment=False,
+        train=False,
         transform=test_transform,
         npoints=opt.npoints
     )
@@ -121,16 +123,17 @@ if __name__ == '__main__':
 
         print('mIoU: ', compute_mIoU(pred_labels, gt_labels))
 
-
         # View result
+        if True:
+            print('View gt labels ..')
+            view_points_labels(points, gt_labels)
 
-        # print('View gt labels ..')
-        # view_points_labels(points, gt_labels)
+        if True:
+            print('View diff labels ..')
+            print(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 ..')
-        # print(pred_labels)
-        # view_points_labels(points, pred_labels)
+        if True:
+            print('View pred labels ..')
+            print(pred_labels)
+            view_points_labels(points, pred_labels)

vis/view.py

@@ -48,7 +48,7 @@ def label2color(labels):
 
     minl, maxl = np.min(labels), np.max(labels)
     for l in range(minl, maxl + 1):
-        colors[labels==l, :] = mini_color_table(l) 
+        colors[labels == l, :] = mini_color_table(l)
 
     return colors