├── .gitignore
├── LICENSE
├── README.md
├── build.sh
├── datasets.py
├── pointnet.py
├── render_balls_so.cpp
├── requirements.txt
├── show3d_balls.py
├── show_seg.py
├── show_seg_s3d.py
├── train_cls.py
├── train_seg.py
└── utils.py


/.gitignore:
--------------------------------------------------------------------------------
1 | ./data/
2 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | The MIT License (MIT) Copyright (c) 2020 Yunxiao Shi
2 | 
3 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
4 | 
5 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
6 | 
7 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
8 | © 2020 GitHub, Inc.
9 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | ## PointNet-PyTorch
 2 | 
 3 | [![Python 3.6+](https://img.shields.io/badge/Python-3.6%2B-blue)](https://www.python.org/)
 4 | [![MIT License](https://img.shields.io/badge/MIT-License-brightgreen)](./LICENSE)
 5 | 
 6 | This is a PyTorch implementation of [PointNet (CVPR 2017)](https://arxiv.org/abs/1612.00593 "PointNet"), with comprehensive experiments.
 7 | 
 8 | ## Installation
 9 | 
10 | It is recommended to use [conda](https://docs.conda.io/en/latest/) to manage your env. For example do  
11 | ```
12 | conda create -n pointnet python=3.6
13 | conda activate pointnet
14 | pip install -r requirements.txt
15 | ```
16 | 
17 | You may also need to install [PyMesh](https://github.com/PyMesh/PyMesh "PyMesh"). See [here](https://github.com/PyMesh/PyMesh#Build) for instructions to install.
18 | 
19 | ## Usage
20 | 
21 | This code implements object classification on ModelNet, shape part segmentation on ShapeNet and indoor scene semantic segmentation on the Stanford 3D dataset.
22 | 
23 | For the missing ```s3d_cat2num.txt``` when training on S3DIS, follow [#3](https://github.com/kentsyx/pointnet-pytorch/issues/3#issuecomment-643061963) to generate it once you have the dataset downloaded.
24 | 
25 | ### ModelNet Classification
26 | 
27 | Download the ModelNet10 dataset from [here](http://3dvision.princeton.edu/projects/2014/3DShapeNets/ModelNet10.zip) or the ModelNet40 dataset from [here](https://lmb.informatik.uni-freiburg.de/resources/datasets/ORION/modelnet40_manually_aligned.tar). Unzip and run 
28 | ```
29 | python train_cls.py -dset modelnet40 -r modelnet_root_dir -np number_of_points_to_sample
30 | ```
31 | 
32 | ### ShapeNet Part Segmentation
33 | 
34 | Download the ShapeNet dataset from [here](https://shapenet.cs.stanford.edu/ericyi/shapenetcore_partanno_segmentation_benchmark_v0.zip). Unzip and run
35 | ```
36 | python train_seg.py -dset shapenet16 -r shapenet_root_dir -np number_of_points_to_sample
37 | ```
38 | 
39 | ### Indoor Scene Semantic Segmentation
40 | 
41 | Download the S3DIS dataset from [here](http://buildingparser.stanford.edu/dataset.html#Download) (you need to submit a request). Unzip and do
42 | ```
43 | cd Stanford3dDataset_v1.2
44 | mkdir train test
45 | mv Area_1 Area_2 Area_3 Area_4 Area_6 train
46 | mv Area_5 test
47 | ```
48 | to create train/test split. Then set ```gen_labels=True``` in the class ```S3dDataset``` in datasets.py and do
49 | ```
50 | python datasets.py
51 | ``` 
52 | to generate labels for the train and test set respectively. __After that always set ```gen_labels=False```__. With labels generated do
53 | ```
54 | python train_seg.py -dset s3dis -r s3dis_root_dir -np number_of_points_to_sample
55 | ```
56 | to start training.
57 | 
58 | ## Visualization
59 | 
60 | First do ```sh build.sh```, then use ```show_seg.py``` to visualize segmented object parts. Below are some example results.
61 | 
62 | <p align="center">
63 | <img src="https://i.ibb.co/rx5KB2x/part.png")
64 | </p>
65 | 
66 | For S3DIS, you have to combine scene components along with their labels into one text file (```cat``` and ```paste``` seems to be an easy way to do this) and then pass it to ```show_seg_s3dis.py```. Below are some example results (removed some clutter classes for better visualization).
67 | 
68 | <p align="center">
69 | <img src="https://i.ibb.co/0Gcy2KG/s3dis.png")
70 | </p>
71 | 
72 | ## Results
73 | 
74 | Certain design choices in the original paper are not implemented here for simplicity. There is some performance gap on ModelNet classification, for ShapeNet and S3DIS seems to be on par with the original paper.
75 | 
76 | <center>
77 | 
78 | | | accuracy | class avg IoU
79 | | :------: | :------: | :------: |
80 | | ModelNet10 | 87.2% | - | 
81 | | ModelNet40 | 85.4% | - | 
82 | | ShapeNet | - | 82.9% |
83 | | S3DIS | 72.1% | 50.6% |
84 | 
85 | </center>
86 | 
87 | ## Acknowledgements
88 | 
89 | [pointnet.pytorch](https://github.com/fxia22/pointnet.pytorch) (many thanks)
90 | 
91 | [original tensorflow implementation](https://github.com/charlesq34/pointnet)
92 | 
93 | ## LICENSE
94 | 
95 | MIT
96 | 


--------------------------------------------------------------------------------
/build.sh:
--------------------------------------------------------------------------------
1 | g++ -std=c++11 render_balls_so.cpp -o render_balls_so.so -shared -fPIC -O2 -D_GLIBCXX_USE_CXX11_ABI=0
2 | 


--------------------------------------------------------------------------------
/datasets.py:
--------------------------------------------------------------------------------
  1 | # -*- coding: utf-8 -*-
  2 | 
  3 | import os
  4 | import random
  5 | import numpy as np
  6 | import torch
  7 | import torch.utils.data as data
  8 | import pymesh
  9 | from tqdm import tqdm
 10 | from utils import shapenet_labels
 11 | 
 12 | def scale_linear_bycolumn(rawdata, high=1.0, low=0.0):
 13 |     mins = np.min(rawdata, axis=0)
 14 |     maxs = np.max(rawdata, axis=0)
 15 |     rng = maxs - mins
 16 |     return high - (high-low)*(maxs-rawdata)/(rng+np.finfo(np.float32).eps)
 17 | 
 18 | 
 19 | class ClsDataset(data.Dataset):
 20 |         '''Object classification on ModelNet'''
 21 |         def __init__(self, root, npoints=1024, train=True):
 22 |                 self.root = root
 23 |                 self.npoints = npoints
 24 |                 self.catfile = os.path.join(self.root, 'modelnet_cat2num.txt')
 25 |                 self.cat = {}
 26 | 
 27 |                 with open(self.catfile, 'r') as f:
 28 |                         for line in f.readlines():
 29 |                                 lns = line.strip().split()
 30 |                                 self.cat[lns[0]] = lns[1]
 31 |                 self.num_classes = len(self.cat)
 32 |                 self.datapath = []
 33 |                 FLAG = 'train' if train else 'test'
 34 |                 for item in os.listdir(self.root):
 35 |                         if os.path.isdir(os.path.join(self.root, item)):
 36 |                                 for f in os.listdir(os.path.join(self.root, item, FLAG)):
 37 |                                     if f.endswith('.off'):
 38 |                                         self.datapath.append((os.path.join(self.root, item, FLAG, f), int(self.cat[item])))
 39 | 
 40 | 
 41 |         def __getitem__(self, idx):
 42 |                 fn = self.datapath[idx]
 43 |                 points = pymesh.load_mesh(fn[0]).vertices
 44 |                 label = fn[1]
 45 |                 replace = True if points.shape[0]<self.npoints else False
 46 |                 choice = np.random.choice(points.shape[0], self.npoints, replace=replace)
 47 |                 points = points[choice, :]
 48 |                 points = scale_linear_bycolumn(points)
 49 |                 points = torch.from_numpy(points.astype(np.float32))
 50 |                 label = torch.from_numpy(np.array([label]).astype(np.int64))
 51 |                 return points, label
 52 | 
 53 | 
 54 |         def __len__(self):
 55 |                 return len(self.datapath)
 56 | 
 57 | 
 58 | class PartDataset(data.Dataset):
 59 |         def __init__(self, root, npoints=2048, class_choice=None, train=True):
 60 |                 '''Part segmentation on ShapeNet'''
 61 |                 self.root = root
 62 |                 self.npoints = npoints
 63 |                 self.catfile = os.path.join(self.root, 'synsetoffset2category.txt')
 64 |                 self.cat = {}
 65 | 
 66 |                 with open(self.catfile, 'r') as f:
 67 |                         for line in f.readlines():
 68 |                                 lns = line.strip().split()
 69 |                                 self.cat[lns[0]] = lns[1]
 70 |                 
 71 |                 if not class_choice is None:
 72 |                     self.cat = {k:v for k, v in self.cat.items() if k in class_choice}
 73 |                     self.num_classes = shapenet_labels[class_choice[0]]
 74 |                 else:
 75 |                     self.num_classes = 50
 76 | 
 77 |                 self.meta = {}
 78 |                 for item in self.cat:
 79 |                         self.meta[item] = []
 80 |                         point_dir = os.path.join(self.root, self.cat[item], 'points')
 81 |                         seg_dir = os.path.join(self.root, self.cat[item], 'points_label')
 82 |                 fns = sorted(os.listdir(point_dir))
 83 |                 if train:
 84 |                         fns = fns[:int(0.9*len(fns))]
 85 |                 else:
 86 |                         fns = fns[int(0.9*len(fns)):]
 87 | 
 88 |                 for fn in fns:
 89 |                         token = (os.path.splitext(os.path.basename(fn))[0])
 90 |                         self.meta[item].append((os.path.join(point_dir, token + '.pts'), os.path.join(seg_dir, token + '.seg')))
 91 | 
 92 |                 self.datapath = []
 93 |                 for item in self.cat:
 94 |                         for fn in self.meta[item]:
 95 |                                 self.datapath.append((item, fn[0], fn[1]))
 96 | 
 97 |                 self.classes = dict(zip(sorted(self.cat), range(len(self.cat))))
 98 | 
 99 | 
100 |         def __getitem__(self, idx):
101 |                 fn = self.datapath[idx]
102 |                 points = np.loadtxt(fn[1]).astype(np.float32)
103 |                 seg = np.loadtxt(fn[2]).astype(np.int64)
104 |                 replace = True if points.shape[0]<self.npoints else False
105 |                 choice = np.random.choice(len(seg), self.npoints, replace=replace)
106 |                 # resample
107 |                 points = points[choice, :]
108 |                 seg = seg[choice]
109 |                 points = torch.from_numpy(points)
110 |                 seg = torch.from_numpy(seg)
111 |                 return points, seg
112 | 
113 | 
114 |         def __len__(self):
115 |                 return len(self.datapath)
116 | 
117 | 
118 | class S3dDataset(data.Dataset):
119 |         '''Semantic segmentation on S3DIS'''
120 |         def __init__(self, root, npoints=4096, train=True, gen_labels=False):
121 |                 self.root = root
122 |                 self.npoints = npoints
123 |                 self.catfile = os.path.join(self.root, 's3d_cat2num.txt')
124 |                 self.cat = {}
125 |                 with open(self.catfile, 'r') as f:
126 |                     for line in f.readlines():
127 |                         lns = line.strip().split()
128 |                         self.cat[lns[0]] = lns[1]
129 |                 self.num_classes = len(self.cat)
130 |                 self.datapath, self.labelspath = [], []
131 |                 FLAG = 'train' if train else 'test'
132 |                 path = os.path.join(self.root, FLAG)
133 |                 for area in os.listdir(path):
134 |                     area_path = os.path.join(path, area)
135 |                     for scene in os.listdir(area_path):
136 |                         if os.path.isdir(os.path.join(area_path, scene)):
137 |                             scene_path = os.path.join(area_path, scene)
138 |                             for scene_component in os.listdir(os.path.join(scene_path, 'Annotations')):
139 |                                 if not scene_component.endswith('_labels.txt'):
140 |                                     self.datapath.append(os.path.join(scene_path, 'Annotations', scene_component))
141 |                 
142 | 
143 |                 if gen_labels: # do this only once
144 |                     pbar = tqdm(total=len(self.datapath))
145 |                     for path in self.datapath:
146 |                         l = path.split('/')
147 |                         labels_path = os.path.join(l[0], l[1], l[2], l[3], l[4])
148 |                         component_name = l[-1].split('.')[0]
149 |                         class_name = l[-1].split('_')[0]
150 |                         with open(path, 'r') as f:
151 |                             for line in f.readlines():
152 |                                 with open(os.path.join(labels_path, component_name + '_labels.txt'), 'a') as g:
153 |                                     g.write(str(self.cat[class_name]) + '\n')
154 |                         pbar.update()
155 |                     
156 | 
157 |         def __getitem__(self, idx):
158 |             fn = self.datapath[idx]
159 |             points = np.loadtxt(fn)[:, :3].astype(np.float32)
160 |             ln = os.path.splitext(fn)[0] + '_labels.txt'
161 |             seg = np.loadtxt(ln).astype(np.int64)
162 |             replace = True if points.shape[0]<self.npoints else False
163 |             choice = np.random.choice(points.shape[0], self.npoints, replace=replace)
164 |             points = points[choice, :]
165 |             points = scale_linear_bycolumn(points)
166 |             seg = seg[choice]
167 |             points = torch.from_numpy(points)
168 |             seg = torch.from_numpy(seg)
169 |             return points, seg
170 | 
171 | 
172 |         def __len__(self):
173 |             return len(self.datapath)
174 | 
175 | 
176 | if __name__ == '__main__':
177 | 
178 |         c = ClsDataset(root='modelnet40_manually_aligned')
179 |         print(ps.type(), ps.size(), l.type(), l.size(), l)
180 | 
181 |         d = PartDataset(root='shapenetcore_partanno_segmentation_benchmark_v0')
182 |         ps, seg = d[10]
183 |         print(ps.type(), ps.size(), seg.type(), seg.size())
184 | 
185 |         s = S3dDataset(root='Stanford3dDataset_v1.2', train=False, gen_labels=False)
186 |         ps, seg = s[100]
187 |         print(ps.type(), ps.size(), seg.type(), seg.size())
188 | 
189 | 
190 | 
191 | 


--------------------------------------------------------------------------------
/pointnet.py:
--------------------------------------------------------------------------------
  1 | # -*- coding: utf-8 -*-
  2 | 
  3 | import numpy as np
  4 | 
  5 | import torch
  6 | import torch.autograd as autograd
  7 | import torch.nn as nn
  8 | import torch.nn.functional as F
  9 | 
 10 | class STN3D(nn.Module):
 11 | 	def __init__(self):
 12 | 		super(STN3D, self).__init__()
 13 | 		self.conv1 = nn.Conv1d(3, 64, 1)
 14 | 		self.conv2 = nn.Conv1d(64, 128, 1)
 15 | 		self.conv3 = nn.Conv1d(128, 1024, 1)
 16 | 		self.fc1 = nn.Linear(1024, 512)
 17 | 		self.fc2 = nn.Linear(512, 256)
 18 | 		self.fc3 = nn.Linear(256, 9)
 19 | 		self.bn1 = nn.BatchNorm1d(64)
 20 | 		self.bn2 = nn.BatchNorm1d(128)
 21 | 		self.bn3 = nn.BatchNorm1d(1024)
 22 | 		self.bn4 = nn.BatchNorm1d(512)
 23 | 		self.bn5 = nn.BatchNorm1d(256)
 24 | 
 25 | 
 26 | 	def forward(self, x):
 27 | 		batchsize = x.size()[0]
 28 | 		x = F.relu(self.bn1(self.conv1(x)))
 29 | 		x = F.relu(self.bn2(self.conv2(x)))
 30 | 		x = F.relu(self.bn3(self.conv3(x)))
 31 | 		x = torch.max(x, 2, keepdim=True)[0]
 32 | 		x = x.view(-1, 1024)
 33 | 
 34 | 		x = F.relu(self.bn4(self.fc1(x)))
 35 | 		x = F.relu(self.bn5(self.fc2(x)))
 36 | 		x = self.fc3(x)
 37 | 
 38 | 		iden = autograd.Variable(torch.from_numpy(np.array([1,0,0,0,1,0,0,0,1]).astype(np.float32))).view(1,9).repeat(batchsize, 1)
 39 | 		if x.is_cuda:
 40 | 			iden = iden.cuda()
 41 | 		x += iden
 42 | 		x = x.view(-1, 3, 3)
 43 | 		return x
 44 | 
 45 | 
 46 | class PointNetfeat(nn.Module):
 47 | 	def __init__(self, global_feat=True):
 48 | 		super(PointNetfeat, self).__init__()
 49 | 		self.stn = STN3D()
 50 | 		self.conv1 = nn.Conv1d(3, 64, 1)
 51 | 		self.conv2 = nn.Conv1d(64, 128, 1)
 52 | 		self.conv3 = nn.Conv1d(128, 1024, 1)
 53 | 		self.bn1 = nn.BatchNorm1d(64)
 54 | 		self.bn2 = nn.BatchNorm1d(128)
 55 | 		self.bn3 = nn.BatchNorm1d(1024)
 56 | 		self.global_feat = global_feat
 57 | 
 58 | 
 59 | 	def forward(self, x):
 60 | 		batchsize = x.size()[0]
 61 | 		n_pts = x.size()[2]
 62 | 		trans = self.stn(x)
 63 | 		x = x.transpose(2, 1)
 64 | 		x = torch.bmm(x, trans)
 65 | 		x = x.transpose(2, 1)
 66 | 		x = F.relu(self.bn1(self.conv1(x)))
 67 | 		pointfeat = x
 68 | 		x = F.relu(self.bn2(self.conv2(x)))
 69 | 		x = self.bn3(self.conv3(x))
 70 | 		x = torch.max(x, 2, keepdim=True)[0]
 71 | 		x = x.view(-1, 1024)
 72 | 		if self.global_feat:
 73 | 			return x, trans
 74 | 		else:
 75 | 			x = x.view(-1, 1024, 1).repeat(1, 1, n_pts)
 76 | 			return torch.cat([x, pointfeat], 1), trans
 77 | 
 78 | 
 79 | class PointNetCls(nn.Module):
 80 | 	def __init__(self, k=2):
 81 | 		super(PointNetCls, self).__init__()
 82 | 		self.k = k
 83 | 		self.feat = PointNetfeat(global_feat=True)
 84 | 		self.fc1 = nn.Linear(1024, 512)
 85 | 		self.fc2 = nn.Linear(512, 256)
 86 | 		self.fc3 = nn.Linear(256, k)
 87 | 		self.bn1 = nn.BatchNorm1d(512)
 88 | 		self.bn2 = nn.BatchNorm1d(256)
 89 | 
 90 | 
 91 | 	def forward(self, x):
 92 | 		x, trans = self.feat(x)
 93 | 		x = F.relu(self.bn1(self.fc1(x)))
 94 | 		x = F.relu(self.bn2(self.fc2(x)))
 95 | 		x = self.fc3(x)
 96 | 		return F.log_softmax(x, dim=1), trans
 97 | 
 98 | 
 99 | class PointNetSeg(nn.Module):
100 | 	def __init__(self, k=2):
101 | 		super(PointNetSeg, self).__init__()
102 | 		self.k = k
103 | 		self.feat = PointNetfeat(global_feat=False)
104 | 		self.conv1 = nn.Conv1d(1088, 512, 1)
105 | 		self.conv2 = nn.Conv1d(512, 256, 1)
106 | 		self.conv3 = nn.Conv1d(256, 128, 1)
107 | 		self.conv4 = nn.Conv1d(128, self.k, 1)
108 | 		self.bn1 = nn.BatchNorm1d(512)
109 | 		self.bn2 = nn.BatchNorm1d(256)
110 | 		self.bn3 = nn.BatchNorm1d(128)
111 | 
112 | 	
113 | 	def forward(self, x):
114 | 		batchsize = x.size()[0]
115 | 		n_pts = x.size()[2]
116 | 		x, trans = self.feat(x)
117 | 		x = F.relu(self.bn1(self.conv1(x)))
118 | 		x = F.relu(self.bn2(self.conv2(x)))
119 | 		x = F.relu(self.bn3(self.conv3(x)))
120 | 		x = self.conv4(x)
121 | 		x = x.transpose(2, 1).contiguous()
122 | 		x = F.log_softmax(x.view(-1, self.k), dim=-1)
123 | 		x = x.view(batchsize, n_pts, self.k)
124 | 		return x, trans
125 | 
126 | 
127 | if __name__ == '__main__':
128 | 
129 | 	sim_data = autograd.Variable(torch.randn(32, 3, 2048))
130 | 	trans = STN3D()
131 | 	out = trans(sim_data)
132 | 	print('stn', out.size())
133 | 
134 | 	pointfeat = PointNetfeat(global_feat=True)
135 | 	out, _ = pointfeat(sim_data)
136 | 	print('global feat', out.size())
137 | 
138 | 	pointfeat = PointNetfeat(global_feat=False)
139 | 	out, _ = pointfeat(sim_data)
140 | 	print('point feat', out.size())
141 | 
142 | 	cls = PointNetCls(k=4)
143 | 	out, _ = cls(sim_data)
144 | 	print('class', out.size())
145 | 
146 | 	seg = PointNetSeg(k=4)
147 | 	out, _ = seg(sim_data)
148 | 	print('seg', out.size())
149 | 
150 | 
151 | 
152 | 
153 | 
154 | 
155 | 
156 | 
157 | 
158 | 
159 | 
160 | 
161 | 
162 | 
163 | 
164 | 
165 | 
166 | 
167 | 
168 | 
169 | 
170 | 
171 | 
172 | 
173 | 
174 | 
175 | 
176 | 
177 | 
178 | 
179 | 
180 | 
181 | 
182 | 
183 | 
184 | 
185 | 
186 | 
187 | 
188 | 
189 | 
190 | 
191 | 
192 | 
193 | 
194 | 
195 | 
196 | 
197 | 
198 | 
199 | 
200 | 
201 | 
202 | 
203 | 
204 | 
205 | 
206 | 


--------------------------------------------------------------------------------
/render_balls_so.cpp:
--------------------------------------------------------------------------------
 1 | #include <cstdio>
 2 | #include <vector>
 3 | #include <algorithm>
 4 | #include <math.h>
 5 | using namespace std;
 6 | 
 7 | struct PointInfo{
 8 | 	int x,y,z;
 9 | 	float r,g,b;
10 | };
11 | 
12 | extern "C"{
13 | 
14 | void render_ball(int h,int w,unsigned char * show,int n,int * xyzs,float * c0,float * c1,float * c2,int r){
15 | 	r=max(r,1);
16 | 	vector<int> depth(h*w,-2100000000);
17 | 	vector<PointInfo> pattern;
18 | 	for (int dx=-r;dx<=r;dx++)
19 | 		for (int dy=-r;dy<=r;dy++)
20 | 			if (dx*dx+dy*dy<r*r){
21 | 				double dz=sqrt(double(r*r-dx*dx-dy*dy));
22 | 				PointInfo pinfo;
23 | 				pinfo.x=dx;
24 | 				pinfo.y=dy;
25 | 				pinfo.z=dz;
26 | 				pinfo.r=dz/r;
27 | 				pinfo.g=dz/r;
28 | 				pinfo.b=dz/r;
29 | 				pattern.push_back(pinfo);
30 | 			}
31 | 	double zmin=0,zmax=0;
32 | 	for (int i=0;i<n;i++){
33 | 		if (i==0){
34 | 			zmin=xyzs[i*3+2]-r;
35 | 			zmax=xyzs[i*3+2]+r;
36 | 		}else{
37 | 			zmin=min(zmin,double(xyzs[i*3+2]-r));
38 | 			zmax=max(zmax,double(xyzs[i*3+2]+r));
39 | 		}
40 | 	}
41 | 	for (int i=0;i<n;i++){
42 | 		int x=xyzs[i*3+0],y=xyzs[i*3+1],z=xyzs[i*3+2];
43 | 		for (int j=0;j<int(pattern.size());j++){
44 | 			int x2=x+pattern[j].x;
45 | 			int y2=y+pattern[j].y;
46 | 			int z2=z+pattern[j].z;
47 | 			if (!(x2<0 || x2>=h || y2<0 || y2>=w) && depth[x2*w+y2]<z2){
48 | 				depth[x2*w+y2]=z2;
49 | 				double intensity=min(1.0,(z2-zmin)/(zmax-zmin)*0.7+0.3);
50 | 				show[(x2*w+y2)*3+0]=pattern[j].b*c2[i]*intensity;
51 | 				show[(x2*w+y2)*3+1]=pattern[j].g*c0[i]*intensity;
52 | 				show[(x2*w+y2)*3+2]=pattern[j].r*c1[i]*intensity;
53 | 			}
54 | 		}
55 | 	}
56 | }
57 | 
58 | }//extern "C"
59 | 


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | pytorch>=0.4.1
2 | torchvision
3 | numpy
4 | tqdm
5 | opencv-python
6 | lera # to monitor training
7 | 
8 | 


--------------------------------------------------------------------------------
/show3d_balls.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | import ctypes as ct
  3 | import cv2
  4 | import sys
  5 | showsz=800
  6 | mousex,mousey=0.5,0.5
  7 | zoom=1.0
  8 | changed=True
  9 | def onmouse(*args):
 10 |     global mousex,mousey,changed
 11 |     y=args[1]
 12 |     x=args[2]
 13 |     mousex=x/float(showsz)
 14 |     mousey=y/float(showsz)
 15 |     changed=True
 16 | cv2.namedWindow('show3d')
 17 | cv2.moveWindow('show3d',0,0)
 18 | cv2.setMouseCallback('show3d',onmouse)
 19 | 
 20 | dll=np.ctypeslib.load_library('render_balls_so','.')
 21 | 
 22 | def showpoints(xyz,c_gt=None, c_pred = None ,waittime=0,showrot=False,magnifyBlue=0,freezerot=False,background=(0, 0, 0),normalizecolor=True,ballradius=2):
 23 |     global showsz,mousex,mousey,zoom,changed
 24 |     xyz=xyz-xyz.mean(axis=0)
 25 |     radius=((xyz**2).sum(axis=-1)**0.5).max()
 26 |     xyz/=(radius*2.2)/showsz
 27 |     if c_gt is None:
 28 |         c0=np.zeros((len(xyz),),dtype='float32')+255
 29 |         c1=np.zeros((len(xyz),),dtype='float32')+255
 30 |         c2=np.zeros((len(xyz),),dtype='float32')+255
 31 |     else:
 32 |         c0=c_gt[:,0]
 33 |         c1=c_gt[:,1]
 34 |         c2=c_gt[:,2]
 35 | 
 36 | 
 37 |     if normalizecolor:
 38 |         c0/=(c0.max()+1e-14)/255.0
 39 |         c1/=(c1.max()+1e-14)/255.0
 40 |         c2/=(c2.max()+1e-14)/255.0
 41 | 
 42 | 
 43 |     c0=np.require(c0,'float32','C')
 44 |     c1=np.require(c1,'float32','C')
 45 |     c2=np.require(c2,'float32','C')
 46 | 
 47 |     show=np.zeros((showsz,showsz,3),dtype='uint8')
 48 |     def render():
 49 |         rotmat=np.eye(3)
 50 |         if not freezerot:
 51 |             xangle=(mousey-0.5)*np.pi*1.2
 52 |         else:
 53 |             xangle=0
 54 |         rotmat=rotmat.dot(np.array([
 55 |             [1.0,0.0,0.0],
 56 |             [0.0,np.cos(xangle),-np.sin(xangle)],
 57 |             [0.0,np.sin(xangle),np.cos(xangle)],
 58 |             ]))
 59 |         if not freezerot:
 60 |             yangle=(mousex-0.5)*np.pi*1.2
 61 |         else:
 62 |             yangle=0
 63 |         rotmat=rotmat.dot(np.array([
 64 |             [np.cos(yangle),0.0,-np.sin(yangle)],
 65 |             [0.0,1.0,0.0],
 66 |             [np.sin(yangle),0.0,np.cos(yangle)],
 67 |             ]))
 68 |         rotmat*=zoom
 69 |         nxyz=xyz.dot(rotmat)+[showsz/2,showsz/2,0]
 70 | 
 71 |         ixyz=nxyz.astype('int32')
 72 |         show[:]=background
 73 |         dll.render_ball(
 74 |             ct.c_int(show.shape[0]),
 75 |             ct.c_int(show.shape[1]),
 76 |             show.ctypes.data_as(ct.c_void_p),
 77 |             ct.c_int(ixyz.shape[0]),
 78 |             ixyz.ctypes.data_as(ct.c_void_p),
 79 |             c0.ctypes.data_as(ct.c_void_p),
 80 |             c1.ctypes.data_as(ct.c_void_p),
 81 |             c2.ctypes.data_as(ct.c_void_p),
 82 |             ct.c_int(ballradius)
 83 |         )
 84 | 
 85 |         if magnifyBlue>0:
 86 |             show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],1,axis=0))
 87 |             if magnifyBlue>=2:
 88 |                 show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],-1,axis=0))
 89 |             show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],1,axis=1))
 90 |             if magnifyBlue>=2:
 91 |                 show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],-1,axis=1))
 92 |         if showrot:
 93 |             cv2.putText(show,'xangle %d'%(int(xangle/np.pi*180)),(30,showsz-30),0,0.5,cv2.cv.CV_RGB(255,0,0))
 94 |             cv2.putText(show,'yangle %d'%(int(yangle/np.pi*180)),(30,showsz-50),0,0.5,cv2.cv.CV_RGB(255,0,0))
 95 |             cv2.putText(show,'zoom %d%%'%(int(zoom*100)),(30,showsz-70),0,0.5,cv2.cv.CV_RGB(255,0,0))
 96 |     changed=True
 97 |     while True:
 98 |         if changed:
 99 |             render()
