ScanNet example

README.md

@@ -111,8 +111,9 @@ Examples in the examples folder include
 * [Assamese handwriting recognition](https://archive.ics.uci.edu/ml/datasets/Online+Handwritten+Assamese+Characters+Dataset#)
 * [Chinese handwriting for recognition](http://www.nlpr.ia.ac.cn/databases/handwriting/Online_database.html)
 * [3D Segmentation](https://shapenet.cs.stanford.edu/iccv17/) using ShapeNet Core-55
+* [ScanNet](http://www.scan-net.org/) 3D Semantic label benchmark
 
-Data will be downloaded/preprocessed on the first run, i.e.
+For example:
 ```
 cd examples/Assamese_handwriting
 python VGGplus.py

examples/ScanNet/README.md

+[ScanNet](http://www.scan-net.org/)
+-------
+
+To train a small U-Net with 5cm-cubed sparse voxels:
+
+1. Download [ScanNet](http://www.scan-net.org/) files
+2. [Split](https://github.com/ScanNet/ScanNet/tree/master/Tasks/Benchmark) the files *vh_clean_2.ply and *_vh_clean_2.labels.ply files into 'train/' and 'val/' folders
+3. Run 'pip install plyfile'
+4. Run 'python prepare_data.py'
+5. Run 'python unet.py'
+
+You can the computational cost (and hopefully accuracy too) by changing m / block_reps / residual_blocks / scale / val_reps in unet.py / data.py.

examples/ScanNet/iou.py

+# Copyright 2016-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch, numpy as np
+
+#Classes relabelled {-100,0,1,...,19}.
+#Predictions will all be in the set {0,1,...,19}
+
+#VALID_CLASS_IDS = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39])
+VALID_CLASS_IDS = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19])
+
+
+CLASS_LABELS = ['wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window', 'bookshelf', 'picture', 'counter', 'desk', 'curtain', 'refrigerator', 'shower curtain', 'toilet', 'sink', 'bathtub', 'otherfurniture']
+UNKNOWN_ID = -100
+
+def confusion_matrix(pred_ids, gt_ids):
+    assert pred_ids.shape == gt_ids.shape, (pred_ids.shape, gt_ids.shape)
+    idxs= gt_ids>=0
+    return np.bincount(pred_ids[idxs]*20+gt_ids[idxs],minlength=400).reshape((20,20)).astype(np.ulonglong)
+
+def get_iou(label_id, confusion):
+    # true positives
+    tp = np.longlong(confusion[label_id, label_id])
+    # false negatives
+    fn = np.longlong(confusion[label_id, :].sum()) - tp
+    # false positives
+    not_ignored = [l for l in VALID_CLASS_IDS if not l == label_id]
+    fp = np.longlong(confusion[not_ignored, label_id].sum())
+
+    denom = (tp + fp + fn)
+    if denom == 0:
+        return float('nan')
+    return (float(tp) / denom, tp, denom)
+
+def evaluate(pred_ids,gt_ids):
+    print('evaluating', gt_ids.size, 'points...')
+    confusion=confusion_matrix(pred_ids,gt_ids)
+    class_ious = {}
+    mean_iou = 0
+    for i in range(len(VALID_CLASS_IDS)):
+        label_name = CLASS_LABELS[i]
+        label_id = VALID_CLASS_IDS[i]
+        class_ious[label_name] = get_iou(label_id, confusion)
+        mean_iou+=class_ious[label_name][0]/20
+
+    print('classes          IoU')
+    print('----------------------------')
+    for i in range(len(VALID_CLASS_IDS)):
+        label_name = CLASS_LABELS[i]
+        print('{0:<14s}: {1:>5.3f}   ({2:>6d}/{3:<6d})'.format(label_name, class_ious[label_name][0], class_ious[label_name][1], class_ious[label_name][2]))
+    print('mean IOU', mean_iou)

examples/ScanNet/prepare_data.py

+# Copyright 2016-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import glob, plyfile, numpy as np, multiprocessing as mp, torch
+
+# Map relevant classes to {0,1,...,19}, and ignored classes to -100
+remapper=np.ones(150)*(-100)
+for i,x in enumerate([1,2,3,4,5,6,7,8,9,10,11,12,14,16,24,28,33,34,36,39]):
+    remapper[x]=i
+
+files=sorted(glob.glob('*/*_vh_clean_2.ply'))
+files2=sorted(glob.glob('*/*_vh_clean_2.labels.ply'))
+assert len(files) == len(files2)
+
+def f(fn):
+    fn2 = fn[:-3]+'labels.ply'
+    a=plyfile.PlyData().read(fn)
+    v=np.array([list(x) for x in a.elements[0]])
+    coords=np.ascontiguousarray(v[:,:3]-v[:,:3].mean(0))
+    colors=np.ascontiguousarray(v[:,3:6])/127.5-1
+    a=plyfile.PlyData().read(fn2)
+    w=remapper[np.array(a.elements[0]['label'])]
+    torch.save((coords,colors,w),fn[:-4]+'.pth')
+    print(fn, fn2)
+
+p = mp.Pool(processes=mp.cpu_count())
+p.map(f,files)
+p.close()
+p.join()

examples/ScanNet/train/.gitignore

+*
+!.gitignore

examples/ScanNet/unet.py

+# Copyright 2016-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Options
+m = 16 # 16 or 32
+residual_blocks=False #True or False
+block_reps = 1 #Conv block repetition factor: 1 or 2
+
+import torch, data, iou
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+import sparseconvnet as scn
+import time
+import os, sys, glob
+import math
+import numpy as np
+
+use_cuda = torch.cuda.is_available()
+exp_name='unet_scale20_m16_rep1_notResidualBlocks'
+
+class Model(nn.Module):
+    def __init__(self):
+        nn.Module.__init__(self)
+        self.sparseModel = scn.Sequential().add(
+           scn.InputLayer(data.dimension,data.full_scale, mode=4)).add(
+           scn.SubmanifoldConvolution(data.dimension, 3, m, 3, False)).add(
+               scn.UNet(data.dimension, block_reps, [m, 2*m, 3*m, 4*m, 5*m, 6*m, 7*m], residual_blocks)).add(
+           scn.BatchNormReLU(m)).add(
+           scn.OutputLayer(data.dimension))
+        self.linear = nn.Linear(m, 20)
+    def forward(self,x):
+        x=self.sparseModel(x)
+        x=self.linear(x)
+        return x
+
+unet=Model()
+if use_cuda:
+    unet=unet.cuda()
+
+training_epochs=512
+training_epoch=scn.checkpoint_restore(unet,exp_name,'unet',use_cuda)
+optimizer = optim.Adam(unet.parameters())
+print('#classifer parameters', sum([x.nelement() for x in unet.parameters()]))
+
+for epoch in range(training_epoch, training_epochs+1):
+    unet.train()
+    stats = {}
+    scn.forward_pass_multiplyAdd_count=0
+    scn.forward_pass_hidden_states=0
+    start = time.time()
+    train_loss=0
+    for i,batch in enumerate(data.train_data_loader):
+        optimizer.zero_grad()
+        if use_cuda:
+            batch['x'][1]=batch['x'][1].cuda()
+            batch['y']=batch['y'].cuda()
+        predictions=unet(batch['x'])
+        loss = torch.nn.functional.cross_entropy(predictions,batch['y'])
+        train_loss+=loss.item()
+        loss.backward()
+        optimizer.step()
+    print(epoch,'Train loss',train_loss/(i+1), 'MegaMulAdd=',scn.forward_pass_multiplyAdd_count/len(data.train)/1e6, 'MegaHidden',scn.forward_pass_hidden_states/len(data.train)/1e6,'time=',time.time() - start,'s')
+    scn.checkpoint_save(unet,exp_name,'unet',epoch, use_cuda)
+
+    if scn.is_power2(epoch):
+        with torch.no_grad():
+            unet.eval()
+            store=torch.zeros(data.valOffsets[-1],20)
+            scn.forward_pass_multiplyAdd_count=0
+            scn.forward_pass_hidden_states=0
+            start = time.time()
+            for rep in range(1,1+data.val_reps):
+                for i,batch in enumerate(data.val_data_loader):
+                    if use_cuda:
+                        batch['x'][1]=batch['x'][1].cuda()
+                        batch['y']=batch['y'].cuda()
+                    predictions=unet(batch['x'])
+                    store.index_add_(0,batch['point_ids'],predictions.cpu())
+                print(epoch,rep,'Val MegaMulAdd=',scn.forward_pass_multiplyAdd_count/len(data.val)/1e6, 'MegaHidden',scn.forward_pass_hidden_states/len(data.val)/1e6,'time=',time.time() - start,'s')
+                iou.evaluate(store.max(1)[1].numpy(),data.valLabels)

examples/ScanNet/val/.gitignore

+*
+!.gitignore