update

dataset.py

@@ -0,0 +1,88 @@
+
+
+class ShapeNetDataset(data.Dataset):
+    def __init__(self,
+                 root,
+                 npoints=2500,
+                 classification=False,
+                 class_choice=None,
+                 split='train',
+                 data_augmentation=True):
+        self.npoints = npoints
+        self.root = root
+        self.catfile = os.path.join(self.root, 'synsetoffset2category.txt')
+        self.cat = {}
+        self.data_augmentation = data_augmentation
+        self.classification = classification
+        self.seg_classes = {}
+        
+        with open(self.catfile, 'r') as f:
+            for line in f:
+                ls = line.strip().split()
+                self.cat[ls[0]] = ls[1]
+        #print(self.cat)
+        if not class_choice is None:
+            self.cat = {k: v for k, v in self.cat.items() if k in class_choice}
+
+        self.id2cat = {v: k for k, v in self.cat.items()}
+
+        self.meta = {}
+        splitfile = os.path.join(self.root, 'train_test_split', 'shuffled_{}_file_list.json'.format(split))
+        #from IPython import embed; embed()
+        filelist = json.load(open(splitfile, 'r'))
+        for item in self.cat:
+            self.meta[item] = []
+
+        for file in filelist:
+            _, category, uuid = file.split('/')
+            if category in self.cat.values():
+                self.meta[self.id2cat[category]].append((os.path.join(self.root, category, 'points', uuid+'.pts'),
+                                        os.path.join(self.root, category, 'points_label', uuid+'.seg')))
+
+        self.datapath = []
+        for item in self.cat:
+            for fn in self.meta[item]:
+                self.datapath.append((item, fn[0], fn[1]))
+
+        self.classes = dict(zip(sorted(self.cat), range(len(self.cat))))
+        print(self.classes)
+        with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../misc/num_seg_classes.txt'), 'r') as f:
+            for line in f:
+                ls = line.strip().split()
+                self.seg_classes[ls[0]] = int(ls[1])
+        self.num_seg_classes = self.seg_classes[list(self.cat.keys())[0]]
+        print(self.seg_classes, self.num_seg_classes)
+
+    def __getitem__(self, index):
+        fn = self.datapath[index]
+        cls = self.classes[self.datapath[index][0]]
+        point_set = np.loadtxt(fn[1]).astype(np.float32)
+        seg = np.loadtxt(fn[2]).astype(np.int64)
+        #print(point_set.shape, seg.shape)
+
+        choice = np.random.choice(len(seg), self.npoints, replace=True)
+        #resample
+        point_set = point_set[choice, :]
+
+        point_set = point_set - np.expand_dims(np.mean(point_set, axis = 0), 0) # center
+        dist = np.max(np.sqrt(np.sum(point_set ** 2, axis = 1)),0)
+        point_set = point_set / dist #scale
+
+        if self.data_augmentation:
+            theta = np.random.uniform(0,np.pi*2)
+            rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]])
+            point_set[:,[0,2]] = point_set[:,[0,2]].dot(rotation_matrix) # random rotation
+            point_set += np.random.normal(0, 0.02, size=point_set.shape) # random jitter
+
+        seg = seg[choice]
+        point_set = torch.from_numpy(point_set)
+        seg = torch.from_numpy(seg)
+        cls = torch.from_numpy(np.array([cls]).astype(np.int64))
+
+        if self.classification:
+            return point_set, cls
+        else:
+            return point_set, seg
+
+    def __len__(self):
+        return len(self.datapath)

geom_solver.py

@@ -17,7 +17,7 @@ class GeomSolver(object):
 
 	def __init__(self):
 		self.min_vertices = 18
-		self.kmeans_th = 150
+		self.kmeans_th = 70
 		self.clr_th = 2.5
 		self.device = 'cuda:0'
 
@@ -44,9 +44,9 @@ class GeomSolver(object):
 			vert_mask = (vert_mask > 0).astype(np.uint8)
 
 			dist = cv2.distanceTransform(1-vert_mask, cv2.DIST_L2, 3)
-			dist[dist > 100] = 100
-			ndist = np.zeros_like(dist)
-			ndist = cv2.normalize(dist, ndist, 0, 1.0, cv2.NORM_MINMAX)
+			# dist[dist > 100] = 100
+			# ndist = np.zeros_like(dist)
+			# ndist = cv2.normalize(dist, ndist, 0, 1.0, cv2.NORM_MINMAX)
 
 			in_this_image = np.array([cki in p.image_ids for p in self.points3D.values()])
 			uv = torch.round(self.pyt_cameras[ki].transform_points(self.verts)[:, :2]).cpu().numpy().astype(int)