100 |             changed=False
101 |         cv2.imshow('show3d',show)
102 |         if waittime==0:
103 |             cmd=cv2.waitKey(10)%256
104 |         else:
105 |             cmd=cv2.waitKey(waittime)%256
106 |         if cmd==ord('q'):
107 |             break
108 |         elif cmd==ord('Q'):
109 |             sys.exit(0)
110 | 
111 |         if cmd==ord('t') or cmd == ord('p'):
112 |             if cmd == ord('t'):
113 |                 if c_gt is None:
114 |                     c0=np.zeros((len(xyz),),dtype='float32')+255
115 |                     c1=np.zeros((len(xyz),),dtype='float32')+255
116 |                     c2=np.zeros((len(xyz),),dtype='float32')+255
117 |                 else:
118 |                     c0=c_gt[:,0]
119 |                     c1=c_gt[:,1]
120 |                     c2=c_gt[:,2]
121 |             else:
122 |                 if c_pred is None:
123 |                     c0=np.zeros((len(xyz),),dtype='float32')+255
124 |                     c1=np.zeros((len(xyz),),dtype='float32')+255
125 |                     c2=np.zeros((len(xyz),),dtype='float32')+255
126 |                 else:
127 |                     c0=c_pred[:,0]
128 |                     c1=c_pred[:,1]
129 |                     c2=c_pred[:,2]
130 |             if normalizecolor:
131 |                 c0/=(c0.max()+1e-14)/255.0
132 |                 c1/=(c1.max()+1e-14)/255.0
133 |                 c2/=(c2.max()+1e-14)/255.0
134 |             c0=np.require(c0,'float32','C')
135 |             c1=np.require(c1,'float32','C')
136 |             c2=np.require(c2,'float32','C')
137 |             changed = True
138 | 
139 | 
140 | 
141 |         if cmd==ord('n'):
142 |             zoom*=1.1
143 |             changed=True
144 |         elif cmd==ord('m'):
145 |             zoom/=1.1
146 |             changed=True
147 |         elif cmd==ord('r'):
148 |             zoom=1.0
149 |             changed=True
150 |         elif cmd==ord('s'):
151 |             cv2.imwrite('show3d.png',show)
152 |         if waittime!=0:
153 |             break
154 |     return cmd
155 | if __name__=='__main__':
156 |     np.random.seed(100)
157 |     showpoints(np.random.randn(2500,3))
158 | 
159 | 


