merged

predict/predict.py

@@ -2,7 +2,10 @@ import sys
 import os
 import shutil
 import math
-from dataset.shapenet import PredictNetPartSegDataset
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../')  # add project root directory
+
+from dataset.shapenet import ShapeNetPartSegDataset
 from model.pointnet2_part_seg import PointNet2PartSegmentNet
 import torch_geometric.transforms as GT
 import torch
@@ -237,7 +240,7 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../')  # add proj
 parser = argparse.ArgumentParser()
 parser.add_argument('--dataset', type=str, default='data', help='dataset path')
 parser.add_argument('--npoints', type=int, default=2048, help='resample points number')
-parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_241.pth', help='model path')
+parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_246.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)
@@ -279,12 +282,12 @@ if __name__ == '__main__':
     print('load dataset ..')
     test_transform = GT.Compose([GT.NormalizeScale(), ])
 
-    test_dataset = PredictNetPartSegDataset(
+    test_dataset = ShapeNetPartSegDataset(
+        mode='predict',
         root_dir=opt.dataset,
-        num_classes=4,
         transform=None,
         npoints=opt.npoints,
-        refresh=True
+        refresh=False
     )
     num_classes = test_dataset.num_classes()