@@ -57,7 +57,7 @@ class GeomSolver(object):
 			dist_points[uv_inl] += dist[uv[:,1], uv[:,0]]
 			visible_counts[uv_inl] += 1
 
-		selected_points = (dist_points / (visible_counts + 1e-6)) <= 10
+		selected_points = (dist_points / (visible_counts + 1e-6)) <= 15
 		selected_points[visible_counts < 1] = False
 
 		pnts = torch.from_numpy(self.xyz[selected_points].astype(np.float32))[None]

pointnet.py

@@ -0,0 +1,213 @@
+from __future__ import print_function
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.utils.data
+from torch.autograd import Variable
+import numpy as np
+import torch.nn.functional as F
+
+
+class STN3d(nn.Module):
+    def __init__(self):
+        super(STN3d, self).__init__()
+        self.conv1 = torch.nn.Conv1d(3, 64, 1)
+        self.conv2 = torch.nn.Conv1d(64, 128, 1)
+        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
+        self.fc1 = nn.Linear(1024, 512)
+        self.fc2 = nn.Linear(512, 256)
+        self.fc3 = nn.Linear(256, 9)
+        self.relu = nn.ReLU()
+
+        self.bn1 = nn.BatchNorm1d(64)
+        self.bn2 = nn.BatchNorm1d(128)
+        self.bn3 = nn.BatchNorm1d(1024)
+        self.bn4 = nn.BatchNorm1d(512)
+        self.bn5 = nn.BatchNorm1d(256)
+
+
+    def forward(self, x):
+        batchsize = x.size()[0]
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = F.relu(self.bn3(self.conv3(x)))
+        x = torch.max(x, 2, keepdim=True)[0]
+        x = x.view(-1, 1024)
+
+        x = F.relu(self.bn4(self.fc1(x)))
+        x = F.relu(self.bn5(self.fc2(x)))
+        x = self.fc3(x)
+
+        iden = Variable(torch.from_numpy(np.array([1,0,0,0,1,0,0,0,1]).astype(np.float32))).view(1,9).repeat(batchsize,1)
+        if x.is_cuda:
+            iden = iden.cuda()
+        x = x + iden
+        x = x.view(-1, 3, 3)
+        return x
+
+
+class STNkd(nn.Module):
+    def __init__(self, k=64):
+        super(STNkd, self).__init__()
+        self.conv1 = torch.nn.Conv1d(k, 64, 1)
+        self.conv2 = torch.nn.Conv1d(64, 128, 1)
+        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
+        self.fc1 = nn.Linear(1024, 512)
+        self.fc2 = nn.Linear(512, 256)
+        self.fc3 = nn.Linear(256, k*k)
+        self.relu = nn.ReLU()
+
+        self.bn1 = nn.BatchNorm1d(64)
+        self.bn2 = nn.BatchNorm1d(128)
+        self.bn3 = nn.BatchNorm1d(1024)
+        self.bn4 = nn.BatchNorm1d(512)
+        self.bn5 = nn.BatchNorm1d(256)
+
+        self.k = k
+
+    def forward(self, x):
+        batchsize = x.size()[0]
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = F.relu(self.bn3(self.conv3(x)))
+        x = torch.max(x, 2, keepdim=True)[0]
+        x = x.view(-1, 1024)
+
+        x = F.relu(self.bn4(self.fc1(x)))
+        x = F.relu(self.bn5(self.fc2(x)))
+        x = self.fc3(x)
+
+        iden = Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32))).view(1,self.k*self.k).repeat(batchsize,1)
+        if x.is_cuda:
+            iden = iden.cuda()
+        x = x + iden
+        x = x.view(-1, self.k, self.k)
+        return x
+
+class PointNetfeat(nn.Module):
+    def __init__(self, global_feat = True, feature_transform = False):
+        super(PointNetfeat, self).__init__()
+        self.stn = STN3d()
+        self.conv1 = torch.nn.Conv1d(3, 64, 1)
+        self.conv2 = torch.nn.Conv1d(64, 128, 1)
+        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
+        self.bn1 = nn.BatchNorm1d(64)
+        self.bn2 = nn.BatchNorm1d(128)
+        self.bn3 = nn.BatchNorm1d(1024)
+        self.global_feat = global_feat
+        self.feature_transform = feature_transform
+        if self.feature_transform:
+            self.fstn = STNkd(k=64)
+
+    def forward(self, x):
+        n_pts = x.size()[2]
+        trans = self.stn(x)
+        x = x.transpose(2, 1)
+        x = torch.bmm(x, trans)
+        x = x.transpose(2, 1)
+        x = F.relu(self.bn1(self.conv1(x)))
+
+        if self.feature_transform:
+            trans_feat = self.fstn(x)
+            x = x.transpose(2,1)
+            x = torch.bmm(x, trans_feat)
+            x = x.transpose(2,1)
+        else:
+            trans_feat = None
+
+        pointfeat = x
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = self.bn3(self.conv3(x))
+        x = torch.max(x, 2, keepdim=True)[0]
+        x = x.view(-1, 1024)
+        if self.global_feat:
+            return x, trans, trans_feat
+        else:
+            x = x.view(-1, 1024, 1).repeat(1, 1, n_pts)
+            return torch.cat([x, pointfeat], 1), trans, trans_feat
+
+class PointNetCls(nn.Module):
+    def __init__(self, k=2, feature_transform=False):
+        super(PointNetCls, self).__init__()
+        self.feature_transform = feature_transform
+        self.feat = PointNetfeat(global_feat=True, feature_transform=feature_transform)
+        self.fc1 = nn.Linear(1024, 512)
+        self.fc2 = nn.Linear(512, 256)
+        self.fc3 = nn.Linear(256, k)
+        self.dropout = nn.Dropout(p=0.3)
+        self.bn1 = nn.BatchNorm1d(512)
+        self.bn2 = nn.BatchNorm1d(256)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x, trans, trans_feat = self.feat(x)
+        x = F.relu(self.bn1(self.fc1(x)))
+        x = F.relu(self.bn2(self.dropout(self.fc2(x))))
+        x = self.fc3(x)
+        return F.log_softmax(x, dim=1), trans, trans_feat
+
+
+class PointNetDenseCls(nn.Module):
+    def __init__(self, k = 2, feature_transform=False):
+        super(PointNetDenseCls, self).__init__()
+        self.k = k
+        self.feature_transform=feature_transform
+        self.feat = PointNetfeat(global_feat=False, feature_transform=feature_transform)
+        self.conv1 = torch.nn.Conv1d(1088, 512, 1)
+        self.conv2 = torch.nn.Conv1d(512, 256, 1)
+        self.conv3 = torch.nn.Conv1d(256, 128, 1)
+        self.conv4 = torch.nn.Conv1d(128, self.k, 1)
+        self.bn1 = nn.BatchNorm1d(512)
+        self.bn2 = nn.BatchNorm1d(256)
+        self.bn3 = nn.BatchNorm1d(128)
+
+    def forward(self, x):
+        batchsize = x.size()[0]
+        n_pts = x.size()[2]
+        x, trans, trans_feat = self.feat(x)
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = F.relu(self.bn2(self.conv2(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), dim=-1)
+        x = x.view(batchsize, n_pts, self.k)
+        return x, trans, trans_feat
+
+def feature_transform_regularizer(trans):
+    d = trans.size()[1]
+    batchsize = trans.size()[0]
+    I = torch.eye(d)[None, :, :]
+    if trans.is_cuda:
+        I = I.cuda()
+    loss = torch.mean(torch.norm(torch.bmm(trans, trans.transpose(2,1)) - I, dim=(1,2)))
+    return loss
+
+if __name__ == '__main__':
+    sim_data = Variable(torch.rand(32,3,2500))
+    trans = STN3d()
+    out = trans(sim_data)
+    print('stn', out.size())
+    print('loss', feature_transform_regularizer(out))
+
+    sim_data_64d = Variable(torch.rand(32, 64, 2500))
+    trans = STNkd(k=64)
+    out = trans(sim_data_64d)
+    print('stn64d', out.size())
+    print('loss', feature_transform_regularizer(out))
+
+    pointfeat = PointNetfeat(global_feat=True)
+    out, _, _ = pointfeat(sim_data)
+    print('global feat', out.size())
+
+    pointfeat = PointNetfeat(global_feat=False)
+    out, _, _ = pointfeat(sim_data)
+    print('point feat', out.size())
+
+    cls = PointNetCls(k = 5)
+    out, _, _ = cls(sim_data)
+    print('class', out.size())
+
+    seg = PointNetDenseCls(k = 3)
+    out, _, _ = seg(sim_data)
+    print('seg', out.size())

train_pointnet.py

@@ -0,0 +1,148 @@
+from __future__ import print_function
+import argparse
+import os
+import random
+import torch
+import torch.nn.parallel
+import torch.optim as optim
+import torch.utils.data
+from pointnet.dataset import ShapeNetDataset, ModelNetDataset
+from pointnet import PointNetCls, feature_transform_regularizer
+import torch.nn.functional as F
+from tqdm import tqdm
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    '--batchSize', type=int, default=32, help='input batch size')
+parser.add_argument(
+    '--num_points', type=int, default=2500, help='input batch size')
+parser.add_argument(
+    '--workers', type=int, help='number of data loading workers', default=4)
+parser.add_argument(
+    '--nepoch', type=int, default=250, help='number of epochs to train for')
+parser.add_argument('--outf', type=str, default='cls', help='output folder')
+parser.add_argument('--model', type=str, default='', help='model path')
+parser.add_argument('--dataset', type=str, required=True, help="dataset path")
+parser.add_argument('--dataset_type', type=str, default='shapenet', help="dataset type shapenet|modelnet40")
+parser.add_argument('--feature_transform', action='store_true', help="use feature transform")
+
+opt = parser.parse_args()
+print(opt)
+
+blue = lambda x: '\033[94m' + x + '\033[0m'
+
+opt.manualSeed = random.randint(1, 10000)  # fix seed
+print("Random Seed: ", opt.manualSeed)
+random.seed(opt.manualSeed)
+torch.manual_seed(opt.manualSeed)
+
+if opt.dataset_type == 'shapenet':
+    dataset = ShapeNetDataset(
+        root=opt.dataset,
+        classification=True,
+        npoints=opt.num_points)
+
+    test_dataset = ShapeNetDataset(
+        root=opt.dataset,
+        classification=True,
+        split='test',
+        npoints=opt.num_points,
+        data_augmentation=False)
+elif opt.dataset_type == 'modelnet40':
+    dataset = ModelNetDataset(
+        root=opt.dataset,
+        npoints=opt.num_points,
+        split='trainval')
+
+    test_dataset = ModelNetDataset(
+        root=opt.dataset,
+        split='test',
+        npoints=opt.num_points,
+        data_augmentation=False)
+else:
+    exit('wrong dataset type')
+
+
+dataloader = torch.utils.data.DataLoader(
+    dataset,
+    batch_size=opt.batchSize,
+    shuffle=True,
+    num_workers=int(opt.workers))
+
+testdataloader = torch.utils.data.DataLoader(
+        test_dataset,
+        batch_size=opt.batchSize,
+        shuffle=True,
+        num_workers=int(opt.workers))
+
+print(len(dataset), len(test_dataset))
+num_classes = len(dataset.classes)
+print('classes', num_classes)
+
+try:
+    os.makedirs(opt.outf)
+except OSError:
+    pass
+
+classifier = PointNetCls(k=num_classes, feature_transform=opt.feature_transform)
+
+if opt.model != '':
+    classifier.load_state_dict(torch.load(opt.model))
+
+
+optimizer = optim.Adam(classifier.parameters(), lr=0.001, betas=(0.9, 0.999))
+scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
+classifier.cuda()
+
+num_batch = len(dataset) / opt.batchSize
+
+for epoch in range(opt.nepoch):
+    scheduler.step()
+    for i, data in enumerate(dataloader, 0):
+        points, target = data
+        target = target[:, 0]
+        points = points.transpose(2, 1)
+        points, target = points.cuda(), target.cuda()
+        optimizer.zero_grad()
+        classifier = classifier.train()
+        pred, trans, trans_feat = classifier(points)
+        loss = F.nll_loss(pred, target)
+        if opt.feature_transform:
+            loss += feature_transform_regularizer(trans_feat) * 0.001
+        loss.backward()
+        optimizer.step()
+        pred_choice = pred.data.max(1)[1]
+        correct = pred_choice.eq(target.data).cpu().sum()
+        print('[%d: %d/%d] train loss: %f accuracy: %f' % (epoch, i, num_batch, loss.item(), correct.item() / float(opt.batchSize)))
+
+        if i % 10 == 0:
+            j, data = next(enumerate(testdataloader, 0))
+            points, target = data
+            target = target[:, 0]
+            points = points.transpose(2, 1)
+            points, target = points.cuda(), target.cuda()
+            classifier = classifier.eval()
+            pred, _, _ = classifier(points)
+            loss = F.nll_loss(pred, target)
+            pred_choice = pred.data.max(1)[1]
+            correct = pred_choice.eq(target.data).cpu().sum()
+            print('[%d: %d/%d] %s loss: %f accuracy: %f' % (epoch, i, num_batch, blue('test'), loss.item(), correct.item()/float(opt.batchSize)))
+
+    torch.save(classifier.state_dict(), '%s/cls_model_%d.pth' % (opt.outf, epoch))
+
+total_correct = 0
+total_testset = 0
+for i,data in tqdm(enumerate(testdataloader, 0)):
+    points, target = data
+    target = target[:, 0]
+    points = points.transpose(2, 1)
+    points, target = points.cuda(), target.cuda()
+    classifier = classifier.eval()
+    pred, _, _ = classifier(points)
+    pred_choice = pred.data.max(1)[1]
+    correct = pred_choice.eq(target.data).cpu().sum()
+    total_correct += correct.item()
+    total_testset += points.size()[0]
+
+print("final accuracy {}".format(total_correct / float(total_testset)))