--------------------------------------------------------------------------------
/show_seg.py:
--------------------------------------------------------------------------------
 1 | # -*- coding: utf-8 -*-
 2 | 
 3 | import argparse
 4 | import os
 5 | import random
 6 | 
 7 | import numpy as np
 8 | import matplotlib.pyplot as plt
 9 | import torch
10 | import torch.nn as nn
11 | 
12 | from datasets import PartDataset
13 | from pointnet import PointNetSeg
14 | from utils import shapenet_labels
15 | from show3d_balls import *
16 | 
17 | parser = argparse.ArgumentParser()
18 | 
19 | parser.add_argument('--model', type=str, default='', help='model path')
20 | parser.add_argument('--idx', type=int, default=0, help='model index')
21 | parser.add_argument('--className', type=str, default='Chair', help='number of classes')
22 | parser.add_argument('--radius', type=int, default=2, help='radius of ball for visualization')
23 | parser.add_argument('--cmap', type=str, default='hsv', help='valid matplotlib cmap')
24 | parser.add_argument('--npoints', type=int, default=2500, help='points to sample')
25 | 
26 | opt = parser.parse_args()
27 | 
28 | idx = opt.idx
29 | 
30 | d = PartDataset(root='shapenetcore_partanno_segmentation_benchmark_v0', class_choice=[opt.className], train=False, npoints=2048)
31 | 
32 | print('model %d/%d' % (idx, len(d)))
33 | 
34 | num_class = d.num_classes
35 | print('number of classes', num_class)
36 | 
37 | point, seg = d[idx]
38 | 
39 | point_np = point.numpy()
40 | 
41 | cmap = plt.cm.get_cmap(opt.cmap, 10)
42 | cmap = np.array([cmap(i) for i in range(10)])[:, :3]
43 | gt = cmap[seg.numpy()-1, :]
44 | 
45 | classifier = PointNetSeg(k=shapenet_labels[opt.className])
46 | classifier.load_state_dict(torch.load(opt.model))
47 | classifier.eval()
48 | 
49 | point = point.transpose(1, 0).contiguous()
50 | point = point.view(1, point.size()[0], point.size()[1])
51 | 
52 | 
53 | pred, _ = classifier(point)
54 | 
55 | pred_choice = pred.data.max(2)[1]
56 | 
57 | pred_color = cmap[pred_choice.numpy()[0], :]
58 | 
59 | showpoints(point_np, gt, gt, ballradius=opt.radius)
60 | 
61 | 
62 | 


