redo

predict/predict.py

@@ -2,7 +2,7 @@ import sys
 import os
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../')  # add project root directory
 
-from dataset.shapenet import PredictNetPartSegDataset, ShapeNetPartSegDataset
+from dataset.shapenet import ShapeNetPartSegDataset
 from model.pointnet2_part_seg import PointNet2PartSegmentNet
 import torch_geometric.transforms as GT
 import torch
@@ -13,7 +13,7 @@ import argparse
 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_249.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)
@@ -24,9 +24,9 @@ if __name__ == '__main__':
     print('Construct 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