fix and close #7 #12

Author: fxia22

.gitignore

@@ -3,3 +3,7 @@ data
 *.pyc
 *.ipynb
 shapenetcore_partanno_segmentation_benchmark_v0/
+*.so
+.idea*
+cls/
+seg/

datasets.py

@@ -60,7 +60,7 @@ def __init__(self, root, npoints = 2500, classification = False, class_choice =
         print(self.classes)
         self.num_seg_classes = 0
         if not self.classification:
-            for i in range(len(self.datapath)/50):
+            for i in range(len(self.datapath)//50):
                 l = len(np.unique(np.loadtxt(self.datapath[i][-1]).astype(np.uint8)))
                 if l > self.num_seg_classes:
                     self.num_seg_classes = l

download.sh

@@ -1,2 +1,3 @@
 wget https://shapenet.cs.stanford.edu/ericyi/shapenetcore_partanno_segmentation_benchmark_v0.zip
 unzip shapenetcore_partanno_segmentation_benchmark_v0.zip
+rm shapenetcore_partanno_segmentation_benchmark_v0.zip

pointnet.py

@@ -105,7 +105,7 @@ def forward(self, x):
         x = F.relu(self.bn1(self.fc1(x)))
         x = F.relu(self.bn2(self.fc2(x)))
         x = self.fc3(x)
-        return F.log_softmax(x), trans
+        return F.log_softmax(x, dim=-1), trans
 
 class PointNetDenseCls(nn.Module):
     def __init__(self, num_points = 2500, k = 2):
@@ -129,7 +129,7 @@ def forward(self, x):
         x = F.relu(self.bn3(self.conv3(x)))
         x = self.conv4(x)
         x = x.transpose(2,1).contiguous()
-        x = F.log_softmax(x.view(-1,self.k))
+        x = F.log_softmax(x.view(-1,self.k), dim=-1)
         x = x.view(batchsize, self.num_points, self.k)
         return x, trans
 

show_cls.py

@@ -46,6 +46,7 @@
     points, target = points.cuda(), target.cuda()
     pred, _ = classifier(points)
     loss = F.nll_loss(pred, target)
+    from IPython import embed; embed()
     pred_choice = pred.data.max(1)[1]
     correct = pred_choice.eq(target.data).cpu().sum()
     print('i:%d  loss: %f accuracy: %f' %(i, loss.data[0], correct/float(32)))

show_seg.py

@@ -51,13 +51,14 @@
 
 classifier = PointNetDenseCls(k = 4)
 classifier.load_state_dict(torch.load(opt.model))
+classifier.eval()
 
 point = point.transpose(1,0).contiguous()
 
 point = Variable(point.view(1, point.size()[0], point.size()[1]))
 pred, _ = classifier(point)
 
-pred_choice = pred.data.max(2)[1][0,:,0]
+pred_choice = pred.data.max(2)[1]
 #print(pred_choice.size())
 pred_color = cmap[pred_choice.numpy(), :]
 

train_classification.py

@@ -83,7 +83,7 @@
         print('[%d: %d/%d] train loss: %f accuracy: %f' %(epoch, i, num_batch, loss.data[0], correct/float(opt.batchSize)))
 
         if i % 10 == 0:
-            j, data = enumerate(testdataloader, 0).next()
+            j, data = next(enumerate(testdataloader, 0))
             points, target = data
             points, target = Variable(points), Variable(target[:,0])
             points = points.transpose(2,1)

train_segmentation.py

@@ -83,7 +83,7 @@
         print('[%d: %d/%d] train loss: %f accuracy: %f' %(epoch, i, num_batch, loss.data[0], correct/float(opt.batchSize * 2500)))
 
         if i % 10 == 0:
-            j, data = enumerate(testdataloader, 0).next()
+            j, data = next(enumerate(testdataloader, 0))
             points, target = data
             points, target = Variable(points), Variable(target)
             points = points.transpose(2,1)