--------------------------------------------------------------------------------
/show_seg_s3d.py:
--------------------------------------------------------------------------------
 1 | # -*- coding: utf-8 -*-
 2 | 
 3 | import argparse
 4 | import os
 5 | import random
 6 | 
 7 | import numpy as np
 8 | import matplotlib.pyplot as plt
 9 | import torch
10 | import torch.nn as nn
11 | 
12 | from pointnet import PointNetSeg
13 | from datasets import S3dDataset
14 | from datasets import scale_linear_bycolumn
15 | 
16 | from show3d_balls import *
17 | 
18 | 
19 | def parse_whole_scene(scene_path, scene_num, npoints=4096):
20 |     scene = np.loadtxt(scene_path).astype(np.float32)
21 |     seg = scene[:, -1].astype(np.int64)
22 |     scene = scene[:, :3]
23 |     replace = False if (scene_num*npoints<scene.shape[0]) else True
24 |     choice = np.random.choice(scene.shape[0], scene_num*npoints, replace=replace)
25 |     scene = scene[choice, :]
26 |     seg = seg[choice]
27 |     scene = scale_linear_bycolumn(scene)
28 |     scene = torch.from_numpy(scene)
29 |     seg = torch.from_numpy(seg)
30 |     return scene, seg
31 | 
32 | parser = argparse.ArgumentParser()
33 | 
34 | parser.add_argument('--scene_path', type=str, default='', help='path to scene')
35 | parser.add_argument('--scene_num', type=int, default='', help='number of scene components')
36 | parser.add_argument('--model', type=str, default='', help='model path')
37 | parser.add_argument('--className', type=str, default='Chair', help='number of classes')
38 | parser.add_argument('--radius', type=int, default=2, help='radius of ball for visualization')
39 | parser.add_argument('--cmap', type=str, default='hsv', help='valid matplotlib cmap')
40 | parser.add_argument('--npoints', type=int, default=2500, help='points to sample')
41 | 
42 | opt = parser.parse_args()
43 | 
44 | num_classes = 14
45 | 
46 | point, seg = parse_whole_scene(opt.scene_path, opt.scene_num, int(opt.npoints))
47 | 
48 | point_np = point.numpy()
49 | 
50 | cmap = plt.cm.get_cmap(opt.cmap, 14)
51 | cmap = np.array([cmap(i) for i in range(14)])[:, :3]
52 | gt = cmap[seg.numpy(), :]
53 | 
54 | classifier = PointNetSeg(k=num_classes)
55 | 
56 | state_dict = torch.load(opt.model)
57 | 
58 | '''
59 | *** load multi-gpu model ***
60 | from collections import OrderedDict
61 | 
62 | new_state_dict = OrderedDict()
63 | for k, v in state_dict.items():
64 |     name = k[7:]
65 |     new_state_dict[name] = v
66 | '''
67 | 
68 | classifier.load_state_dict(state_dict)
69 | classifier.eval()
70 | 
71 | point = point.transpose(1, 0).contiguous()
72 | point = point.view(1, point.size()[0], point.size()[1])
73 | 
74 | 
75 | pred, _ = classifier(point)
76 | 
77 | pred_choice = pred.data.max(2)[1]
78 | 
79 | pred_color = cmap[pred_choice.numpy()[0], :]
80 | 
81 | showpoints(point_np, gt, gt, ballradius=opt.radius)
82 | 
83 | 


