Skip to content

Commit 3700c44

Browse files
committed
a bunch of fixes
1 parent f5cabcc commit 3700c44

File tree

3 files changed

+12
-10
lines changed

3 files changed

+12
-10
lines changed

datasets.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ def __init__(self, root, npoints = 2500, classification = False, class_choice =
5656
self.datapath.append((item, fn[0], fn[1]))
5757

5858

59-
self.classes = dict(zip(self.cat, range(len(self.cat))))
59+
self.classes = dict(zip(sorted(self.cat), range(len(self.cat))))
6060
print(self.classes)
6161
self.num_seg_classes = 0
6262
if not self.classification:

show_cls.py

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,29 +24,31 @@
2424
parser = argparse.ArgumentParser()
2525

2626
parser.add_argument('--model', type=str, default = '', help='model path')
27+
parser.add_argument('--num_points', type=int, default=2500, help='input batch size')
2728

2829

2930
opt = parser.parse_args()
3031
print (opt)
3132

32-
test_dataset = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0' , train = False, classification = True)
33+
test_dataset = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0' , train = False, classification = True, npoints = opt.num_points)
3334

34-
testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle = False)
35+
testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle = True)
3536

3637

37-
classifier = PointNetCls(k = len(test_dataset.classes))
38+
classifier = PointNetCls(k = len(test_dataset.classes), num_points = opt.num_points)
3839
classifier.cuda()
3940
classifier.load_state_dict(torch.load(opt.model))
4041
classifier.eval()
4142

43+
4244
for i, data in enumerate(testdataloader, 0):
4345
points, target = data
44-
points, target = Variable(points), Variable(target[:,0])
45-
points = points.transpose(2,1)
46+
points, target = Variable(points), Variable(target[:, 0])
47+
points = points.transpose(2, 1)
4648
points, target = points.cuda(), target.cuda()
4749
pred, _ = classifier(points)
4850
loss = F.nll_loss(pred, target)
49-
from IPython import embed; embed()
51+
5052
pred_choice = pred.data.max(1)[1]
5153
correct = pred_choice.eq(target.data).cpu().sum()
5254
print('i:%d loss: %f accuracy: %f' %(i, loss.data[0], correct/float(32)))

show_seg.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -57,12 +57,12 @@
5757

5858
point = Variable(point.view(1, point.size()[0], point.size()[1]))
5959
pred, _ = classifier(point)
60-
6160
pred_choice = pred.data.max(2)[1]
61+
print(pred_choice)
62+
6263
#print(pred_choice.size())
63-
pred_color = cmap[pred_choice.numpy(), :]
64+
pred_color = cmap[pred_choice.numpy()[0], :]
6465

6566
#print(pred_color.shape)
66-
6767
showpoints(point_np, gt, pred_color)
6868

0 commit comments

Comments
 (0)