From 6b8abaa50f02fa45f84395b385fd5b6e6db627f9 Mon Sep 17 00:00:00 2001 From: Dell Du <18588220928@163.com> Date: Wed, 5 Dec 2018 10:45:10 +0800 Subject: [PATCH 01/11] Test base on pytorch-1.0 and fix bugs. --- README.md | 2 +- show_cls.py | 3 +-- train_classification.py | 3 +-- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 2c3bfbf66..4c09bbc6e 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ This repo is implementation for PointNet(https://arxiv.org/abs/1612.00593) in py bash build.sh #build C++ code for visualization bash download.sh #download dataset python train_classification.py #train 3D model classification -python python train_segmentation.py # train 3D model segmentaion +python train_segmentation.py # train 3D model segmentaion python show_seg.py --model seg/seg_model_20.pth # show segmentation results ``` diff --git a/show_cls.py b/show_cls.py index 6d054265e..793af9b2b 100644 --- a/show_cls.py +++ b/show_cls.py @@ -34,8 +34,7 @@ testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle = True) - -classifier = PointNetCls(k = len(test_dataset.classes), num_points = opt.num_points) +classifier = PointNetCls(k=len(test_dataset.classes)) classifier.cuda() classifier.load_state_dict(torch.load(opt.model)) classifier.eval() diff --git a/train_classification.py b/train_classification.py index a0b8c1188..264076fc7 100644 --- a/train_classification.py +++ b/train_classification.py @@ -55,8 +55,7 @@ pass -classifier = PointNetCls(k = num_classes, num_points = opt.num_points) - +classifier = PointNetCls(k=num_classes) if opt.model != '': classifier.load_state_dict(torch.load(opt.model)) From f6eeca147ae0cb9588ebb3f3284ea7210cf09110 Mon Sep 17 00:00:00 2001 From: Dell Du <18588220928@163.com> Date: Wed, 5 Dec 2018 10:47:23 +0800 Subject: [PATCH 02/11] Fix bug. --- show_cls.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/show_cls.py b/show_cls.py index 793af9b2b..43581fe95 100644 --- a/show_cls.py +++ b/show_cls.py @@ -50,4 +50,4 @@ 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))) + print('i:%d loss: %f accuracy: %f' % (i, loss.data.item(), correct / float(32))) From 05f2da94c312d560880370da5e3a022c9611ab69 Mon Sep 17 00:00:00 2001 From: Dell Du <18588220928@163.com> Date: Wed, 5 Dec 2018 11:03:42 +0800 Subject: [PATCH 03/11] Delete unused code and re-format. --- datasets.py | 22 ++--- pointnet.py | 18 ++-- show3d_balls.py | 194 +++++++++++++++++++++------------------- show_cls.py | 21 ++--- show_seg.py | 31 +++---- train_classification.py | 69 ++++++++------ train_segmentation.py | 76 +++++++++------- 7 files changed, 216 insertions(+), 215 deletions(-) diff --git a/datasets.py b/datasets.py index d96d6753a..fe8e01279 100644 --- a/datasets.py +++ b/datasets.py @@ -1,21 +1,17 @@ from __future__ import print_function import torch.utils.data as data -from PIL import Image import os import os.path -import errno import torch -import json -import codecs import numpy as np -import sys -import torchvision.transforms as transforms -import argparse -import json - class PartDataset(data.Dataset): - def __init__(self, root, npoints = 2500, classification = False, class_choice = None, train = True): + def __init__(self, + root, + npoints=2500, + classification=False, + class_choice=None, + train=True): self.npoints = npoints self.root = root self.catfile = os.path.join(self.root, 'synsetoffset2category.txt') @@ -28,8 +24,8 @@ def __init__(self, root, npoints = 2500, classification = False, class_choice = 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} + if not class_choice is None: + self.cat = {k: v for k, v in self.cat.items() if k in class_choice} self.meta = {} for item in self.cat: @@ -59,7 +55,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 diff --git a/pointnet.py b/pointnet.py index 37091cd84..7e24c7381 100644 --- a/pointnet.py +++ b/pointnet.py @@ -1,20 +1,10 @@ from __future__ import print_function -import argparse -import os -import random import torch import torch.nn as nn import torch.nn.parallel -import torch.backends.cudnn as cudnn -import torch.optim as optim import torch.utils.data -import torchvision.transforms as transforms -import torchvision.utils as vutils from torch.autograd import Variable -from PIL import Image import numpy as np -import matplotlib.pyplot as plt -import pdb import torch.nn.functional as F @@ -67,13 +57,14 @@ def __init__(self, global_feat = True): self.bn2 = nn.BatchNorm1d(128) self.bn3 = nn.BatchNorm1d(1024) self.global_feat = global_feat + + def forward(self, x): - batchsize = x.size()[0] n_pts = x.size()[2] trans = self.stn(x) - x = x.transpose(2,1) + x = x.transpose(2, 1) x = torch.bmm(x, trans) - x = x.transpose(2,1) + x = x.transpose(2, 1) x = F.relu(self.bn1(self.conv1(x))) pointfeat = x x = F.relu(self.bn2(self.conv2(x))) @@ -96,6 +87,7 @@ def __init__(self, k = 2): self.bn1 = nn.BatchNorm1d(512) self.bn2 = nn.BatchNorm1d(256) self.relu = nn.ReLU() + def forward(self, x): x, trans = self.feat(x) x = F.relu(self.bn1(self.fc1(x))) diff --git a/show3d_balls.py b/show3d_balls.py index 33a615838..a9b41ac3c 100644 --- a/show3d_balls.py +++ b/show3d_balls.py @@ -2,142 +2,148 @@ import ctypes as ct import cv2 import sys -showsz=800 -mousex,mousey=0.5,0.5 -zoom=1.0 -changed=True +showsz = 800 +mousex, mousey = 0.5, 0.5 +zoom = 1.0 +changed = True + def onmouse(*args): - global mousex,mousey,changed - y=args[1] - x=args[2] - mousex=x/float(showsz) - mousey=y/float(showsz) - changed=True + global mousex, mousey, changed + y = args[1] + x = args[2] + mousex = x / float(showsz) + mousey = y / float(showsz) + changed = True + cv2.namedWindow('show3d') -cv2.moveWindow('show3d',0,0) -cv2.setMouseCallback('show3d',onmouse) +cv2.moveWindow('show3d', 0, 0) +cv2.setMouseCallback('show3d', onmouse) -dll=np.ctypeslib.load_library('render_balls_so','.') +dll = np.ctypeslib.load_library('render_balls_so', '.') -def showpoints(xyz,c_gt=None, c_pred = None ,waittime=0,showrot=False,magnifyBlue=0,freezerot=False,background=(0,0,0),normalizecolor=True,ballradius=10): - global showsz,mousex,mousey,zoom,changed +def showpoints(xyz,c_gt=None, c_pred = None, waittime=0, + showrot=False, magnifyBlue=0, freezerot=False, background=(0,0,0), + normalizecolor=True, ballradius=10): + global showsz, mousex, mousey, zoom, changed xyz=xyz-xyz.mean(axis=0) radius=((xyz**2).sum(axis=-1)**0.5).max() xyz/=(radius*2.2)/showsz if c_gt is None: - c0=np.zeros((len(xyz),),dtype='float32')+255 - c1=np.zeros((len(xyz),),dtype='float32')+255 - c2=np.zeros((len(xyz),),dtype='float32')+255 + c0 = np.zeros((len(xyz), ), dtype='float32') + 255 + c1 = np.zeros((len(xyz), ), dtype='float32') + 255 + c2 = np.zeros((len(xyz), ), dtype='float32') + 255 else: - c0=c_gt[:,0] - c1=c_gt[:,1] - c2=c_gt[:,2] + c0 = c_gt[:, 0] + c1 = c_gt[:, 1] + c2 = c_gt[:, 2] if normalizecolor: - c0/=(c0.max()+1e-14)/255.0 - c1/=(c1.max()+1e-14)/255.0 - c2/=(c2.max()+1e-14)/255.0 + c0 /= (c0.max() + 1e-14) / 255.0 + c1 /= (c1.max() + 1e-14) / 255.0 + c2 /= (c2.max() + 1e-14) / 255.0 - c0=np.require(c0,'float32','C') - c1=np.require(c1,'float32','C') - c2=np.require(c2,'float32','C') + c0 = np.require(c0, 'float32', 'C') + c1 = np.require(c1, 'float32', 'C') + c2 = np.require(c2, 'float32', 'C') - show=np.zeros((showsz,showsz,3),dtype='uint8') + show = np.zeros((showsz, showsz, 3), dtype='uint8') def render(): rotmat=np.eye(3) if not freezerot: xangle=(mousey-0.5)*np.pi*1.2 else: xangle=0 - rotmat=rotmat.dot(np.array([ - [1.0,0.0,0.0], - [0.0,np.cos(xangle),-np.sin(xangle)], - [0.0,np.sin(xangle),np.cos(xangle)], + rotmat = rotmat.dot( + np.array([ + [1.0, 0.0, 0.0], + [0.0, np.cos(xangle), -np.sin(xangle)], + [0.0, np.sin(xangle), np.cos(xangle)], ])) if not freezerot: - yangle=(mousex-0.5)*np.pi*1.2 + yangle = (mousex - 0.5) * np.pi * 1.2 else: - yangle=0 - rotmat=rotmat.dot(np.array([ - [np.cos(yangle),0.0,-np.sin(yangle)], - [0.0,1.0,0.0], - [np.sin(yangle),0.0,np.cos(yangle)], + yangle = 0 + rotmat = rotmat.dot( + np.array([ + [np.cos(yangle), 0.0, -np.sin(yangle)], + [0.0, 1.0, 0.0], + [np.sin(yangle), 0.0, np.cos(yangle)], ])) - rotmat*=zoom - nxyz=xyz.dot(rotmat)+[showsz/2,showsz/2,0] + rotmat *= zoom + nxyz = xyz.dot(rotmat) + [showsz / 2, showsz / 2, 0] - ixyz=nxyz.astype('int32') - show[:]=background + ixyz = nxyz.astype('int32') + show[:] = background dll.render_ball( - ct.c_int(show.shape[0]), - ct.c_int(show.shape[1]), - show.ctypes.data_as(ct.c_void_p), - ct.c_int(ixyz.shape[0]), - ixyz.ctypes.data_as(ct.c_void_p), - c0.ctypes.data_as(ct.c_void_p), - c1.ctypes.data_as(ct.c_void_p), - c2.ctypes.data_as(ct.c_void_p), - ct.c_int(ballradius) - ) + ct.c_int(show.shape[0]), ct.c_int(show.shape[1]), + show.ctypes.data_as(ct.c_void_p), ct.c_int(ixyz.shape[0]), + ixyz.ctypes.data_as(ct.c_void_p), c0.ctypes.data_as(ct.c_void_p), + c1.ctypes.data_as(ct.c_void_p), c2.ctypes.data_as(ct.c_void_p), + ct.c_int(ballradius)) - if magnifyBlue>0: - show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],1,axis=0)) - if magnifyBlue>=2: - show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],-1,axis=0)) - show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],1,axis=1)) - if magnifyBlue>=2: - show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],-1,axis=1)) + if magnifyBlue > 0: + show[:, :, 0] = np.maximum(show[:, :, 0], np.roll( + show[:, :, 0], 1, axis=0)) + if magnifyBlue >= 2: + show[:, :, 0] = np.maximum(show[:, :, 0], + np.roll(show[:, :, 0], -1, axis=0)) + show[:, :, 0] = np.maximum(show[:, :, 0], np.roll( + show[:, :, 0], 1, axis=1)) + if magnifyBlue >= 2: + show[:, :, 0] = np.maximum(show[:, :, 0], + np.roll(show[:, :, 0], -1, axis=1)) if showrot: - cv2.putText(show,'xangle %d'%(int(xangle/np.pi*180)),(30,showsz-30),0,0.5,cv2.cv.CV_RGB(255,0,0)) - cv2.putText(show,'yangle %d'%(int(yangle/np.pi*180)),(30,showsz-50),0,0.5,cv2.cv.CV_RGB(255,0,0)) - cv2.putText(show,'zoom %d%%'%(int(zoom*100)),(30,showsz-70),0,0.5,cv2.cv.CV_RGB(255,0,0)) - changed=True + cv2.putText(show, 'xangle %d' % (int(xangle / np.pi * 180)), + (30, showsz - 30), 0, 0.5, cv2.cv.CV_RGB(255, 0, 0)) + cv2.putText(show, 'yangle %d' % (int(yangle / np.pi * 180)), + (30, showsz - 50), 0, 0.5, cv2.cv.CV_RGB(255, 0, 0)) + cv2.putText(show, 'zoom %d%%' % (int(zoom * 100)), (30, showsz - 70), 0, + 0.5, cv2.cv.CV_RGB(255, 0, 0)) + changed = True while True: if changed: render() - changed=False - cv2.imshow('show3d',show) - if waittime==0: - cmd=cv2.waitKey(10)%256 + changed = False + cv2.imshow('show3d', show) + if waittime == 0: + cmd = cv2.waitKey(10) % 256 else: - cmd=cv2.waitKey(waittime)%256 - if cmd==ord('q'): + cmd = cv2.waitKey(waittime) % 256 + if cmd == ord('q'): break - elif cmd==ord('Q'): + elif cmd == ord('Q'): sys.exit(0) - if cmd==ord('t') or cmd == ord('p'): + if cmd == ord('t') or cmd == ord('p'): if cmd == ord('t'): if c_gt is None: - c0=np.zeros((len(xyz),),dtype='float32')+255 - c1=np.zeros((len(xyz),),dtype='float32')+255 - c2=np.zeros((len(xyz),),dtype='float32')+255 + c0 = np.zeros((len(xyz), ), dtype='float32') + 255 + c1 = np.zeros((len(xyz), ), dtype='float32') + 255 + c2 = np.zeros((len(xyz), ), dtype='float32') + 255 else: - c0=c_gt[:,0] - c1=c_gt[:,1] - c2=c_gt[:,2] + c0 = c_gt[:, 0] + c1 = c_gt[:, 1] + c2 = c_gt[:, 2] else: if c_pred is None: - c0=np.zeros((len(xyz),),dtype='float32')+255 - c1=np.zeros((len(xyz),),dtype='float32')+255 - c2=np.zeros((len(xyz),),dtype='float32')+255 + c0 = np.zeros((len(xyz), ), dtype='float32') + 255 + c1 = np.zeros((len(xyz), ), dtype='float32') + 255 + c2 = np.zeros((len(xyz), ), dtype='float32') + 255 else: - c0=c_pred[:,0] - c1=c_pred[:,1] - c2=c_pred[:,2] + c0 = c_pred[:, 0] + c1 = c_pred[:, 1] + c2 = c_pred[:, 2] if normalizecolor: - c0/=(c0.max()+1e-14)/255.0 - c1/=(c1.max()+1e-14)/255.0 - c2/=(c2.max()+1e-14)/255.0 - c0=np.require(c0,'float32','C') - c1=np.require(c1,'float32','C') - c2=np.require(c2,'float32','C') + c0 /= (c0.max() + 1e-14) / 255.0 + c1 /= (c1.max() + 1e-14) / 255.0 + c2 /= (c2.max() + 1e-14) / 255.0 + c0 = np.require(c0, 'float32', 'C') + c1 = np.require(c1, 'float32', 'C') + c2 = np.require(c2, 'float32', 'C') changed = True - - if cmd==ord('n'): zoom*=1.1 changed=True @@ -152,7 +158,7 @@ def render(): if waittime!=0: break return cmd -if __name__=='__main__': - np.random.seed(100) - showpoints(np.random.randn(2500,3)) +if __name__ == '__main__': + np.random.seed(100) + showpoints(np.random.randn(2500, 3)) diff --git a/show_cls.py b/show_cls.py index 43581fe95..47168d9c0 100644 --- a/show_cls.py +++ b/show_cls.py @@ -1,22 +1,12 @@ from __future__ import print_function import argparse -import os -import random -import numpy as np import torch -import torch.nn as nn import torch.nn.parallel -import torch.backends.cudnn as cudnn -import torch.optim as optim import torch.utils.data -import torchvision.datasets as dset -import torchvision.transforms as transforms -import torchvision.utils as vutils from torch.autograd import Variable from datasets import PartDataset from pointnet import PointNetCls import torch.nn.functional as F -import matplotlib.pyplot as plt #showpoints(np.random.randn(2500,3), c1 = np.random.uniform(0,1,size = (2500))) @@ -28,11 +18,16 @@ opt = parser.parse_args() -print (opt) +print(opt) -test_dataset = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0' , train = False, classification = True, npoints = opt.num_points) +test_dataset = PartDataset( + root='shapenetcore_partanno_segmentation_benchmark_v0', + train=False, + classification=True, + npoints=opt.num_points) -testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle = True) +testdataloader = torch.utils.data.DataLoader( + test_dataset, batch_size=32, shuffle=True) classifier = PointNetCls(k=len(test_dataset.classes)) classifier.cuda() diff --git a/show_seg.py b/show_seg.py index aef3dfa75..58b493adf 100644 --- a/show_seg.py +++ b/show_seg.py @@ -1,22 +1,13 @@ from __future__ import print_function -from show3d_balls import * +from show3d_balls import showpoints import argparse -import os -import random import numpy as np import torch -import torch.nn as nn import torch.nn.parallel -import torch.backends.cudnn as cudnn -import torch.optim as optim import torch.utils.data -import torchvision.datasets as dset -import torchvision.transforms as transforms -import torchvision.utils as vutils from torch.autograd import Variable from datasets import PartDataset from pointnet import PointNetDenseCls -import torch.nn.functional as F import matplotlib.pyplot as plt @@ -24,19 +15,22 @@ parser = argparse.ArgumentParser() -parser.add_argument('--model', type=str, default = '', help='model path') -parser.add_argument('--idx', type=int, default = 0, help='model index') +parser.add_argument('--model', type=str, default='', help='model path') +parser.add_argument('--idx', type=int, default=0, help='model index') opt = parser.parse_args() -print (opt) +print(opt) -d = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', class_choice = ['Chair'], train = False) +d = PartDataset( + root='shapenetcore_partanno_segmentation_benchmark_v0', + class_choice=['Airplane'], + train=False) idx = opt.idx -print("model %d/%d" %( idx, len(d))) +print("model %d/%d" % (idx, len(d))) point, seg = d[idx] print(point.size(), seg.size()) @@ -46,14 +40,14 @@ cmap = plt.cm.get_cmap("hsv", 10) -cmap = np.array([cmap(i) for i in range(10)])[:,:3] +cmap = np.array([cmap(i) for i in range(10)])[:, :3] gt = cmap[seg.numpy() - 1, :] -classifier = PointNetDenseCls(k = 4) +classifier = PointNetDenseCls(k=4) classifier.load_state_dict(torch.load(opt.model)) classifier.eval() -point = point.transpose(1,0).contiguous() +point = point.transpose(1, 0).contiguous() point = Variable(point.view(1, point.size()[0], point.size()[1])) pred, _ = classifier(point) @@ -65,4 +59,3 @@ #print(pred_color.shape) showpoints(point_np, gt, pred_color) - diff --git a/train_classification.py b/train_classification.py index 264076fc7..d99e839a0 100644 --- a/train_classification.py +++ b/train_classification.py @@ -2,16 +2,10 @@ import argparse import os import random -import numpy as np import torch -import torch.nn as nn import torch.nn.parallel -import torch.backends.cudnn as cudnn import torch.optim as optim import torch.utils.data -import torchvision.datasets as dset -import torchvision.transforms as transforms -import torchvision.utils as vutils from torch.autograd import Variable from datasets import PartDataset from pointnet import PointNetCls @@ -20,30 +14,47 @@ 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=25, 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( + '--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=25, 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') opt = parser.parse_args() -print (opt) +print(opt) -blue = lambda x:'\033[94m' + x + '\033[0m' +blue = lambda x: '\033[94m' + x + '\033[0m' -opt.manualSeed = random.randint(1, 10000) # fix seed +opt.manualSeed = random.randint(1, 10000) # fix seed print("Random Seed: ", opt.manualSeed) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) -dataset = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', classification = True, npoints = opt.num_points) -dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize, - shuffle=True, num_workers=int(opt.workers)) - -test_dataset = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', classification = True, train = False, npoints = opt.num_points) -testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=opt.batchSize, - shuffle=True, num_workers=int(opt.workers)) +dataset = PartDataset( + root='shapenetcore_partanno_segmentation_benchmark_v0', + classification=True, + npoints=opt.num_points) +dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=opt.batchSize, + shuffle=True, + num_workers=int(opt.workers)) + +test_dataset = PartDataset( + root='shapenetcore_partanno_segmentation_benchmark_v0', + classification=True, + train=False, + npoints=opt.num_points) +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) @@ -64,13 +75,13 @@ optimizer = optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) classifier.cuda() -num_batch = len(dataset)/opt.batchSize +num_batch = len(dataset) / opt.batchSize for epoch in range(opt.nepoch): for i, data in enumerate(dataloader, 0): points, target = data - points, target = Variable(points), Variable(target[:,0]) - points = points.transpose(2,1) + points, target = Variable(points), Variable(target[:, 0]) + points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() @@ -80,19 +91,19 @@ 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))) + 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 - points, target = Variable(points), Variable(target[:,0]) - points = points.transpose(2,1) + points, target = Variable(points), Variable(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))) + 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)) diff --git a/train_segmentation.py b/train_segmentation.py index 2fd6c5a66..f6368ef51 100644 --- a/train_segmentation.py +++ b/train_segmentation.py @@ -2,46 +2,55 @@ import argparse import os import random -import numpy as np import torch -import torch.nn as nn import torch.nn.parallel -import torch.backends.cudnn as cudnn import torch.optim as optim import torch.utils.data -import torchvision.datasets as dset -import torchvision.transforms as transforms -import torchvision.utils as vutils from torch.autograd import Variable from datasets import PartDataset from pointnet import PointNetDenseCls import torch.nn.functional as F - parser = argparse.ArgumentParser() -parser.add_argument('--batchSize', type=int, default=32, 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=25, help='number of epochs to train for') -parser.add_argument('--outf', type=str, default='seg', help='output folder') -parser.add_argument('--model', type=str, default = '', help='model path') +parser.add_argument( + '--batchSize', type=int, default=32, 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=25, help='number of epochs to train for') +parser.add_argument('--outf', type=str, default='seg', help='output folder') +parser.add_argument('--model', type=str, default='', help='model path') opt = parser.parse_args() -print (opt) +print(opt) -opt.manualSeed = random.randint(1, 10000) # fix seed +opt.manualSeed = random.randint(1, 10000) # fix seed print("Random Seed: ", opt.manualSeed) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) -dataset = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', classification = False, class_choice = ['Chair']) -dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize, - shuffle=True, num_workers=int(opt.workers)) - -test_dataset = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', classification = False, class_choice = ['Chair'], train = False) -testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=opt.batchSize, - shuffle=True, num_workers=int(opt.workers)) +dataset = PartDataset( + root='shapenetcore_partanno_segmentation_benchmark_v0', + classification=False, + class_choice=['Chair']) +dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=opt.batchSize, + shuffle=True, + num_workers=int(opt.workers)) + +test_dataset = PartDataset( + root='shapenetcore_partanno_segmentation_benchmark_v0', + classification=False, + class_choice=['Chair'], + train=False) +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 = dataset.num_seg_classes @@ -51,10 +60,9 @@ except OSError: pass -blue = lambda x:'\033[94m' + x + '\033[0m' - +blue = lambda x: '\033[94m' + x + '\033[0m' -classifier = PointNetDenseCls(k = num_classes) +classifier = PointNetDenseCls(k=num_classes) if opt.model != '': classifier.load_state_dict(torch.load(opt.model)) @@ -62,41 +70,41 @@ optimizer = optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) classifier.cuda() -num_batch = len(dataset)/opt.batchSize +num_batch = len(dataset) / opt.batchSize for epoch in range(opt.nepoch): for i, data in enumerate(dataloader, 0): points, target = data points, target = Variable(points), Variable(target) - points = points.transpose(2,1) - points, target = points.cuda(), target.cuda() + points = points.transpose(2, 1) + points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() pred, _ = classifier(points) pred = pred.view(-1, num_classes) - target = target.view(-1,1)[:,0] - 1 + target = target.view(-1, 1)[:, 0] - 1 #print(pred.size(), target.size()) loss = F.nll_loss(pred, target) 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 * 2500))) - + print('[%d: %d/%d] train loss: %f accuracy: %f' % (epoch, i, num_batch, loss.item(), correct.item()/float(opt.batchSize * 2500))) + if i % 10 == 0: j, data = next(enumerate(testdataloader, 0)) points, target = data points, target = Variable(points), Variable(target) - points = points.transpose(2,1) + points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() classifier = classifier.eval() pred, _ = classifier(points) pred = pred.view(-1, num_classes) - target = target.view(-1,1)[:,0] - 1 + target = target.view(-1, 1)[:, 0] - 1 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 * 2500))) - + print('[%d: %d/%d] %s loss: %f accuracy: %f' % (epoch, i, num_batch, blue('test'), loss.item(), correct.item()/float(opt.batchSize * 2500))) + torch.save(classifier.state_dict(), '%s/seg_model_%d.pth' % (opt.outf, epoch)) \ No newline at end of file From 033bb3f6133094ebb8e75a318979e1bc4360c02e Mon Sep 17 00:00:00 2001 From: Dell Du <18588220928@163.com> Date: Wed, 5 Dec 2018 11:08:17 +0800 Subject: [PATCH 04/11] Add shortcuts. --- show_seg.sh | 2 ++ train_cls.sh | 2 ++ train_seg.sh | 2 ++ 3 files changed, 6 insertions(+) create mode 100755 show_seg.sh create mode 100755 train_cls.sh create mode 100755 train_seg.sh diff --git a/show_seg.sh b/show_seg.sh new file mode 100755 index 000000000..92ee663cb --- /dev/null +++ b/show_seg.sh @@ -0,0 +1,2 @@ +python show_seg.py --model seg/seg_model_20.pth + diff --git a/train_cls.sh b/train_cls.sh new file mode 100755 index 000000000..9051da944 --- /dev/null +++ b/train_cls.sh @@ -0,0 +1,2 @@ +python train_classification.py + diff --git a/train_seg.sh b/train_seg.sh new file mode 100755 index 000000000..bd12e816c --- /dev/null +++ b/train_seg.sh @@ -0,0 +1,2 @@ +python train_segmentation.py + From 0ce9a4badb033c4d0de042c1acb409e1e59ee69e Mon Sep 17 00:00:00 2001 From: fxia22 Date: Mon, 17 Dec 2018 19:17:36 -0800 Subject: [PATCH 05/11] change folder structure --- build.sh | 1 - datasets.py => pointnet/dataset.py | 6 ++++-- pointnet.py => pointnet/model.py | 0 scripts/build.sh | 5 +++++ download.sh => scripts/download.sh | 5 +++++ setup.py | 9 +++++++++ show_seg.sh | 2 -- train_cls.sh | 2 -- train_seg.sh | 2 -- render_balls_so.cpp => utils/render_balls_so.cpp | 0 show3d_balls.py => utils/show3d_balls.py | 0 show_cls.py => utils/show_cls.py | 4 ++-- show_seg.py => utils/show_seg.py | 4 ++-- train_classification.py => utils/train_classification.py | 4 ++-- train_segmentation.py => utils/train_segmentation.py | 6 +++--- 15 files changed, 32 insertions(+), 18 deletions(-) delete mode 100644 build.sh rename datasets.py => pointnet/dataset.py (93%) rename pointnet.py => pointnet/model.py (100%) create mode 100755 scripts/build.sh rename download.sh => scripts/download.sh (75%) mode change 100644 => 100755 create mode 100644 setup.py delete mode 100755 show_seg.sh delete mode 100755 train_cls.sh delete mode 100755 train_seg.sh rename render_balls_so.cpp => utils/render_balls_so.cpp (100%) rename show3d_balls.py => utils/show3d_balls.py (100%) rename show_cls.py => utils/show_cls.py (94%) rename show_seg.py => utils/show_seg.py (94%) rename train_classification.py => utils/train_classification.py (97%) rename train_segmentation.py => utils/train_segmentation.py (97%) diff --git a/build.sh b/build.sh deleted file mode 100644 index 3bceac2a8..000000000 --- a/build.sh +++ /dev/null @@ -1 +0,0 @@ -g++ -std=c++11 render_balls_so.cpp -o render_balls_so.so -shared -fPIC -O2 -D_GLIBCXX_USE_CXX11_ABI=0 diff --git a/datasets.py b/pointnet/dataset.py similarity index 93% rename from datasets.py rename to pointnet/dataset.py index fe8e01279..953b585e0 100644 --- a/datasets.py +++ b/pointnet/dataset.py @@ -4,6 +4,7 @@ import os.path import torch import numpy as np +import sys class PartDataset(data.Dataset): def __init__(self, @@ -86,13 +87,14 @@ def __len__(self): if __name__ == '__main__': + datapath = sys.argv[1] print('test') - d = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', class_choice = ['Chair']) + d = PartDataset(root = datapath, class_choice = ['Chair']) print(len(d)) ps, seg = d[0] print(ps.size(), ps.type(), seg.size(),seg.type()) - d = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', classification = True) + d = PartDataset(root = datapath, classification = True) print(len(d)) ps, cls = d[0] print(ps.size(), ps.type(), cls.size(),cls.type()) diff --git a/pointnet.py b/pointnet/model.py similarity index 100% rename from pointnet.py rename to pointnet/model.py diff --git a/scripts/build.sh b/scripts/build.sh new file mode 100755 index 000000000..8c723ed21 --- /dev/null +++ b/scripts/build.sh @@ -0,0 +1,5 @@ +SCRIPT=`realpath $0` +SCRIPTPATH=`dirname $SCRIPT` +echo $SCRIPTPATH + +g++ -std=c++11 $SCRIPTPATH/../utils/render_balls_so.cpp -o $SCRIPTPATH/../utils/render_balls_so.so -shared -fPIC -O2 -D_GLIBCXX_USE_CXX11_ABI=0 diff --git a/download.sh b/scripts/download.sh old mode 100644 new mode 100755 similarity index 75% rename from download.sh rename to scripts/download.sh index d9cde8f69..abc9b71f3 --- a/download.sh +++ b/scripts/download.sh @@ -1,3 +1,8 @@ +SCRIPT=`realpath $0` +SCRIPTPATH=`dirname $SCRIPT` + +cd $SCRIPTPATH/.. wget https://shapenet.cs.stanford.edu/ericyi/shapenetcore_partanno_segmentation_benchmark_v0.zip --no-check-certificate unzip shapenetcore_partanno_segmentation_benchmark_v0.zip rm shapenetcore_partanno_segmentation_benchmark_v0.zip +cd - diff --git a/setup.py b/setup.py new file mode 100644 index 000000000..967cf9550 --- /dev/null +++ b/setup.py @@ -0,0 +1,9 @@ +# install using 'pip install -e .' + +from setuptools import setup + +setup(name='pointnet', + packages=['pointnet'], + package_dir={'pointnet': 'pointnet'}, + install_requires=['torch'], + version='0.0.1') diff --git a/show_seg.sh b/show_seg.sh deleted file mode 100755 index 92ee663cb..000000000 --- a/show_seg.sh +++ /dev/null @@ -1,2 +0,0 @@ -python show_seg.py --model seg/seg_model_20.pth - diff --git a/train_cls.sh b/train_cls.sh deleted file mode 100755 index 9051da944..000000000 --- a/train_cls.sh +++ /dev/null @@ -1,2 +0,0 @@ -python train_classification.py - diff --git a/train_seg.sh b/train_seg.sh deleted file mode 100755 index bd12e816c..000000000 --- a/train_seg.sh +++ /dev/null @@ -1,2 +0,0 @@ -python train_segmentation.py - diff --git a/render_balls_so.cpp b/utils/render_balls_so.cpp similarity index 100% rename from render_balls_so.cpp rename to utils/render_balls_so.cpp diff --git a/show3d_balls.py b/utils/show3d_balls.py similarity index 100% rename from show3d_balls.py rename to utils/show3d_balls.py diff --git a/show_cls.py b/utils/show_cls.py similarity index 94% rename from show_cls.py rename to utils/show_cls.py index 47168d9c0..5dced8b97 100644 --- a/show_cls.py +++ b/utils/show_cls.py @@ -4,8 +4,8 @@ import torch.nn.parallel import torch.utils.data from torch.autograd import Variable -from datasets import PartDataset -from pointnet import PointNetCls +from pointnet.dataset import PartDataset +from pointnet.model import PointNetCls import torch.nn.functional as F diff --git a/show_seg.py b/utils/show_seg.py similarity index 94% rename from show_seg.py rename to utils/show_seg.py index 58b493adf..302610f3f 100644 --- a/show_seg.py +++ b/utils/show_seg.py @@ -6,8 +6,8 @@ import torch.nn.parallel import torch.utils.data from torch.autograd import Variable -from datasets import PartDataset -from pointnet import PointNetDenseCls +from pointnet.dataset import PartDataset +from pointnet.model import PointNetDenseCls import matplotlib.pyplot as plt diff --git a/train_classification.py b/utils/train_classification.py similarity index 97% rename from train_classification.py rename to utils/train_classification.py index d99e839a0..571d1313d 100644 --- a/train_classification.py +++ b/utils/train_classification.py @@ -7,8 +7,8 @@ import torch.optim as optim import torch.utils.data from torch.autograd import Variable -from datasets import PartDataset -from pointnet import PointNetCls +from pointnet.dataset import PartDataset +from pointnet.model import PointNetCls import torch.nn.functional as F diff --git a/train_segmentation.py b/utils/train_segmentation.py similarity index 97% rename from train_segmentation.py rename to utils/train_segmentation.py index f6368ef51..6d6e7757c 100644 --- a/train_segmentation.py +++ b/utils/train_segmentation.py @@ -7,8 +7,8 @@ import torch.optim as optim import torch.utils.data from torch.autograd import Variable -from datasets import PartDataset -from pointnet import PointNetDenseCls +from pointnet.dataset import PartDataset +from pointnet.model import PointNetDenseCls import torch.nn.functional as F @@ -107,4 +107,4 @@ 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 * 2500))) - torch.save(classifier.state_dict(), '%s/seg_model_%d.pth' % (opt.outf, epoch)) \ No newline at end of file + torch.save(classifier.state_dict(), '%s/seg_model_%d.pth' % (opt.outf, epoch)) From 523f07fffc7268d94a21b3f5e70371569f3b5967 Mon Sep 17 00:00:00 2001 From: fxia22 Date: Sat, 2 Mar 2019 15:46:39 -0800 Subject: [PATCH 06/11] add some benchmarking code --- .gitignore | 1 + pointnet/__init__.py | 0 pointnet/dataset.py | 16 ++++++++++- pointnet/model.py | 5 ++-- utils/show_seg.py | 17 +++++------- utils/train_classification.py | 35 +++++++++++++++++------- utils/train_segmentation.py | 50 ++++++++++++++++++++++++++++------- 7 files changed, 92 insertions(+), 32 deletions(-) create mode 100644 pointnet/__init__.py diff --git a/.gitignore b/.gitignore index 2614c1af6..fe0185166 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,4 @@ shapenetcore_partanno_segmentation_benchmark_v0/ .idea* cls/ seg/ +*.egg-info/ diff --git a/pointnet/__init__.py b/pointnet/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pointnet/dataset.py b/pointnet/dataset.py index 953b585e0..6a0aff2db 100644 --- a/pointnet/dataset.py +++ b/pointnet/dataset.py @@ -12,11 +12,13 @@ def __init__(self, npoints=2500, classification=False, class_choice=None, - train=True): + train=True, + 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 @@ -73,10 +75,22 @@ def __getitem__(self, index): 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: diff --git a/pointnet/model.py b/pointnet/model.py index 7e24c7381..7be4b7cfa 100644 --- a/pointnet/model.py +++ b/pointnet/model.py @@ -84,6 +84,7 @@ def __init__(self, k = 2): 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() @@ -91,9 +92,9 @@ def __init__(self, k = 2): def forward(self, x): x, trans = self.feat(x) x = F.relu(self.bn1(self.fc1(x))) - x = F.relu(self.bn2(self.fc2(x))) + x = F.relu(self.bn2(self.dropout(self.fc2(x)))) x = self.fc3(x) - return F.log_softmax(x, dim=0), trans + return F.log_softmax(x, dim=1), trans class PointNetDenseCls(nn.Module): def __init__(self, k = 2): diff --git a/utils/show_seg.py b/utils/show_seg.py index 302610f3f..b6b01c624 100644 --- a/utils/show_seg.py +++ b/utils/show_seg.py @@ -17,34 +17,31 @@ parser.add_argument('--model', type=str, default='', help='model path') parser.add_argument('--idx', type=int, default=0, help='model index') - - +parser.add_argument('--dataset', type=str, default='', help='dataset path') +parser.add_argument('--class_choice', type=str, default='', help='class choice') opt = parser.parse_args() print(opt) d = PartDataset( - root='shapenetcore_partanno_segmentation_benchmark_v0', - class_choice=['Airplane'], + root=opt.dataset, + class_choice=[opt.class_choice], train=False) idx = opt.idx print("model %d/%d" % (idx, len(d))) - point, seg = d[idx] print(point.size(), seg.size()) - point_np = point.numpy() - - cmap = plt.cm.get_cmap("hsv", 10) cmap = np.array([cmap(i) for i in range(10)])[:, :3] gt = cmap[seg.numpy() - 1, :] -classifier = PointNetDenseCls(k=4) -classifier.load_state_dict(torch.load(opt.model)) +state_dict = torch.load(opt.model) +classifier = PointNetDenseCls(k= state_dict['conv4.weight'].size()[0] ) +classifier.load_state_dict(state_dict) classifier.eval() point = point.transpose(1, 0).contiguous() diff --git a/utils/train_classification.py b/utils/train_classification.py index 571d1313d..07f758143 100644 --- a/utils/train_classification.py +++ b/utils/train_classification.py @@ -6,11 +6,10 @@ import torch.nn.parallel import torch.optim as optim import torch.utils.data -from torch.autograd import Variable from pointnet.dataset import PartDataset from pointnet.model import PointNetCls import torch.nn.functional as F - +from tqdm import tqdm parser = argparse.ArgumentParser() @@ -21,9 +20,10 @@ parser.add_argument( '--workers', type=int, help='number of data loading workers', default=4) parser.add_argument( - '--nepoch', type=int, default=25, help='number of epochs to train for') + '--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") opt = parser.parse_args() print(opt) @@ -36,7 +36,7 @@ torch.manual_seed(opt.manualSeed) dataset = PartDataset( - root='shapenetcore_partanno_segmentation_benchmark_v0', + root=opt.dataset, classification=True, npoints=opt.num_points) dataloader = torch.utils.data.DataLoader( @@ -46,7 +46,7 @@ num_workers=int(opt.workers)) test_dataset = PartDataset( - root='shapenetcore_partanno_segmentation_benchmark_v0', + root=opt.dataset, classification=True, train=False, npoints=opt.num_points) @@ -65,22 +65,23 @@ except OSError: pass - classifier = PointNetCls(k=num_classes) if opt.model != '': classifier.load_state_dict(torch.load(opt.model)) -optimizer = optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) +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 - points, target = Variable(points), Variable(target[:, 0]) + target = target[:, 0] points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() @@ -96,7 +97,7 @@ if i % 10 == 0: j, data = next(enumerate(testdataloader, 0)) points, target = data - points, target = Variable(points), Variable(target[:, 0]) + target = target[:, 0] points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() classifier = classifier.eval() @@ -107,3 +108,19 @@ 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))) \ No newline at end of file diff --git a/utils/train_segmentation.py b/utils/train_segmentation.py index 6d6e7757c..b6be26576 100644 --- a/utils/train_segmentation.py +++ b/utils/train_segmentation.py @@ -6,10 +6,11 @@ import torch.nn.parallel import torch.optim as optim import torch.utils.data -from torch.autograd import Variable from pointnet.dataset import PartDataset from pointnet.model import PointNetDenseCls import torch.nn.functional as F +from tqdm import tqdm +import numpy as np parser = argparse.ArgumentParser() @@ -21,6 +22,8 @@ '--nepoch', type=int, default=25, help='number of epochs to train for') parser.add_argument('--outf', type=str, default='seg', 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('--class_choice', type=str, default='Chair', help="class_choice") opt = parser.parse_args() @@ -32,9 +35,9 @@ torch.manual_seed(opt.manualSeed) dataset = PartDataset( - root='shapenetcore_partanno_segmentation_benchmark_v0', + root=opt.dataset, classification=False, - class_choice=['Chair']) + class_choice=[opt.class_choice]) dataloader = torch.utils.data.DataLoader( dataset, batch_size=opt.batchSize, @@ -42,9 +45,9 @@ num_workers=int(opt.workers)) test_dataset = PartDataset( - root='shapenetcore_partanno_segmentation_benchmark_v0', + root=opt.dataset, classification=False, - class_choice=['Chair'], + class_choice=[opt.class_choice], train=False) testdataloader = torch.utils.data.DataLoader( test_dataset, @@ -67,15 +70,16 @@ if opt.model != '': classifier.load_state_dict(torch.load(opt.model)) -optimizer = optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) +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 - points, target = Variable(points), Variable(target) points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() @@ -94,17 +98,43 @@ if i % 10 == 0: j, data = next(enumerate(testdataloader, 0)) points, target = data - points, target = Variable(points), Variable(target) points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() classifier = classifier.eval() pred, _ = classifier(points) pred = pred.view(-1, num_classes) target = target.view(-1, 1)[:, 0] - 1 - 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 * 2500))) - torch.save(classifier.state_dict(), '%s/seg_model_%d.pth' % (opt.outf, epoch)) + torch.save(classifier.state_dict(), '%s/seg_model_%s_%d.pth' % (opt.outf, opt.class_choice, epoch)) + +## benchmark mIOU +shape_ious = [] +for i,data in tqdm(enumerate(testdataloader, 0)): + points, target = data + points = points.transpose(2, 1) + points, target = points.cuda(), target.cuda() + classifier = classifier.eval() + pred, _ = classifier(points) + pred_choice = pred.data.max(2)[1] + + pred_np = pred_choice.cpu().data.numpy() + target_np = target.cpu().data.numpy() - 1 + + for shape_idx in range(target_np.shape[0]): + parts = np.unique(target_np[shape_idx]) + part_ious = [] + for part in parts: + I = np.sum(np.logical_and(pred_np[shape_idx] == part, target_np[shape_idx] == part)) + U = np.sum(np.logical_or(pred_np[shape_idx] == part, target_np[shape_idx] == part)) + if U == 0: + iou = 0 + else: + iou = I / float(U) + part_ious.append(iou) + shape_ious.append(np.mean(part_ious)) + +print("mIOU for class {}: {}".format(opt.class_choice, np.mean(shape_ious))) \ No newline at end of file From 3c7e2cd4708ffb24e6aafdcf99af1f4b601b82b4 Mon Sep 17 00:00:00 2001 From: fxia22 Date: Mon, 4 Mar 2019 21:23:23 -0800 Subject: [PATCH 07/11] use shapenet split --- pointnet/dataset.py | 92 +++++++++++++++++++++++------------ pointnet/num_seg_classes.txt | 16 ++++++ utils/show_cls.py | 4 +- utils/show_seg.py | 4 +- utils/train_classification.py | 8 +-- utils/train_segmentation.py | 12 ++--- 6 files changed, 92 insertions(+), 44 deletions(-) create mode 100644 pointnet/num_seg_classes.txt diff --git a/pointnet/dataset.py b/pointnet/dataset.py index 6a0aff2db..7150201e8 100644 --- a/pointnet/dataset.py +++ b/pointnet/dataset.py @@ -5,23 +5,59 @@ import torch import numpy as np import sys - -class PartDataset(data.Dataset): +from tqdm import tqdm +import json + +def get_segmentation_classes(root): + catfile = os.path.join(root, 'synsetoffset2category.txt') + cat = {} + meta = {} + + with open(catfile, 'r') as f: + for line in f: + ls = line.strip().split() + cat[ls[0]] = ls[1] + + for item in cat: + dir_seg = os.path.join(root, cat[item], 'points_label') + dir_point = os.path.join(root, cat[item], 'points') + fns = sorted(os.listdir(dir_point)) + meta[item] = [] + for fn in fns: + token = (os.path.splitext(os.path.basename(fn))[0]) + meta[item].append((os.path.join(dir_point, token + '.pts'), os.path.join(dir_seg, token + '.seg'))) + + with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'num_seg_classes.txt'), 'w') as f: + for item in cat: + datapath = [] + num_seg_classes = 0 + for fn in meta[item]: + datapath.append((item, fn[0], fn[1])) + + for i in tqdm(range(len(datapath))): + l = len(np.unique(np.loadtxt(datapath[i][-1]).astype(np.uint8))) + if l > num_seg_classes: + num_seg_classes = l + + print("category {} num segmentation classes {}".format(item, num_seg_classes)) + f.write("{}\t{}\n".format(item, num_seg_classes)) + +class ShapeNetDataset(data.Dataset): def __init__(self, root, npoints=2500, classification=False, class_choice=None, - train=True, + 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() @@ -30,40 +66,34 @@ def __init__(self, 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: - #print('category', item) self.meta[item] = [] - dir_point = os.path.join(self.root, self.cat[item], 'points') - dir_seg = os.path.join(self.root, self.cat[item], 'points_label') - #print(dir_point, dir_seg) - fns = sorted(os.listdir(dir_point)) - if train: - fns = fns[:int(len(fns) * 0.9)] - else: - fns = fns[int(len(fns) * 0.9):] - - #print(os.path.basename(fns)) - for fn in fns: - token = (os.path.splitext(os.path.basename(fn))[0]) - self.meta[item].append((os.path.join(dir_point, token + '.pts'), os.path.join(dir_seg, token + '.seg'))) + + 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) - self.num_seg_classes = 0 - if not self.classification: - 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 - #print(self.num_seg_classes) - + with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), '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] @@ -103,12 +133,14 @@ def __len__(self): if __name__ == '__main__': datapath = sys.argv[1] print('test') - d = PartDataset(root = datapath, class_choice = ['Chair']) + d = ShapeNetDataset(root = datapath, class_choice = ['Chair']) print(len(d)) ps, seg = d[0] print(ps.size(), ps.type(), seg.size(),seg.type()) - d = PartDataset(root = datapath, classification = True) + d = ShapeNetDataset(root = datapath, classification = True) print(len(d)) ps, cls = d[0] print(ps.size(), ps.type(), cls.size(),cls.type()) + + #get_segmentation_classes(datapath) \ No newline at end of file diff --git a/pointnet/num_seg_classes.txt b/pointnet/num_seg_classes.txt new file mode 100644 index 000000000..bf3dd6358 --- /dev/null +++ b/pointnet/num_seg_classes.txt @@ -0,0 +1,16 @@ +Airplane 4 +Bag 2 +Cap 2 +Car 4 +Chair 4 +Earphone 3 +Guitar 3 +Knife 2 +Lamp 4 +Laptop 2 +Motorbike 6 +Mug 2 +Pistol 3 +Rocket 3 +Skateboard 3 +Table 3 diff --git a/utils/show_cls.py b/utils/show_cls.py index 5dced8b97..ed9348986 100644 --- a/utils/show_cls.py +++ b/utils/show_cls.py @@ -4,7 +4,7 @@ import torch.nn.parallel import torch.utils.data from torch.autograd import Variable -from pointnet.dataset import PartDataset +from pointnet.dataset import ShapeNetDataset from pointnet.model import PointNetCls import torch.nn.functional as F @@ -20,7 +20,7 @@ opt = parser.parse_args() print(opt) -test_dataset = PartDataset( +test_dataset = ShapeNetDataset( root='shapenetcore_partanno_segmentation_benchmark_v0', train=False, classification=True, diff --git a/utils/show_seg.py b/utils/show_seg.py index b6b01c624..89fd9c7c6 100644 --- a/utils/show_seg.py +++ b/utils/show_seg.py @@ -6,7 +6,7 @@ import torch.nn.parallel import torch.utils.data from torch.autograd import Variable -from pointnet.dataset import PartDataset +from pointnet.dataset import ShapeNetDataset from pointnet.model import PointNetDenseCls import matplotlib.pyplot as plt @@ -23,7 +23,7 @@ opt = parser.parse_args() print(opt) -d = PartDataset( +d = ShapeNetDataset( root=opt.dataset, class_choice=[opt.class_choice], train=False) diff --git a/utils/train_classification.py b/utils/train_classification.py index 07f758143..86ab5be4b 100644 --- a/utils/train_classification.py +++ b/utils/train_classification.py @@ -6,7 +6,7 @@ import torch.nn.parallel import torch.optim as optim import torch.utils.data -from pointnet.dataset import PartDataset +from pointnet.dataset import ShapeNetDataset from pointnet.model import PointNetCls import torch.nn.functional as F from tqdm import tqdm @@ -35,7 +35,7 @@ random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) -dataset = PartDataset( +dataset = ShapeNetDataset( root=opt.dataset, classification=True, npoints=opt.num_points) @@ -45,10 +45,10 @@ shuffle=True, num_workers=int(opt.workers)) -test_dataset = PartDataset( +test_dataset = ShapeNetDataset( root=opt.dataset, classification=True, - train=False, + split='test', npoints=opt.num_points) testdataloader = torch.utils.data.DataLoader( test_dataset, diff --git a/utils/train_segmentation.py b/utils/train_segmentation.py index b6be26576..7e79ea1c6 100644 --- a/utils/train_segmentation.py +++ b/utils/train_segmentation.py @@ -6,7 +6,7 @@ import torch.nn.parallel import torch.optim as optim import torch.utils.data -from pointnet.dataset import PartDataset +from pointnet.dataset import ShapeNetDataset from pointnet.model import PointNetDenseCls import torch.nn.functional as F from tqdm import tqdm @@ -34,7 +34,7 @@ random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) -dataset = PartDataset( +dataset = ShapeNetDataset( root=opt.dataset, classification=False, class_choice=[opt.class_choice]) @@ -44,11 +44,11 @@ shuffle=True, num_workers=int(opt.workers)) -test_dataset = PartDataset( +test_dataset = ShapeNetDataset( root=opt.dataset, classification=False, class_choice=[opt.class_choice], - train=False) + split='test') testdataloader = torch.utils.data.DataLoader( test_dataset, batch_size=opt.batchSize, @@ -125,13 +125,13 @@ target_np = target.cpu().data.numpy() - 1 for shape_idx in range(target_np.shape[0]): - parts = np.unique(target_np[shape_idx]) + parts = range(num_classes)#np.unique(target_np[shape_idx]) part_ious = [] for part in parts: I = np.sum(np.logical_and(pred_np[shape_idx] == part, target_np[shape_idx] == part)) U = np.sum(np.logical_or(pred_np[shape_idx] == part, target_np[shape_idx] == part)) if U == 0: - iou = 0 + iou = 1 #If the union of groundtruth and prediction points is empty, then count part IoU as 1 else: iou = I / float(U) part_ious.append(iou) From 5787367aea054e436e147ed2d1f2f3e6c9c80728 Mon Sep 17 00:00:00 2001 From: fxia22 Date: Mon, 4 Mar 2019 22:18:21 -0800 Subject: [PATCH 08/11] add modelnet script --- misc/modelnet_id.txt | 40 +++++++++++ {pointnet => misc}/num_seg_classes.txt | 0 pointnet/dataset.py | 99 ++++++++++++++++++++++---- utils/train_classification.py | 44 ++++++++---- 4 files changed, 154 insertions(+), 29 deletions(-) create mode 100644 misc/modelnet_id.txt rename {pointnet => misc}/num_seg_classes.txt (100%) diff --git a/misc/modelnet_id.txt b/misc/modelnet_id.txt new file mode 100644 index 000000000..f494da764 --- /dev/null +++ b/misc/modelnet_id.txt @@ -0,0 +1,40 @@ +airplane 0 +bathtub 1 +bed 2 +bench 3 +bookshelf 4 +bottle 5 +bowl 6 +car 7 +chair 8 +cone 9 +cup 10 +curtain 11 +desk 12 +door 13 +dresser 14 +flower_pot 15 +glass_box 16 +guitar 17 +keyboard 18 +lamp 19 +laptop 20 +mantel 21 +monitor 22 +night_stand 23 +person 24 +piano 25 +plant 26 +radio 27 +range_hood 28 +sink 29 +sofa 30 +stairs 31 +stool 32 +table 33 +tent 34 +toilet 35 +tv_stand 36 +vase 37 +wardrobe 38 +xbox 39 diff --git a/pointnet/num_seg_classes.txt b/misc/num_seg_classes.txt similarity index 100% rename from pointnet/num_seg_classes.txt rename to misc/num_seg_classes.txt diff --git a/pointnet/dataset.py b/pointnet/dataset.py index 7150201e8..678460a36 100644 --- a/pointnet/dataset.py +++ b/pointnet/dataset.py @@ -7,6 +7,7 @@ import sys from tqdm import tqdm import json +from plyfile import PlyData, PlyElement def get_segmentation_classes(root): catfile = os.path.join(root, 'synsetoffset2category.txt') @@ -27,7 +28,7 @@ def get_segmentation_classes(root): token = (os.path.splitext(os.path.basename(fn))[0]) meta[item].append((os.path.join(dir_point, token + '.pts'), os.path.join(dir_seg, token + '.seg'))) - with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'num_seg_classes.txt'), 'w') as f: + with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../misc/num_seg_classes.txt'), 'w') as f: for item in cat: datapath = [] num_seg_classes = 0 @@ -42,6 +43,16 @@ def get_segmentation_classes(root): print("category {} num segmentation classes {}".format(item, num_seg_classes)) f.write("{}\t{}\n".format(item, num_seg_classes)) +def gen_modelnet_id(root): + classes = [] + with open(os.path.join(root, 'train.txt'), 'r') as f: + for line in f: + classes.append(line.strip().split('/')[0]) + classes = np.unique(classes) + with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../misc/modelnet_id.txt'), 'w') as f: + for i in range(len(classes)): + f.write('{}\t{}\n'.format(classes[i], i)) + class ShapeNetDataset(data.Dataset): def __init__(self, root, @@ -88,7 +99,7 @@ def __init__(self, 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__)), 'num_seg_classes.txt'), 'r') as f: + 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]) @@ -129,18 +140,76 @@ def __getitem__(self, index): def __len__(self): return len(self.datapath) +class ModelNetDataset(data.Dataset): + def __init__(self, + root, + npoints=2500, + split='train', + data_augmentation=True): + self.npoints = npoints + self.root = root + self.split = split + self.data_augmentation = data_augmentation + self.fns = [] + with open(os.path.join(root, '{}.txt'.format(self.split)), 'r') as f: + for line in f: + self.fns.append(line.strip()) + + self.cat = {} + with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../misc/modelnet_id.txt'), 'r') as f: + for line in f: + ls = line.strip().split() + self.cat[ls[0]] = int(ls[1]) + + print(self.cat) + self.classes = list(self.cat.keys()) + + def __getitem__(self, index): + fn = self.fns[index] + cls = self.cat[fn.split('/')[0]] + with open(os.path.join(self.root, fn), 'rb') as f: + plydata = PlyData.read(f) + pts = np.vstack([plydata['vertex']['x'], plydata['vertex']['y'], plydata['vertex']['z']]).T + choice = np.random.choice(len(pts), self.npoints, replace=True) + point_set = pts[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 + + point_set = torch.from_numpy(point_set.astype(np.float32)) + cls = torch.from_numpy(np.array([cls]).astype(np.int64)) + return point_set, cls + + + def __len__(self): + return len(self.fns) if __name__ == '__main__': - datapath = sys.argv[1] - print('test') - d = ShapeNetDataset(root = datapath, class_choice = ['Chair']) - print(len(d)) - ps, seg = d[0] - print(ps.size(), ps.type(), seg.size(),seg.type()) - - d = ShapeNetDataset(root = datapath, classification = True) - print(len(d)) - ps, cls = d[0] - print(ps.size(), ps.type(), cls.size(),cls.type()) - - #get_segmentation_classes(datapath) \ No newline at end of file + dataset = sys.argv[1] + datapath = sys.argv[2] + + if dataset == 'shapenet': + d = ShapeNetDataset(root = datapath, class_choice = ['Chair']) + print(len(d)) + ps, seg = d[0] + print(ps.size(), ps.type(), seg.size(),seg.type()) + + d = ShapeNetDataset(root = datapath, classification = True) + print(len(d)) + ps, cls = d[0] + print(ps.size(), ps.type(), cls.size(),cls.type()) + # get_segmentation_classes(datapath) + + if dataset == 'modelnet': + gen_modelnet_id(datapath) + d = ModelNetDataset(root=datapath) + print(len(d)) + print(d[0]) + diff --git a/utils/train_classification.py b/utils/train_classification.py index 86ab5be4b..451874c83 100644 --- a/utils/train_classification.py +++ b/utils/train_classification.py @@ -6,7 +6,7 @@ import torch.nn.parallel import torch.optim as optim import torch.utils.data -from pointnet.dataset import ShapeNetDataset +from pointnet.dataset import ShapeNetDataset, ModelNetDataset from pointnet.model import PointNetCls import torch.nn.functional as F from tqdm import tqdm @@ -24,6 +24,7 @@ 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") opt = parser.parse_args() print(opt) @@ -35,26 +36,41 @@ random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) -dataset = ShapeNetDataset( - root=opt.dataset, - classification=True, - npoints=opt.num_points) +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) +elif opt.dataset_type == 'modelnet40': + dataset = ModelNetDataset( + root=opt.dataset, + npoints=opt.num_points) + + test_dataset = ModelNetDataset( + root=opt.dataset, + split='test', + npoints=opt.num_points) +else: + exit('wrong dataset type') + + dataloader = torch.utils.data.DataLoader( dataset, batch_size=opt.batchSize, shuffle=True, num_workers=int(opt.workers)) -test_dataset = ShapeNetDataset( - root=opt.dataset, - classification=True, - split='test', - npoints=opt.num_points) testdataloader = torch.utils.data.DataLoader( - test_dataset, - batch_size=opt.batchSize, - shuffle=True, - num_workers=int(opt.workers)) + test_dataset, + batch_size=opt.batchSize, + shuffle=True, + num_workers=int(opt.workers)) print(len(dataset), len(test_dataset)) num_classes = len(dataset.classes) From 44f0112aa1873c118c91bb5f48000cd97a3b0e3d Mon Sep 17 00:00:00 2001 From: fxia22 Date: Tue, 5 Mar 2019 11:28:32 -0800 Subject: [PATCH 09/11] results added --- README.md | 50 +++++++++++++++++++++++++++++++---- setup.py | 4 ++- utils/show_seg.py | 2 +- utils/train_classification.py | 9 ++++--- utils/train_segmentation.py | 3 ++- 5 files changed, 57 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 4c09bbc6e..755059706 100644 --- a/README.md +++ b/README.md @@ -1,24 +1,64 @@ # PointNet.pytorch -This repo is implementation for PointNet(https://arxiv.org/abs/1612.00593) in pytorch. The model is in `pointnet.py`. +This repo is implementation for PointNet(https://arxiv.org/abs/1612.00593) in pytorch. The model is in `pointnet/model.py`. +It is tested with pytorch-1.0. # Download data and running ``` +git clone https://github.com/fxia22/pointnet.pytorch +cd pointnet.pytorch +pip install -e . +``` + +Download and build visualization tool +``` +cd script bash build.sh #build C++ code for visualization bash download.sh #download dataset -python train_classification.py #train 3D model classification -python train_segmentation.py # train 3D model segmentaion +``` -python show_seg.py --model seg/seg_model_20.pth # show segmentation results +Training +``` +cd utils +python train_classification.py --dataset --nepoch= --dataset_type +python train_segmentation.py --dataset --nepoch= ``` # Performance -Without heavy tuning, PointNet can achieve 80-90% performance in classification and segmentaion on this [dataset](http://web.stanford.edu/~ericyi/project_page/part_annotation/index.html). Sample segmentation result: ![seg](https://raw.githubusercontent.com/fxia22/pointnet.pytorch/master/misc/show3d.png?token=AE638Oy51TL2HDCaeCF273X_-Bsy6-E2ks5Y_BUzwA%3D%3D) +## Classification performance + +On ModelNet40: + +| | Overall Acc | +| :---: | :---: | +| Original implementation | 89.2 | +| this implementation(w/o feature transform) | TBA | +| this implementation(w/ feature transform) | TBA | + +On [A subset of shapenet](http://web.stanford.edu/~ericyi/project_page/part_annotation/index.html) + +| | Overall Acc | +| :---: | :---: | +| Original implementation | N/A | +| this implementation(w/o feature transform) | 98.1 | +| this implementation(w/ feature transform) | TBA | + +## Segmentation performance + +Segmentation on [A subset of shapenet](http://web.stanford.edu/~ericyi/project_page/part_annotation/index.html). + +| Class(mIOU) | Airplane | Bag| Cap|Car|Chair|Earphone|Guitar|Knife|Lamp|Laptop|Motorbike|Mug|Pistol|Rocket|Skateboard|Table +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| Original implementation | 83.4 | 78.7 | 82.5| 74.9 |89.6| 73.0| 91.5| 85.9| 80.8| 95.3| 65.2| 93.0| 81.2| 57.9| 72.8| 80.6| +| this implementation(w/o feature transform) | 73.5 | 71.3 | 64.3 | 61.1 | 87.2 | 69.5 | 86.1|81.6| 77.4|92.7|41.3|86.5|78.2|41.2|61.0|81.1| +| this implementation(w/ feature transform) | TBA | + +Note that this implementation trains each class separately, so classes with fewer data will have slightly lower performance than reference implementation. # Links diff --git a/setup.py b/setup.py index 967cf9550..9e71bd7b2 100644 --- a/setup.py +++ b/setup.py @@ -5,5 +5,7 @@ setup(name='pointnet', packages=['pointnet'], package_dir={'pointnet': 'pointnet'}, - install_requires=['torch'], + install_requires=['torch', + 'tqdm', + 'plyfile'], version='0.0.1') diff --git a/utils/show_seg.py b/utils/show_seg.py index 89fd9c7c6..3926876c7 100644 --- a/utils/show_seg.py +++ b/utils/show_seg.py @@ -40,7 +40,7 @@ gt = cmap[seg.numpy() - 1, :] state_dict = torch.load(opt.model) -classifier = PointNetDenseCls(k= state_dict['conv4.weight'].size()[0] ) +classifier = PointNetDenseCls(k= state_dict['conv4.weight'].size()[0]) classifier.load_state_dict(state_dict) classifier.eval() diff --git a/utils/train_classification.py b/utils/train_classification.py index 451874c83..5427cb579 100644 --- a/utils/train_classification.py +++ b/utils/train_classification.py @@ -46,16 +46,19 @@ root=opt.dataset, classification=True, split='test', - npoints=opt.num_points) + npoints=opt.num_points, + data_augmentation=False) elif opt.dataset_type == 'modelnet40': dataset = ModelNetDataset( root=opt.dataset, - npoints=opt.num_points) + npoints=opt.num_points, + split='trainval') test_dataset = ModelNetDataset( root=opt.dataset, split='test', - npoints=opt.num_points) + npoints=opt.num_points, + data_augmentation=False) else: exit('wrong dataset type') diff --git a/utils/train_segmentation.py b/utils/train_segmentation.py index 7e79ea1c6..0ae5ca263 100644 --- a/utils/train_segmentation.py +++ b/utils/train_segmentation.py @@ -48,7 +48,8 @@ root=opt.dataset, classification=False, class_choice=[opt.class_choice], - split='test') + split='test', + data_augmentation=False) testdataloader = torch.utils.data.DataLoader( test_dataset, batch_size=opt.batchSize, From bf91ed3a6bd2f133ea7151c43bddbe7af972b527 Mon Sep 17 00:00:00 2001 From: fxia22 Date: Tue, 5 Mar 2019 22:21:35 -0800 Subject: [PATCH 10/11] add feature transform --- README.md | 10 ++-- pointnet/model.py | 100 ++++++++++++++++++++++++++++------ utils/show_cls.py | 7 ++- utils/show_seg.py | 5 +- utils/train_classification.py | 13 +++-- utils/train_segmentation.py | 14 +++-- 6 files changed, 112 insertions(+), 37 deletions(-) diff --git a/README.md b/README.md index 755059706..db58a919b 100644 --- a/README.md +++ b/README.md @@ -25,10 +25,9 @@ python train_classification.py --dataset --nepoch= python train_segmentation.py --dataset --nepoch= ``` -# Performance +Use `--feature_transform` to use feature transform. -Sample segmentation result: -![seg](https://raw.githubusercontent.com/fxia22/pointnet.pytorch/master/misc/show3d.png?token=AE638Oy51TL2HDCaeCF273X_-Bsy6-E2ks5Y_BUzwA%3D%3D) +# Performance ## Classification performance @@ -37,7 +36,7 @@ On ModelNet40: | | Overall Acc | | :---: | :---: | | Original implementation | 89.2 | -| this implementation(w/o feature transform) | TBA | +| this implementation(w/o feature transform) | 86.4 | | this implementation(w/ feature transform) | TBA | On [A subset of shapenet](http://web.stanford.edu/~ericyi/project_page/part_annotation/index.html) @@ -60,6 +59,9 @@ Segmentation on [A subset of shapenet](http://web.stanford.edu/~ericyi/project_ Note that this implementation trains each class separately, so classes with fewer data will have slightly lower performance than reference implementation. +Sample segmentation result: +![seg](https://raw.githubusercontent.com/fxia22/pointnet.pytorch/master/misc/show3d.png?token=AE638Oy51TL2HDCaeCF273X_-Bsy6-E2ks5Y_BUzwA%3D%3D) + # Links - [Project Page](http://stanford.edu/~rqi/pointnet/) diff --git a/pointnet/model.py b/pointnet/model.py index 7be4b7cfa..9e5838e96 100644 --- a/pointnet/model.py +++ b/pointnet/model.py @@ -46,8 +46,46 @@ def forward(self, x): 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): + def __init__(self, global_feat = True, feature_transform = False): super(PointNetfeat, self).__init__() self.stn = STN3d() self.conv1 = torch.nn.Conv1d(3, 64, 1) @@ -57,7 +95,9 @@ def __init__(self, global_feat = True): 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] @@ -66,21 +106,31 @@ def forward(self, x): 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 + return x, trans, trans_feat else: x = x.view(-1, 1024, 1).repeat(1, 1, n_pts) - return torch.cat([x, pointfeat], 1), trans + return torch.cat([x, pointfeat], 1), trans, trans_feat class PointNetCls(nn.Module): - def __init__(self, k = 2): + def __init__(self, k=2, feature_transform=False): super(PointNetCls, self).__init__() - self.feat = PointNetfeat(global_feat=True) + 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) @@ -90,17 +140,18 @@ def __init__(self, k = 2): self.relu = nn.ReLU() def forward(self, x): - x, trans = self.feat(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 + return F.log_softmax(x, dim=1), trans, trans_feat class PointNetDenseCls(nn.Module): - def __init__(self, k = 2): + def __init__(self, k = 2, feature_transform=False): super(PointNetDenseCls, self).__init__() self.k = k - self.feat = PointNetfeat(global_feat=False) + 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) @@ -112,7 +163,7 @@ def __init__(self, k = 2): def forward(self, x): batchsize = x.size()[0] n_pts = x.size()[2] - x, trans = self.feat(x) + 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))) @@ -120,27 +171,42 @@ def forward(self, 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 + return x, trans, trans_feat +def feature_transform_reguliarzer(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_reguliarzer(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_reguliarzer(out)) + pointfeat = PointNetfeat(global_feat=True) - out, _ = pointfeat(sim_data) + out, _, _ = pointfeat(sim_data) print('global feat', out.size()) pointfeat = PointNetfeat(global_feat=False) - out, _ = pointfeat(sim_data) + out, _, _ = pointfeat(sim_data) print('point feat', out.size()) cls = PointNetCls(k = 5) - out, _ = cls(sim_data) + out, _, _ = cls(sim_data) print('class', out.size()) seg = PointNetDenseCls(k = 3) - out, _ = seg(sim_data) + out, _, _ = seg(sim_data) print('seg', out.size()) diff --git a/utils/show_cls.py b/utils/show_cls.py index ed9348986..d52ef0878 100644 --- a/utils/show_cls.py +++ b/utils/show_cls.py @@ -22,9 +22,10 @@ test_dataset = ShapeNetDataset( root='shapenetcore_partanno_segmentation_benchmark_v0', - train=False, + split='test', classification=True, - npoints=opt.num_points) + npoints=opt.num_points, + data_augmentation=False) testdataloader = torch.utils.data.DataLoader( test_dataset, batch_size=32, shuffle=True) @@ -40,7 +41,7 @@ points, target = Variable(points), Variable(target[:, 0]) points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() - pred, _ = classifier(points) + pred, _, _ = classifier(points) loss = F.nll_loss(pred, target) pred_choice = pred.data.max(1)[1] diff --git a/utils/show_seg.py b/utils/show_seg.py index 3926876c7..a849751a3 100644 --- a/utils/show_seg.py +++ b/utils/show_seg.py @@ -26,7 +26,8 @@ d = ShapeNetDataset( root=opt.dataset, class_choice=[opt.class_choice], - train=False) + split='test', + data_augmentation=False) idx = opt.idx @@ -47,7 +48,7 @@ point = point.transpose(1, 0).contiguous() point = Variable(point.view(1, point.size()[0], point.size()[1])) -pred, _ = classifier(point) +pred, _, _ = classifier(point) pred_choice = pred.data.max(2)[1] print(pred_choice) diff --git a/utils/train_classification.py b/utils/train_classification.py index 5427cb579..f85edbe19 100644 --- a/utils/train_classification.py +++ b/utils/train_classification.py @@ -7,7 +7,7 @@ import torch.optim as optim import torch.utils.data from pointnet.dataset import ShapeNetDataset, ModelNetDataset -from pointnet.model import PointNetCls +from pointnet.model import PointNetCls, feature_transform_reguliarzer import torch.nn.functional as F from tqdm import tqdm @@ -25,6 +25,7 @@ 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) @@ -84,7 +85,7 @@ except OSError: pass -classifier = PointNetCls(k=num_classes) +classifier = PointNetCls(k=num_classes, feature_transform=opt.feature_transform) if opt.model != '': classifier.load_state_dict(torch.load(opt.model)) @@ -105,8 +106,10 @@ points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() - pred, _ = classifier(points) + pred, trans, trans_feat = classifier(points) loss = F.nll_loss(pred, target) + if opt.feature_transform: + loss += feature_transform_reguliarzer(trans_feat) * 0.001 loss.backward() optimizer.step() pred_choice = pred.data.max(1)[1] @@ -120,7 +123,7 @@ points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() classifier = classifier.eval() - pred, _ = classifier(points) + pred, _, _ = classifier(points) loss = F.nll_loss(pred, target) pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.data).cpu().sum() @@ -136,7 +139,7 @@ points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() classifier = classifier.eval() - pred, _ = classifier(points) + pred, _, _ = classifier(points) pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.data).cpu().sum() total_correct += correct.item() diff --git a/utils/train_segmentation.py b/utils/train_segmentation.py index 0ae5ca263..237c298c6 100644 --- a/utils/train_segmentation.py +++ b/utils/train_segmentation.py @@ -7,7 +7,7 @@ import torch.optim as optim import torch.utils.data from pointnet.dataset import ShapeNetDataset -from pointnet.model import PointNetDenseCls +from pointnet.model import PointNetDenseCls, feature_transform_reguliarzer import torch.nn.functional as F from tqdm import tqdm import numpy as np @@ -24,7 +24,7 @@ parser.add_argument('--model', type=str, default='', help='model path') parser.add_argument('--dataset', type=str, required=True, help="dataset path") parser.add_argument('--class_choice', type=str, default='Chair', help="class_choice") - +parser.add_argument('--feature_transform', action='store_true', help="use feature transform") opt = parser.parse_args() print(opt) @@ -66,7 +66,7 @@ blue = lambda x: '\033[94m' + x + '\033[0m' -classifier = PointNetDenseCls(k=num_classes) +classifier = PointNetDenseCls(k=num_classes, feature_transform=opt.feature_transform) if opt.model != '': classifier.load_state_dict(torch.load(opt.model)) @@ -85,11 +85,13 @@ points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() - pred, _ = classifier(points) + pred, trans, trans_feat = classifier(points) pred = pred.view(-1, num_classes) target = target.view(-1, 1)[:, 0] - 1 #print(pred.size(), target.size()) loss = F.nll_loss(pred, target) + if opt.feature_transform: + loss += feature_transform_reguliarzer(trans_feat) * 0.001 loss.backward() optimizer.step() pred_choice = pred.data.max(1)[1] @@ -102,7 +104,7 @@ points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() classifier = classifier.eval() - pred, _ = classifier(points) + pred, _, _ = classifier(points) pred = pred.view(-1, num_classes) target = target.view(-1, 1)[:, 0] - 1 loss = F.nll_loss(pred, target) @@ -119,7 +121,7 @@ points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() classifier = classifier.eval() - pred, _ = classifier(points) + pred, _, _ = classifier(points) pred_choice = pred.data.max(2)[1] pred_np = pred_choice.cpu().data.numpy() From fab247e8875dbcc909995df20330a525033e7c05 Mon Sep 17 00:00:00 2001 From: fxia22 Date: Wed, 6 Mar 2019 10:15:29 -0800 Subject: [PATCH 11/11] add feature transform results --- README.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index db58a919b..fafd9d25a 100644 --- a/README.md +++ b/README.md @@ -37,7 +37,7 @@ On ModelNet40: | :---: | :---: | | Original implementation | 89.2 | | this implementation(w/o feature transform) | 86.4 | -| this implementation(w/ feature transform) | TBA | +| this implementation(w/ feature transform) | 86.8 | On [A subset of shapenet](http://web.stanford.edu/~ericyi/project_page/part_annotation/index.html) @@ -45,7 +45,7 @@ On [A subset of shapenet](http://web.stanford.edu/~ericyi/project_page/part_anno | :---: | :---: | | Original implementation | N/A | | this implementation(w/o feature transform) | 98.1 | -| this implementation(w/ feature transform) | TBA | +| this implementation(w/ feature transform) | 96.8 | ## Segmentation performance @@ -55,7 +55,6 @@ Segmentation on [A subset of shapenet](http://web.stanford.edu/~ericyi/project_ | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | | Original implementation | 83.4 | 78.7 | 82.5| 74.9 |89.6| 73.0| 91.5| 85.9| 80.8| 95.3| 65.2| 93.0| 81.2| 57.9| 72.8| 80.6| | this implementation(w/o feature transform) | 73.5 | 71.3 | 64.3 | 61.1 | 87.2 | 69.5 | 86.1|81.6| 77.4|92.7|41.3|86.5|78.2|41.2|61.0|81.1| -| this implementation(w/ feature transform) | TBA | Note that this implementation trains each class separately, so classes with fewer data will have slightly lower performance than reference implementation.