--------------------------------------------------------------------------------
/train_cls.py:
--------------------------------------------------------------------------------
  1 | # -*- coding: utf-8 -*-
  2 | 
  3 | import argparse
  4 | import os
  5 | import pdb
  6 | 
  7 | import numpy as np
  8 | import torch
  9 | import torch.nn as nn
 10 | import torch.nn.functional as F
 11 | import torch.optim as optim
 12 | import lera
 13 | 
 14 | from datasets import ClsDataset
 15 | from pointnet import PointNetCls
 16 | 
 17 | 
 18 | def train(config):
 19 |     print('Random seed: %d' % int(config.seed))
 20 |     torch.manual_seed(config.seed)
 21 |     
 22 |     torch.backends.cudnn.benchmark = True
 23 | 
 24 |     dset = config.dataset
 25 |     if dset == 'modelnet10' or dset == 'modelnet40':
 26 |         dataset = ClsDataset(root=config.root, npoints=config.npoints, train=True)
 27 |         test_dataset = ClsDataset(root=config.root, npoints=config.npoints, train=False)
 28 |     else:
 29 |         raise NotImplementedError('Dataset not supported.')
 30 |     
 31 |     print('Selected %s' % dset)
 32 |     dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.batchsize, shuffle=True, 
 33 |                 num_workers=config.workers)
 34 |     test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=config.batchsize, shuffle=True, 
 35 |         num_workers=config.workers)
 36 | 
 37 |     num_classes = dataset.num_classes
 38 |     print('number of classes: %d' % num_classes)
 39 |     print('train set size: %d | test set size: %d' % (len(dataset), len(test_dataset)))
 40 |     try:
 41 |         os.makedirs(config.outf)
 42 |     except:
 43 |         pass
 44 | 
 45 |     blue = lambda x: '\033[94m' + x + '\033[0m'
 46 |     yellow = lambda x: '\033[93m' + x + '\033[0m'
 47 |     red = lambda x: '\033[91m' + x + '\033[0m'
 48 | 
 49 |     classifier = PointNetCls(k=num_classes)
 50 | 
 51 |     if config.model != '':
 52 |         classifier.load_state_dict(torch.load(config.model))
 53 | 
 54 |     optimizer = optim.SGD(classifier.parameters(), lr=config.lr, momentum=config.momentum)
 55 | 
 56 |     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 57 |     classifier.to(device)
 58 |     if config.mgpu:
 59 |         classifier = torch.nn.DataParallel(classifier, device_ids=config.gpuids)
 60 | 
 61 |     num_batch = len(dataset) / config.batchsize
 62 | 
 63 |     lera.log_hyperparams({
 64 |         'title': dset, 
 65 |         'batchsize': config.batchsize, 
 66 |         'epochs': config.nepochs, 
 67 |         'npoints': config.npoints, 
 68 |         'optimizer': 'SGD', 
 69 |         'lr': config.lr, 
 70 |         })
 71 | 
 72 |     for epoch in range(config.nepochs):
 73 |         train_acc_epoch, test_acc_epoch = [], []
 74 |         for i, data in enumerate(dataloader):
 75 |             points, labels = data
 76 |             points = points.transpose(2, 1)
 77 |             labels = labels[:, 0]
 78 |             points, labels = points.to(device), labels.to(device)
 79 |             optimizer.zero_grad()
 80 |             classifier = classifier.train()
 81 |             pred, _ = classifier(points)
 82 |             pred = pred.view(-1, num_classes)
 83 |             # print(pred.size(), labels.size())
 84 |             loss = F.nll_loss(pred, labels)
 85 |             loss.backward()
 86 |             optimizer.step()
 87 |             pred_choice = pred.data.max(1)[1]
 88 |             correct = pred_choice.eq(labels.data).cpu().sum()
 89 |             train_acc = correct.item() / float(config.batchsize)
 90 |             print('epoch %d: %d/%d | train loss: %f | train acc: %f' % (epoch+1, i+1, num_batch+1, loss.item(), train_acc))
 91 |             train_acc_epoch.append(train_acc)
 92 |             lera.log({
 93 |                 'train loss': loss.item(), 
 94 |                 'train acc': train_acc
 95 |                 })
 96 | 
 97 |             if (i+1) % 10 == 0:
 98 |                 j, data = next(enumerate(test_dataloader, 0))
 99 |                 points, labels = data
100 |                 points = points.transpose(2, 1)
101 |                 labels = labels[:, 0]
102 |                 points, labels = points.to(device), labels.to(device)
103 |                 classifier = classifier.eval()
104 |                 with torch.no_grad():
105 |                     pred, _ = classifier(points)
106 |                 pred = pred.view(-1, num_classes)
107 |                 loss = F.nll_loss(pred, labels)
108 |                 pred_choice = pred.data.max(1)[1]
109 |                 correct = pred_choice.eq(labels.data).cpu().sum()
110 |                 test_acc = correct.item() / float(config.batchsize)
111 |                 print(blue('epoch %d: %d/%d | test loss: %f | test acc: %f') % (epoch+1, i+1, num_batch+1, loss.item(), test_acc))
112 |                 test_acc_epoch.append(test_acc)
113 |                 lera.log({
114 |                     'test loss': loss.item(), 
115 |                     'test acc': test_acc
116 |                     })
117 |         print(yellow('epoch %d | mean train acc: %f') % (epoch+1, np.mean(train_acc_epoch)))
118 |         print(red('epoch %d | mean test acc: %f') % (epoch+1, np.mean(test_acc_epoch)))
119 |         lera.log({
120 |             'train acc epoch': np.mean(train_acc_epoch), 
121 |             'test acc epoch': np.mean(test_acc_epoch)})
122 |         torch.save(classifier.state_dict(), '%s/%s_model_%d.pth' % (config.outf, config.dataset, epoch))
123 | 
124 | 
125 | if __name__ == '__main__':
126 | 
127 |     parser = argparse.ArgumentParser()
128 |     parser.add_argument('-s', '--seed', type=int, help='random seed')
129 |     parser.add_argument('-dset', '--dataset', type=str, required=True, help='dataset to train on, one of modelnet10 and modelnet40)
130 |     parser.add_argument('-r', '--root', type=str, required=True, help='path to dataset')
131 |     parser.add_argument('-np', '--npoints', type=int, help='number of points to sample')
132 |     parser.add_argument('-bs', '--batchsize', type=int, default=32, help='batch size')
133 |     parser.add_argument('-ws', '--workers', type=int, default=4, help='number of workers')
134 |     parser.add_argument('-out', '--outf', type=str, default='./checkpoints', help='path to save model checkpoints')
135 |     parser.add_argument('--model', type=str, default='', help='checkpoint dir')
136 |     parser.add_argument('--lr', type=float, default=0.001, help='learning rate')
137 |     parser.add_argument('--momentum', type=float, default=0.9, help='SGD momentum')
138 |     parser.add_argument('--mgpu', type=bool, default=False, help='whether to utilize multiple gpus')
139 |     parser.add_argument('--gpuids', nargs='+', type=int, help='which gpus to use')
140 |     parser.add_argument('--nepochs', type=int, default=100, help='epochs to train')
141 |     config = parser.parse_args()
142 | 
143 |     train(config)
144 | 
145 | 
146 | 


--------------------------------------------------------------------------------
/train_seg.py:
--------------------------------------------------------------------------------
  1 | # -*- coding: utf-8 -*-
  2 | 
  3 | import argparse
  4 | import os
  5 | import pdb
  6 | 
  7 | import numpy as np
  8 | import torch
  9 | import torch.nn as nn
 10 | import torch.nn.functional as F
 11 | import torch.optim as optim
 12 | import lera
 13 | 
 14 | from datasets import PartDataset
 15 | from datasets import S3dDataset
 16 | from pointnet import PointNetSeg
 17 | 
 18 | def train(config):
 19 |     print('Random seed: %d' % int(config.seed))
 20 |     torch.manual_seed(config.seed)
 21 |     
 22 |     torch.backends.cudnn.benchmark = True
 23 | 
 24 |     dset = config.dataset
 25 |     if dset == 'shapenet16':
 26 |         dataset = PartDataset(root=config.root, npoints=config.npoints, class_choice=[config.classname], train=True)
 27 |         test_dataset = PartDataset(root=config.root, npoints=config.npoints, class_choice=[config.classname], train=False)
 28 |     elif dset == 's3dis':
 29 |         dataset = S3dDataset(root=config.root, npoints=config.npoints, train=True)
 30 |         test_dataset = S3dDataset(root=config.root, npoints=config.npoints, train=False)
 31 |     else:
 32 |         raise NotImplementedError('Dataset not supported.')
 33 |     
 34 |     print('Selected %s' % dset)
 35 |     dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.batchsize, shuffle=True, 
 36 |                 num_workers=config.workers)
 37 |     test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=config.batchsize, shuffle=True, 
 38 |         num_workers=config.workers)
 39 | 
 40 |     num_classes = dataset.num_classes
 41 |     print('number of classes: %d' % num_classes)
 42 |     print('train set size: %d | test set size: %d' % (len(dataset), len(test_dataset)))
 43 |     try:
 44 |         os.makedirs(config.outf)
 45 |     except:
 46 |         pass
 47 | 
 48 |     blue = lambda x: '\033[94m' + x + '\033[0m'
 49 |     yellow = lambda x: '\033[93m' + x + '\033[0m'
 50 |     red = lambda x: '\033[91m' + x + '\033[0m'
 51 | 
 52 |     classifier = PointNetSeg(k=num_classes)
 53 | 
 54 |     if config.model != '':
 55 |         classifier.load_state_dict(torch.load(config.model))
 56 | 
 57 |     optimizer = optim.SGD(classifier.parameters(), lr=config.lr, momentum=config.momentum)
 58 | 
 59 |     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 60 |     classifier.to(device)
 61 |     if config.mgpu:
 62 |         classifier = torch.nn.DataParallel(classifier, device_ids=config.gpuids)
 63 | 
 64 |     num_batch = len(dataset) / config.batchsize
 65 | 
 66 |     lera.log_hyperparams({
 67 |         'title': dset,
 68 |         'classname': config.classname, 
 69 |         'batchsize': config.batchsize, 
 70 |         'epochs': config.nepochs, 
 71 |         'optimizer': 'SGD', 
 72 |         'lr': config.lr, 
 73 |         'npoints': config.npoints
 74 |         })
 75 | 
 76 |     for epoch in range(config.nepochs):
 77 |         train_acc_epoch, train_iou_epoch, test_acc_epoch, test_iou_epoch = [], [], [], []
 78 |         for i, data in enumerate(dataloader):
 79 |             points, labels = data
 80 |             points = points.transpose(2, 1)
 81 |             points, labels = points.to(device), labels.to(device)
 82 |             optimizer.zero_grad()
 83 |             classifier = classifier.train()
 84 |             pred, _ = classifier(points)
 85 |             pred = pred.view(-1, num_classes)
 86 |             # print(pred.size(), labels.size())
 87 |             if dset == 'shapenet16':
 88 |                 labels = labels.view(-1, 1)[:, 0] - 1
 89 |             else:
 90 |                 labels = labels.view(-1, 1)[:, 0]
 91 |             # print(pred.size(), labels.size())
 92 |             # pdb.set_trace()
 93 |             loss = F.nll_loss(pred, labels)
 94 |             loss.backward()
 95 |             optimizer.step()
 96 |             pred_choice = pred.data.max(1)[1]
 97 |             correct = pred_choice.eq(labels.data).cpu().sum()
 98 |             train_acc = correct.item() / float(config.batchsize*config.npoints)
 99 |             train_iou = correct.item() / float(2*config.batchsize*config.npoints-correct.item())
100 |             print('epoch %d: %d/%d | train loss: %f | train acc: %f | train iou: %f' % (epoch+1, i+1, num_batch+1, loss.item(), train_acc, train_iou))
101 |             train_acc_epoch.append(train_acc)
102 |             train_iou_epoch.append(train_iou)
103 |             lera.log({
104 |                 'train loss': loss.item(), 
105 |                 'train acc': train_acc, 
106 |                 'train IoU': train_iou}
107 |                 )
108 | 
109 |             if (i+1) % 10 == 0:
110 |                 j, data = next(enumerate(test_dataloader, 0))
111 |                 points, labels = data
112 |                 points = points.transpose(2, 1)
113 |                 points, labels = points.to(device), labels.to(device)
114 |                 classifier = classifier.eval()
115 |                 with torch.no_grad():
116 |                     pred, _ = classifier(points)
117 |                 pred = pred.view(-1, num_classes)
118 |                 if dset == 'shapenet16':
119 |                     labels = labels.view(-1, 1)[:, 0] - 1
120 |                 else:
121 |                     labels = labels.view(-1, 1)[:, 0]
122 |                 loss = F.nll_loss(pred, labels)
123 |                 pred_choice = pred.data.max(1)[1]
124 |                 correct = pred_choice.eq(labels.data).cpu().sum()
125 |                 test_acc = correct.item() / float(config.batchsize*config.npoints)
126 |                 test_iou = correct.item() / float(2*config.batchsize*config.npoints-correct.item())
127 |                 print(blue('epoch %d: %d/%d | test loss: %f | test acc: %f | test iou: %f') % (epoch+1, i+1, num_batch+1, loss.item(), test_acc, test_iou))
128 |                 test_acc_epoch.append(test_acc)
129 |                 test_iou_epoch.append(test_iou)
130 |                 lera.log({
131 |                     'test loss': loss.item(), 
132 |                     'test acc': test_acc, 
133 |                     'test IoU': test_iou})
134 |         print(yellow('epoch %d | mean train acc: %f | mean train IoU: %f') % (epoch+1, np.mean(train_acc_epoch), np.mean(train_iou_epoch)))
135 |         print(red('epoch %d | mean test acc: %f | mean test IoU: %f') % (epoch+1, np.mean(test_acc_epoch), np.mean(test_iou_epoch)))
136 |         lera.log({
137 |             'mean train acc': np.mean(train_acc_epoch), 
138 |             'mean train iou': np.mean(train_iou_epoch), 
139 |             'mean test acc': np.mean(test_acc_epoch), 
140 |             'mean test iou': np.mean(test_iou_epoch)})
141 |         torch.save(classifier.state_dict(), '%s/%s_model_%d.pth' % (config.outf, config.dataset, epoch))
142 | 
143 | 
144 | if __name__ == '__main__':
145 | 
146 |     parser = argparse.ArgumentParser()
147 |     parser.add_argument('-s', '--seed', type=int, help='random seed')
148 |     parser.add_argument('-dset', '--dataset', type=str, required=True, help='dataset to train on, one of modelnet, shapenet16 and s3dis')
149 |     parser.add_argument('-c', '--classname', type=str, default='Chair', help='one of 16 categories on shapenet16')
150 |     parser.add_argument('-r', '--root', type=str, required=True, help='path to dataset')
151 |     parser.add_argument('-np', '--npoints', type=int, help='number of points to sample')
152 |     parser.add_argument('-bs', '--batchsize', type=int, default=32, help='batch size')
153 |     parser.add_argument('-ws', '--workers', type=int, default=4, help='number of workers')
154 |     parser.add_argument('-out', '--outf', type=str, default='./checkpoints', help='path to save model checkpoints')
155 |     parser.add_argument('--model', type=str, default='', help='checkpoint dir')
156 |     parser.add_argument('--lr', type=float, default=0.001, help='learning rate')
157 |     parser.add_argument('--momentum', type=float, default=0.9, help='momentum in SGD')
158 |     parser.add_argument('--mgpu', type=bool, default=False, help='whether to utilize multiple gpus')
159 |     parser.add_argument('--gpuids', nargs='+', type=int, help='which gpus to use')
160 |     parser.add_argument('--nepochs', type=int, default=100, help='epochs to train')
161 |     config = parser.parse_args()
162 | 
163 |     train(config)
164 | 
165 | 
166 | 


--------------------------------------------------------------------------------
/utils.py:
--------------------------------------------------------------------------------
 1 | # -*- coding: utf-8 -*-
 2 | 
 3 | import os
 4 | 
 5 | def modelnet_cat2num(modelnet_root):
 6 |     for i, item in enumerate(os.listdir(modelnet_root)):
 7 |         with open(os.path.join(modelnet_root, 'modelnet_cat2num.txt'), 'a') as f:
 8 |             f.write(item + ' ' + str(i) + '\n')
 9 | 
10 | 
11 | shapenet_labels = {'Airplane': 4, 
12 |         'Bag': 2, 
13 |         'Cap': 2, 
14 |         'Car': 4, 
15 |         'Chair': 4, 
16 |         'Earphone': 3, 
17 |         'Guitar': 3, 
18 |         'Knife': 2, 
19 |         'Lamp': 4, 
20 |         'Laptop': 2, 
21 |         'Motorbike': 6, 
22 |         'Mug': 2, 
23 |         'Pistol': 3, 
24 |         'Rocket': 3, 
25 |         'Skateboard': 3, 
26 |         'Table': 3
27 |         }
28 | 
29 | def s3d_cat2num(s3d_root):
30 |     for item in s3d_cat2num:
31 |         with open(os.path.join(s3d_root, 's3d_cat2num.txt'), 'a') as f:
32 |             f.write(item + ' ' + str(s3d_cat2num[item]) + '\n')
33 |             
34 | 
35 | if __name__ == '__main__':
36 | 
37 |     '''
38 |     modelnet_root = './ModelNet10'
39 |     modelnet_cat2num(modelnet_root)
40 | 
41 |     s3d_root = './Stanford3dDataset_v1.2'
42 |     s3d_cat2num(s3d_root)
43 |     '''
44 | 


--------------------------------------------------------------------------------