├── .gitignore
├── sphere.npy
├── gaussian.npy
├── image
├── DGCNN.jpg
├── acc1.png
├── acc2.png
├── acc3.png
├── acc4.png
├── acc5.png
├── shapenetpartpart_train0_chair.png
├── shapenetpartpart_train58_airplane.png
└── shapenetpartpart_train38_skateboard.png
├── tensorboard
├── Segment_dgcnn_seg_k40_1024_b32
│ └── events.out.tfevents.1579308861.server231
├── Classify_dgcnn_cls_k40_1024_b32
│ └── events.out.tfevents.1579684482.server231
├── Classify_dgcnn_seg_k40_1024_b32
│ └── events.out.tfevents.1579565090.server231
├── Reconstruct_dgcnn_cls_k20_1024_b16
│ └── events.out.tfevents.1579401032.server231
├── Reconstruct_dgcnn_seg_k40_1024_b32
│ └── events.out.tfevents.1579272092.server231
├── Classify_dgcnn_cls_k40_1024_b32_part
│ └── events.out.tfevents.1579698812.server231
├── Segment_dgcnn_seg_k40_1024_b32_nolabel
│ └── events.out.tfevents.1579336965.server231
└── Reconstruct_dgcnn_cls_k20_1024_b16_part
│ └── events.out.tfevents.1579685815.server231
├── LICENSE
├── utils.py
├── svm.py
├── main.py
├── dataset.py
├── loss.py
├── inference.py
├── reconstruction.py
├── README.md
├── classification.py
├── visualization.py
├── segmentation.py
└── model.py
/.gitignore:
--------------------------------------------------------------------------------
1 | .DS_Store
2 | */.DS_Store
3 |
--------------------------------------------------------------------------------
/sphere.npy:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/sphere.npy
--------------------------------------------------------------------------------
/gaussian.npy:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/gaussian.npy
--------------------------------------------------------------------------------
/image/DGCNN.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/image/DGCNN.jpg
--------------------------------------------------------------------------------
/image/acc1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/image/acc1.png
--------------------------------------------------------------------------------
/image/acc2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/image/acc2.png
--------------------------------------------------------------------------------
/image/acc3.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/image/acc3.png
--------------------------------------------------------------------------------
/image/acc4.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/image/acc4.png
--------------------------------------------------------------------------------
/image/acc5.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/image/acc5.png
--------------------------------------------------------------------------------
/image/shapenetpartpart_train0_chair.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/image/shapenetpartpart_train0_chair.png
--------------------------------------------------------------------------------
/image/shapenetpartpart_train58_airplane.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/image/shapenetpartpart_train58_airplane.png
--------------------------------------------------------------------------------
/image/shapenetpartpart_train38_skateboard.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/image/shapenetpartpart_train38_skateboard.png
--------------------------------------------------------------------------------
/tensorboard/Segment_dgcnn_seg_k40_1024_b32/events.out.tfevents.1579308861.server231:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/tensorboard/Segment_dgcnn_seg_k40_1024_b32/events.out.tfevents.1579308861.server231
--------------------------------------------------------------------------------
/tensorboard/Classify_dgcnn_cls_k40_1024_b32/events.out.tfevents.1579684482.server231:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/tensorboard/Classify_dgcnn_cls_k40_1024_b32/events.out.tfevents.1579684482.server231
--------------------------------------------------------------------------------
/tensorboard/Classify_dgcnn_seg_k40_1024_b32/events.out.tfevents.1579565090.server231:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/tensorboard/Classify_dgcnn_seg_k40_1024_b32/events.out.tfevents.1579565090.server231
--------------------------------------------------------------------------------
/tensorboard/Reconstruct_dgcnn_cls_k20_1024_b16/events.out.tfevents.1579401032.server231:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/tensorboard/Reconstruct_dgcnn_cls_k20_1024_b16/events.out.tfevents.1579401032.server231
--------------------------------------------------------------------------------
/tensorboard/Reconstruct_dgcnn_seg_k40_1024_b32/events.out.tfevents.1579272092.server231:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/tensorboard/Reconstruct_dgcnn_seg_k40_1024_b32/events.out.tfevents.1579272092.server231
--------------------------------------------------------------------------------
/tensorboard/Classify_dgcnn_cls_k40_1024_b32_part/events.out.tfevents.1579698812.server231:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/tensorboard/Classify_dgcnn_cls_k40_1024_b32_part/events.out.tfevents.1579698812.server231
--------------------------------------------------------------------------------
/tensorboard/Segment_dgcnn_seg_k40_1024_b32_nolabel/events.out.tfevents.1579336965.server231:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/tensorboard/Segment_dgcnn_seg_k40_1024_b32_nolabel/events.out.tfevents.1579336965.server231
--------------------------------------------------------------------------------
/tensorboard/Reconstruct_dgcnn_cls_k20_1024_b16_part/events.out.tfevents.1579685815.server231:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/antao97/PointCloudSegmentation/HEAD/tensorboard/Reconstruct_dgcnn_cls_k20_1024_b16_part/events.out.tfevents.1579685815.server231
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2020 An Tao
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 |
--------------------------------------------------------------------------------
/utils.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- coding: utf-8 -*-
3 | """
4 | @Author: An Tao
5 | @Contact: ta19@mails.tsinghua.edu.cn
6 | @File: utils.py
7 | @Time: 2020/1/2 10:26 AM
8 | """
9 |
10 | import os
11 | import sys
12 |
13 |
14 | class AverageMeter(object):
15 | """Computes and stores the average and current value"""
16 | def __init__(self):
17 | self.reset()
18 |
19 | def reset(self):
20 | self.val = 0
21 | self.avg = 0
22 | self.sum = 0
23 | self.count = 0
24 |
25 | def update(self, val, n=1):
26 | self.val = val
27 | self.sum += val * n
28 | self.count += n
29 | self.avg = self.sum / self.count
30 |
31 |
32 | def learning_rate_decay(optimizer, t, lr_0):
33 | for param_group in optimizer.param_groups:
34 | lr = lr_0 / np.sqrt(1 + lr_0 * param_group['weight_decay'] * t)
35 | param_group['lr'] = lr
36 |
37 |
38 | class Logger(object):
39 | def __init__(self, fpath=None):
40 | self.console = sys.stdout
41 | self.file = None
42 | if fpath is not None:
43 | self.file = open(fpath, 'w')
44 |
45 | def __del__(self):
46 | self.close()
47 |
48 | def __enter__(self):
49 | pass
50 |
51 | def __exit__(self, *args):
52 | self.close()
53 |
54 | def write(self, msg):
55 | self.console.write(msg)
56 | if self.file is not None:
57 | self.file.write(msg)
58 |
59 | def flush(self):
60 | self.console.flush()
61 | if self.file is not None:
62 | self.file.flush()
63 | os.fsync(self.file.fileno())
64 |
65 | def close(self):
66 | self.console.close()
67 | if self.file is not None:
68 | self.file.close()
--------------------------------------------------------------------------------
/svm.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- coding: utf-8 -*-
3 | """
4 | @Author: An Tao
5 | @Contact: ta19@mails.tsinghua.edu.cn
6 | @File: svm.py
7 | @Time: 2020/1/2 10:26 AM
8 | """
9 |
10 | import os
11 | import h5py
12 | import numpy as np
13 | from glob import glob
14 | from sklearn.svm import LinearSVC
15 |
16 |
17 | class SVM(object):
18 | def __init__(self, feature_dir):
19 | self.feature_dir = feature_dir
20 |
21 | self.train_path = glob(os.path.join(self.feature_dir, 'train*.h5'))
22 | self.test_path = glob(os.path.join(self.feature_dir, 'test*.h5'))
23 |
24 | print("Loading feature dataset...")
25 | train_data = []
26 | train_label = []
27 | for path in self.train_path:
28 | f = h5py.File(path, 'r+')
29 | data = f['data'][:].astype('float32')
30 | label = f['label'][:].astype('int64')
31 | f.close()
32 | train_data.append(data)
33 | train_label.append(label)
34 | self.train_data = np.concatenate(train_data, axis=0)
35 | self.train_label = np.concatenate(train_label, axis=0)
36 | print("Training set size:", np.size(self.train_data, 0))
37 |
38 | test_data = []
39 | test_label = []
40 | for path in self.test_path:
41 | f = h5py.File(path, 'r+')
42 | data = f['data'][:].astype('float32')
43 | label = f['label'][:].astype('int64')
44 | f.close()
45 | test_data.append(data)
46 | test_label.append(label)
47 | self.test_data = np.concatenate(test_data, axis=0)
48 | self.test_label = np.concatenate(test_label, axis=0)
49 | print("Testing set size:", np.size(self.test_data, 0))
50 |
51 | def run(self):
52 | clf = LinearSVC(random_state=0)
53 | clf.fit(self.train_data, self.train_label)
54 | result = clf.predict(self.test_data)
55 | accuracy = np.sum(result==self.test_label).astype(float) / np.size(self.test_label)
56 | print("Transfer linear SVM accuracy: {:.2f}%".format(accuracy*100))
57 |
58 |
59 |
--------------------------------------------------------------------------------
/main.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- coding: utf-8 -*-
3 | """
4 | @Author: An Tao
5 | @Contact: ta19@mails.tsinghua.edu.cn
6 | @File: main.py
7 | @Time: 2020/1/2 10:26 AM
8 | """
9 |
10 | import argparse
11 |
12 | from reconstruction import Reconstruction
13 | from classification import Classification
14 | from segmentation import Segmentation
15 | from inference import Inference
16 | from svm import SVM
17 |
18 |
19 | def get_parser():
20 | parser = argparse.ArgumentParser(description='Unsupervised Point Cloud Feature Learning')
21 | parser.add_argument('--exp_name', type=str, default=None, metavar='N',
22 | help='Name of the experiment')
23 | parser.add_argument('--task', type=str, default='reconstruct', metavar='N',
24 | choices=['reconstruct', 'classify', 'segment'],
25 | help='Experiment task, [reconstruct, classify, segment]')
26 | parser.add_argument('--seg_no_class_label', action='store_true',
27 | help='Do not use class labels in segmentation')
28 | parser.add_argument('--loss', type=str, default='softmax', metavar='N',
29 | choices=['softmax', 'triplet'],
30 | help='Loss to use, [softmax, triplet]')
31 | parser.add_argument('--margin', type=float, default=None,
32 | help='Margin for triplet loss')
33 | parser.add_argument('--encoder', type=str, default='foldingnet', metavar='N',
34 | choices=['foldnet', 'dgcnn_cls', 'dgcnn_seg'],
35 | help='Encoder to use, [foldingnet, dgcnn_cls, dgcnn_seg]')
36 | parser.add_argument('--dropout', type=float, default=0.5,
37 | help='Dropout rate')
38 | parser.add_argument('--feat_dims', type=int, default=512, metavar='N',
39 | help='Number of dims for feature ')
40 | parser.add_argument('--k', type=int, default=None, metavar='N',
41 | help='Num of nearest neighbors to use for KNN')
42 | parser.add_argument('--shape', type=str, default='sphere', metavar='N',
43 | choices=['plane', 'sphere', 'gaussian'],
44 | help='Shape of points to input decoder, [plane, sphere, gaussian]')
45 | parser.add_argument('--dataset', type=str, default='shapenetcorev2', metavar='N',
46 | choices=['shapenetpart', 'modelnet40', 'modelnet10', 'shapenetpartpart'],
47 | help='Dataset to use, [shapenetpart, modelnet40, modelnet10, shapenetpartpart]')
48 | parser.add_argument('--class_choice', type=str, default=None, metavar='N',
49 | choices=['airplane', 'bag', 'cap', 'car', 'chair',
50 | 'earphone', 'guitar', 'knife', 'lamp', 'laptop',
51 | 'motorbike', 'mug', 'pistol', 'rocket', 'skateboard', 'table'])
52 | parser.add_argument('--no_scheduler', action='store_true',
53 | help='Do not use scheduler in training')
54 | parser.add_argument('--use_rotate', action='store_true',
55 | help='Rotate the pointcloud before training')
56 | parser.add_argument('--use_translate', action='store_true',
57 | help='Translate the pointcloud before training')
58 | parser.add_argument('--use_jitter', action='store_true',
59 | help='Jitter the pointcloud before training')
60 | parser.add_argument('--dataset_root', type=str, default='../dataset', help="Dataset root path")
61 | parser.add_argument('--gpu', type=str, help='Id of gpu device to be used', default='0')
62 | parser.add_argument('--batch_size', type=int, default=16, metavar='batch_size',
63 | help='Size of batch)')
64 | parser.add_argument('--workers', type=int, help='Number of data loading workers', default=16)
65 | parser.add_argument('--epochs', type=int, default=None, metavar='N',
66 | help='Number of episode to train ')
67 | parser.add_argument('--snapshot_interval', type=int, default=10, metavar='N',
68 | help='Save snapshot interval ')
69 | parser.add_argument('--no_cuda', action='store_true',
70 | help='Enables CUDA training')
71 | parser.add_argument('--eval', action='store_true',
72 | help='Evaluate the model')
73 | parser.add_argument('--num_points', type=int, default=2048,
74 | help='Num of points to use')
75 | parser.add_argument('--model_path', type=str, default='', metavar='N',
76 | help='Path to load model')
77 | args = parser.parse_args()
78 | return args
79 |
80 |
81 | if __name__ == '__main__':
82 | args = get_parser()
83 | if args.eval == False:
84 | if args.task == 'reconstruct':
85 | reconstruction = Reconstruction(args)
86 | reconstruction.run()
87 | elif args.task == 'classify':
88 | classification = Classification(args)
89 | classification.run()
90 | elif args.task == 'segment':
91 | segmentation = Segmentation(args)
92 | segmentation.run()
93 | else:
94 | inference = Inference(args)
95 | feature_dir = inference.run()
96 | svm = SVM(feature_dir)
97 | svm.run()
98 |
--------------------------------------------------------------------------------
/dataset.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- coding: utf-8 -*-
3 | """
4 | @Author: An Tao
5 | @Contact: ta19@mails.tsinghua.edu.cn
6 | @File: dataset.py
7 | @Time: 2020/1/2 10:26 AM
8 | """
9 |
10 | import os
11 | import torch
12 | import json
13 | import h5py
14 | from glob import glob
15 | import numpy as np
16 | import torch.utils.data as data
17 |
18 |
19 | shapenetpart_seg_num = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3]
20 | shapenetpart_seg_start_index = [0, 4, 6, 8, 12, 16, 19, 22, 24, 28, 30, 36, 38, 41, 44, 47]
21 |
22 |
23 | def translate_pointcloud(pointcloud):
24 | xyz1 = np.random.uniform(low=2./3., high=3./2., size=[3])
25 | xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3])
26 |
27 | translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype('float32')
28 | return translated_pointcloud
29 |
30 |
31 | def jitter_pointcloud(pointcloud, sigma=0.01, clip=0.02):
32 | N, C = pointcloud.shape
33 | pointcloud += np.clip(sigma * np.random.randn(N, C), -1*clip, clip)
34 | return pointcloud
35 |
36 |
37 | def rotate_pointcloud(pointcloud):
38 | theta = np.pi*2 * np.random.choice(24) / 24
39 | rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]])
40 | pointcloud[:,[0,2]] = pointcloud[:,[0,2]].dot(rotation_matrix) # random rotation (x,z)
41 | return pointcloud
42 |
43 |
44 | class Dataset(data.Dataset):
45 | def __init__(self, root, dataset_name='modelnet40', class_choice=None,
46 | num_points=2048, split='train', load_name=False,
47 | segmentation=False, random_rotate=False, random_jitter=False,
48 | random_translate=False):
49 |
50 | assert dataset_name.lower() in ['shapenetcorev2', 'shapenetpart',
51 | 'modelnet10', 'modelnet40', 'shapenetpartpart']
52 | assert num_points <= 2048
53 |
54 | if dataset_name in ['shapenetcorev2', 'shapenetpart', 'shapenetpartpart']:
55 | assert split.lower() in ['train', 'test', 'val', 'trainval', 'all']
56 | else:
57 | assert split.lower() in ['train', 'test', 'all']
58 |
59 | if dataset_name not in ['shapenetcorev2', 'shapenetpart', 'shapenetpartpart'] and segmentation == True:
60 | raise AssertionError
61 |
62 | self.root = os.path.join(root, dataset_name + '_' + '*hdf5_2048')
63 | self.dataset_name = dataset_name
64 | self.class_choice = class_choice
65 | self.num_points = num_points
66 | self.split = split
67 | self.load_name = load_name
68 | if self.dataset_name == 'shapenetpartpart':
69 | self.segmentation = True
70 | else:
71 | self.segmentation = segmentation
72 | self.random_rotate = random_rotate
73 | self.random_jitter = random_jitter
74 | self.random_translate = random_translate
75 |
76 | self.path_h5py_all = []
77 | self.path_name_all = []
78 | self.path_file_all = []
79 |
80 | if self.split in ['train','trainval','all']:
81 | self.get_path('train')
82 | if self.dataset_name in ['shapenetcorev2', 'shapenetpart', 'shapenetpartpart']:
83 | if self.split in ['val','trainval','all']:
84 | self.get_path('val')
85 | if self.split in ['test', 'all']:
86 | self.get_path('test')
87 |
88 | self.path_h5py_all.sort()
89 | data, label, seg = self.load_h5py(self.path_h5py_all)
90 |
91 | if self.load_name or self.class_choice != None:
92 | self.path_name_all.sort()
93 | self.name = self.load_json(self.path_name_all) # load label name
94 |
95 | self.data = np.concatenate(data, axis=0)
96 | self.label = np.concatenate(label, axis=0)
97 | if self.segmentation:
98 | self.seg = np.concatenate(seg, axis=0)
99 |
100 | if self.class_choice != None:
101 | indices = (self.name == class_choice).squeeze()
102 | self.data = self.data[indices]
103 | self.label = self.label[indices]
104 | if self.segmentation:
105 | self.seg = self.seg[indices]
106 | self.seg_num_all = shapenetpart_seg_num[id_choice]
107 | self.seg_start_index = shapenetpart_seg_start_index[id_choice]
108 | elif self.segmentation:
109 | self.seg_num_all = 50
110 | self.seg_start_index = 0
111 |
112 | def get_path(self, type):
113 | path_h5py = os.path.join(self.root, '*%s*.h5'%type)
114 | self.path_h5py_all += glob(path_h5py)
115 | if self.load_name:
116 | path_json = os.path.join(self.root, '%s*_id2name.json'%type)
117 | self.path_name_all += glob(path_json)
118 | return
119 |
120 | def load_h5py(self, path):
121 | all_data = []
122 | all_label = []
123 | all_seg = []
124 | for h5_name in path:
125 | f = h5py.File(h5_name, 'r+')
126 | data = f['data'][:].astype('float32')
127 | label = f['label'][:].astype('int64')
128 | if self.segmentation:
129 | seg = f['seg'][:].astype('int64')
130 | f.close()
131 | all_data.append(data)
132 | all_label.append(label)
133 | if self.segmentation:
134 | all_seg.append(seg)
135 | return all_data, all_label, all_seg
136 |
137 | def load_json(self, path):
138 | all_data = []
139 | for json_name in path:
140 | j = open(json_name, 'r+')
141 | data = json.load(j)
142 | all_data += data
143 | return all_data
144 |
145 | def __getitem__(self, item):
146 | point_set = self.data[item][:self.num_points]
147 | label = self.label[item]
148 | if self.load_name:
149 | name = self.name[item] # get label name
150 |
151 | if self.random_rotate:
152 | point_set = rotate_pointcloud(point_set)
153 | if self.random_jitter:
154 | point_set = jitter_pointcloud(point_set)
155 | if self.random_translate:
156 | point_set = translate_pointcloud(point_set)
157 |
158 | # convert numpy array to pytorch Tensor
159 | point_set = torch.from_numpy(point_set)
160 | label = torch.from_numpy(np.array([label]).astype(np.int64))
161 | label = label.squeeze(0)
162 |
163 | if self.segmentation:
164 | seg = self.seg[item]
165 | seg = torch.from_numpy(seg)
166 | if self.dataset_name == 'shapenetpartpart':
167 | return point_set, seg.unsqueeze(1)[0]
168 | else:
169 | return point_set, label, seg
170 | else:
171 | return point_set, label
172 |
173 | def __len__(self):
174 | return self.data.shape[0]
--------------------------------------------------------------------------------
/loss.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- coding: utf-8 -*-
3 | """
4 | @Author: An Tao
5 | @Contact: ta19@mails.tsinghua.edu.cn
6 | @File: loss.py
7 | @Time: 2020/1/2 10:26 AM
8 | """
9 |
10 | import copy
11 | import torch
12 | import torch.nn as nn
13 | import torch.nn.functional as F
14 | from torch.autograd import Variable
15 |
16 |
17 | def batch_pairwise_dist(x, y):
18 | bs, num_points_x, points_dim = x.size()
19 | _, num_points_y, _ = y.size()
20 | xx = torch.bmm(x, x.transpose(2, 1))
21 | yy = torch.bmm(y, y.transpose(2, 1))
22 | zz = torch.bmm(x, y.transpose(2, 1))
23 | diag_ind_x = torch.arange(0, num_points_x)
24 | diag_ind_y = torch.arange(0, num_points_y)
25 | if x.get_device() != -1:
26 | diag_ind_x = diag_ind_x.cuda(x.get_device())
27 | diag_ind_y = diag_ind_y.cuda(x.get_device())
28 | rx = xx[:, diag_ind_x, diag_ind_x].unsqueeze(1).expand_as(zz.transpose(2, 1))
29 | ry = yy[:, diag_ind_y, diag_ind_y].unsqueeze(1).expand_as(zz)
30 | P = (rx.transpose(2, 1) + ry - 2 * zz)
31 | return P
32 |
33 |
34 | class ChamferLoss(nn.Module):
35 | def __init__(self):
36 | super(ChamferLoss, self).__init__()
37 | self.use_cuda = torch.cuda.is_available()
38 |
39 | def forward(self, preds, gts):
40 | P = batch_pairwise_dist(gts, preds)
41 | mins, _ = torch.min(P, 1)
42 | loss_1 = torch.sum(mins)
43 | mins, _ = torch.min(P, 2)
44 | loss_2 = torch.sum(mins)
45 | return loss_1 + loss_2
46 |
47 |
48 | class CrossEntropyLoss(nn.Module):
49 | def __init__(self, smoothing=True):
50 | super(CrossEntropyLoss, self).__init__()
51 | self.smoothing = smoothing
52 |
53 | def forward(self, preds, gts):
54 | gts = gts.contiguous().view(-1)
55 |
56 | if self.smoothing:
57 | eps = 0.2
58 | n_class = preds.size(1)
59 |
60 | one_hot = torch.zeros_like(preds).scatter(1, gts.view(-1, 1), 1)
61 | one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1)
62 | log_prb = F.log_softmax(preds, dim=1)
63 |
64 | loss = -(one_hot * log_prb).sum(dim=1).mean()
65 | else:
66 | loss = F.cross_entropy(preds, gts, reduction='mean')
67 |
68 | return loss
69 |
70 |
71 | class TripletLoss(nn.Module):
72 | def __init__(self, margin=None, normalize_feature=True):
73 | super(TripletLoss, self).__init__()
74 | self.margin = margin
75 | self.normalize_feature = normalize_feature
76 | if self.margin is not None:
77 | self.ranking_loss = nn.MarginRankingLoss(margin=margin)
78 | else:
79 | self.ranking_loss = nn.SoftMarginLoss()
80 |
81 | def normalize(self, x, axis=-1):
82 | """Normalizing to unit length along the specified dimension.
83 | Args:
84 | x: pytorch Variable
85 | Returns:
86 | x: pytorch Variable, same shape as input
87 | """
88 | x = 1. * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12)
89 | return x
90 |
91 | def euclidean_dist(self, x, y):
92 | """
93 | Args:
94 | x: pytorch Variable, with shape [m, d]
95 | y: pytorch Variable, with shape [n, d]
96 | Returns:
97 | dist: pytorch Variable, with shape [m, n]
98 | """
99 | m, n = x.size(0), y.size(0)
100 | xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)
101 | yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t()
102 | dist = xx + yy
103 | dist.addmm_(1, -2, x, y.t())
104 | dist = dist.clamp(min=1e-12).sqrt() # for numerical stability
105 | return dist
106 |
107 | def hard_example_mining(self, dist_mat, labels, return_inds=False):
108 | """For each anchor, find the hardest positive and negative sample.
109 | Args:
110 | dist_mat: pytorch Variable, pair wise distance between samples, shape [B, N, N]
111 | labels: pytorch LongTensor, with shape [B, N]
112 | return_inds: whether to return the indices. Save time if `False`(?)
113 | Returns:
114 | dist_ap: pytorch Variable, distance(anchor, positive); shape [N]
115 | dist_an: pytorch Variable, distance(anchor, negative); shape [N]
116 | p_inds: pytorch LongTensor, with shape [N];
117 | indices of selected hard positive samples; 0 <= p_inds[i] <= N - 1
118 | n_inds: pytorch LongTensor, with shape [N];
119 | indices of selected hard negative samples; 0 <= n_inds[i] <= N - 1
120 | NOTE: Only consider the case in which all labels have same num of samples,
121 | thus we can cope with all anchors in parallel.
122 | """
123 |
124 | assert len(dist_mat.size()) == 3
125 | assert dist_mat.size(1) == dist_mat.size(2)
126 | B = dist_mat.size(0)
127 | N = dist_mat.size(1)
128 |
129 | # shape [B, N, N]
130 | is_pos = labels.unsqueeze(2).expand(B, N, N).eq(labels.unsqueeze(2).expand(B, N, N).transpose(2,1))
131 | is_neg = labels.unsqueeze(2).expand(B, N, N).ne(labels.unsqueeze(2).expand(B, N, N).transpose(2,1))
132 |
133 | # `dist_ap` means distance(anchor, positive)
134 | # both `dist_ap` and `relative_p_inds` with shape [B, N, 1]
135 | dist_mat_pos = torch.zeros(dist_mat.size(), dtype=torch.float32, device=dist_mat.device)
136 | dist_mat_pos[is_pos] = dist_mat[is_pos]
137 | dist_ap, relative_p_inds = torch.max(dist_mat_pos, 2, keepdim=True)
138 | # `dist_an` means distance(anchor, negative)
139 | # both `dist_an` and `relative_n_inds` with shape [B, N, 1]
140 | dist_mat_neg = torch.zeros(dist_mat.size(), dtype=torch.float32, device=dist_mat.device).fill_(1000)
141 | dist_mat_neg[is_neg] = dist_mat[is_neg]
142 | dist_an, relative_n_inds = torch.min(dist_mat_neg, 2, keepdim=True)
143 |
144 | # shape [B, N]
145 | dist_ap = dist_ap.squeeze(2)
146 | dist_an = dist_an.squeeze(2)
147 |
148 | return dist_ap, dist_an
149 |
150 | def forward(self, preds, gts, new_device):
151 | """
152 | Args:
153 | preds: pytorch Variable, shape [B, N, C]
154 | gts: pytorch LongTensor, with shape [B, N]
155 | Returns:
156 | loss: pytorch Variable, with shape [1]
157 | p_inds: pytorch LongTensor, with shape [N];
158 | indices of selected hard positive samples; 0 <= p_inds[i] <= N - 1
159 | n_inds: pytorch LongTensor, with shape [N];
160 | indices of selected hard negative samples; 0 <= n_inds[i] <= N - 1
161 | ==================
162 | For Debugging, etc
163 | ==================
164 | dist_ap: pytorch Variable, distance(anchor, positive); shape [N]
165 | dist_an: pytorch Variable, distance(anchor, negative); shape [N]
166 | dist_mat: pytorch Variable, pairwise euclidean distance; shape [N, N]
167 | """
168 | if self.normalize_feature:
169 | preds = self.normalize(preds, axis=-1)
170 | preds = preds.cuda(new_device)
171 | gts = gts.cuda(new_device)
172 | # shape [B, N, N]
173 | dist_mat = batch_pairwise_dist(preds, preds)
174 | dist_mat = dist_mat.clamp(min=1e-12).sqrt() # for numerical stability
175 | dist_ap, dist_an = self.hard_example_mining(dist_mat, gts)
176 | y = Variable(dist_an.data.new().resize_as_(dist_an.data).fill_(1))
177 | if self.margin is not None:
178 | loss = self.ranking_loss(dist_an, dist_ap, y)
179 | else:
180 | loss = self.ranking_loss(dist_an - dist_ap, y)
181 | return loss, dist_ap, dist_an
182 |
183 |
--------------------------------------------------------------------------------
/inference.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- coding: utf-8 -*-
3 | """
4 | @Author: An Tao
5 | @Contact: ta19@mails.tsinghua.edu.cn
6 | @File: inference.py
7 | @Time: 2020/1/2 10:26 AM
8 | """
9 |
10 | import os
11 | import sys
12 | import time
13 | import shutil
14 | import torch
15 | import numpy as np
16 | import h5py
17 |
18 | from tensorboardX import SummaryWriter
19 |
20 | from model import ReconstructionNet, ClassificationNet, SegmentationNet
21 | from dataset import Dataset
22 | from utils import Logger
23 |
24 |
25 | class Inference(object):
26 | def __init__(self, args):
27 | self.batch_size = args.batch_size
28 | self.no_cuda = args.no_cuda
29 | self.task = args.task
30 |
31 | # create exp directory
32 | file = [f for f in args.model_path.split('/')]
33 | if args.exp_name != None:
34 | self.experiment_id = args.exp_name
35 | else:
36 | self.experiment_id = time.strftime('%m%d%H%M%S')
37 | cache_root = 'cache/%s' % self.experiment_id
38 | os.makedirs(cache_root, exist_ok=True)
39 | self.feature_dir = os.path.join(cache_root, 'features/')
40 | sys.stdout = Logger(os.path.join(cache_root, 'log.txt'))
41 |
42 | # check directory
43 | if not os.path.exists(self.feature_dir):
44 | os.makedirs(self.feature_dir)
45 | else:
46 | shutil.rmtree(self.feature_dir)
47 | os.makedirs(self.feature_dir)
48 |
49 | # print args
50 | print(str(args))
51 |
52 | # get gpu id
53 | gids = ''.join(args.gpu.split())
54 | self.gpu_ids = [int(gid) for gid in gids.split(',')]
55 | self.first_gpu = self.gpu_ids[0]
56 |
57 | # generate dataset
58 | self.infer_dataset_train = Dataset(
59 | root=args.dataset_root,
60 | dataset_name=args.dataset,
61 | split='train',
62 | num_points=args.num_points,
63 | )
64 | self.infer_dataset_test = Dataset(
65 | root=args.dataset_root,
66 | dataset_name=args.dataset,
67 | split='test',
68 | num_points=args.num_points,
69 | )
70 | self.infer_loader_train = torch.utils.data.DataLoader(
71 | self.infer_dataset_train,
72 | batch_size=args.batch_size,
73 | shuffle=False,
74 | num_workers=args.workers
75 | )
76 | self.infer_loader_test = torch.utils.data.DataLoader(
77 | self.infer_dataset_test,
78 | batch_size=args.batch_size,
79 | shuffle=False,
80 | num_workers=args.workers
81 | )
82 | print("Inference set size (train):", self.infer_loader_train.dataset.__len__())
83 | print("Inference set size (test):", self.infer_loader_test.dataset.__len__())
84 |
85 | # initialize model
86 | if args.task == "reconstruct":
87 | self.model = ReconstructionNet(args)
88 | elif args.task == "classify":
89 | self.model = ClassificationNet(args)
90 | elif args.task == "segment":
91 | self.model = SegmentationNet(args)
92 | if args.model_path != '':
93 | self._load_pretrain(args.model_path)
94 |
95 | # load model to gpu
96 | if not args.no_cuda:
97 | if len(self.gpu_ids) != 1: # multiple gpus
98 | self.model = torch.nn.DataParallel(self.model.cuda(self.first_gpu), self.gpu_ids)
99 | else:
100 | self.model = self.model.cuda(self.gpu_ids[0])
101 |
102 | def run(self):
103 | self.model.eval()
104 |
105 | # generate train set for SVM
106 | loss_buf = []
107 | feature_train = []
108 | lbs_train = []
109 | n = 0
110 | for iter, (pts, lbs) in enumerate(self.infer_loader_train):
111 | if not self.no_cuda:
112 | pts = pts.cuda(self.first_gpu)
113 | lbs = lbs.cuda(self.first_gpu)
114 | if self.task == "reconstruct":
115 | output, feature = self.model(pts)
116 | elif self.task in ["classify", "segment"]:
117 | feature = self.model(pts)
118 | feature_train.append(feature.detach().cpu().numpy().squeeze(1))
119 | lbs_train.append(lbs.cpu().numpy().squeeze(1))
120 | if ((iter+1) * self.batch_size % 2048) == 0 \
121 | or (iter+1) == len(self.infer_loader_train):
122 | feature_train = np.concatenate(feature_train, axis=0)
123 | lbs_train = np.concatenate(lbs_train, axis=0)
124 | f = h5py.File(os.path.join(self.feature_dir, 'train' + str(n) + '.h5'),'w')
125 | f['data'] = feature_train
126 | f['label'] = lbs_train
127 | f.close()
128 | print("Train set {} for SVM saved.".format(n))
129 | feature_train = []
130 | lbs_train = []
131 | n += 1
132 | if self.task == "reconstruct":
133 | if len(self.gpu_ids) != 1: # multiple gpus
134 | loss = self.model.module.get_loss(pts, output)
135 | else:
136 | loss = self.model.get_loss(pts, output)
137 | loss_buf.append(loss.detach().cpu().numpy())
138 | if self.task == "reconstruct":
139 | print(f'Avg loss {np.mean(loss_buf)}')
140 | print("Finish generating train set for SVM.")
141 |
142 | # generate test set for SVM
143 | loss_buf = []
144 | feature_test = []
145 | lbs_test = []
146 | n = 0
147 | for iter, (pts, lbs) in enumerate(self.infer_loader_test):
148 | if not self.no_cuda:
149 | pts = pts.cuda(self.first_gpu)
150 | lbs = lbs.cuda(self.first_gpu)
151 | if self.task == "reconstruct":
152 | output, feature = self.model(pts)
153 | elif self.task in ["classify", "segment"]:
154 | feature = self.model(pts)
155 | feature_test.append(feature.detach().cpu().numpy().squeeze(1))
156 | lbs_test.append(lbs.cpu().numpy().squeeze(1))
157 | if ((iter+1) * self.batch_size % 2048) == 0 \
158 | or (iter+1) == len(self.infer_loader_test):
159 | feature_test = np.concatenate(feature_test, axis=0)
160 | lbs_test = np.concatenate(lbs_test, axis=0)
161 | f = h5py.File(os.path.join(self.feature_dir, 'test' + str(n) + '.h5'),'w')
162 | f['data'] = feature_test
163 | f['label'] = lbs_test
164 | f.close()
165 | print("Test set {} for SVM saved.".format(n))
166 | feature_test = []
167 | lbs_test = []
168 | n += 1
169 | if self.task == "reconstruct":
170 | if len(self.gpu_ids) != 1: # multiple gpus
171 | loss = self.model.module.get_loss(pts, output)
172 | else:
173 | loss = self.model.get_loss(pts, output)
174 | loss_buf.append(loss.detach().cpu().numpy())
175 | if self.task == "reconstruct":
176 | print(f'Avg loss {np.mean(loss_buf)}')
177 | print("Finish generating test set for SVM.")
178 |
179 | return self.feature_dir
180 |
181 |
182 | def _load_pretrain(self, pretrain):
183 | state_dict = torch.load(pretrain, map_location='cpu')
184 | from collections import OrderedDict
185 | new_state_dict = OrderedDict()
186 | for key, val in state_dict.items():
187 | if key[:6] == 'module':
188 | name = key[7:] # remove 'module.'
189 | else:
190 | name = key
191 | if key[:10] == 'classifier':
192 | continue
193 | if key[:9] == 'segmenter':
194 | continue
195 | new_state_dict[name] = val
196 | self.model.load_state_dict(new_state_dict)
197 | print(f"Load model from {pretrain}")
198 |
--------------------------------------------------------------------------------
/reconstruction.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- coding: utf-8 -*-
3 | """
4 | @Author: An Tao
5 | @Contact: ta19@mails.tsinghua.edu.cn
6 | @File: reconstruction.py
7 | @Time: 2020/1/2 10:26 AM
8 | """
9 |
10 | import os
11 | import sys
12 | import time
13 | import shutil
14 | import torch
15 | import torch.optim as optim
16 | import numpy as np
17 |
18 | from tensorboardX import SummaryWriter
19 |
20 | from model import ReconstructionNet
21 | from dataset import Dataset
22 | from utils import Logger
23 |
24 |
25 | class Reconstruction(object):
26 | def __init__(self, args):
27 | self.dataset_name = args.dataset
28 | if args.epochs != None:
29 | self.epochs = args.epochs
30 | elif args.encoder == 'foldnet':
31 | self.epochs = 278
32 | elif args.encoder == 'dgcnn_cls':
33 | self.epochs = 250
34 | elif args.encoder == 'dgcnn_seg':
35 | self.epochs = 290
36 | self.batch_size = args.batch_size
37 | self.snapshot_interval = args.snapshot_interval
38 | self.no_cuda = args.no_cuda
39 | self.model_path = args.model_path
40 |
41 | # create exp directory
42 | file = [f for f in args.model_path.split('/')]
43 | if args.exp_name != None:
44 | self.experiment_id = "Reconstruct_" + args.exp_name
45 | elif file[-2] == 'models':
46 | self.experiment_id = file[-3]
47 | else:
48 | self.experiment_id = "Reconstruct" + time.strftime('%m%d%H%M%S')
49 | snapshot_root = 'snapshot/%s' % self.experiment_id
50 | tensorboard_root = 'tensorboard/%s' % self.experiment_id
51 | self.save_dir = os.path.join(snapshot_root, 'models/')
52 | self.tboard_dir = tensorboard_root
53 |
54 | # check arguments
55 | if self.model_path == '':
56 | if not os.path.exists(self.save_dir):
57 | os.makedirs(self.save_dir)
58 | else:
59 | choose = input("Remove " + self.save_dir + " ? (y/n)")
60 | if choose == "y":
61 | shutil.rmtree(self.save_dir)
62 | os.makedirs(self.save_dir)
63 | else:
64 | sys.exit(0)
65 | if not os.path.exists(self.tboard_dir):
66 | os.makedirs(self.tboard_dir)
67 | else:
68 | shutil.rmtree(self.tboard_dir)
69 | os.makedirs(self.tboard_dir)
70 | sys.stdout = Logger(os.path.join(snapshot_root, 'log.txt'))
71 | self.writer = SummaryWriter(log_dir=self.tboard_dir)
72 |
73 | # print args
74 | print(str(args))
75 |
76 | # get gpu id
77 | gids = ''.join(args.gpu.split())
78 | self.gpu_ids = [int(gid) for gid in gids.split(',')]
79 | self.first_gpu = self.gpu_ids[0]
80 |
81 | # generate dataset
82 | self.train_dataset = Dataset(
83 | root=args.dataset_root,
84 | dataset_name=args.dataset,
85 | split='all',
86 | num_points=args.num_points,
87 | random_translate=args.use_translate,
88 | random_rotate=True,
89 | random_jitter=args.use_jitter
90 | )
91 | self.train_loader = torch.utils.data.DataLoader(
92 | self.train_dataset,
93 | batch_size=args.batch_size,
94 | shuffle=True,
95 | num_workers=args.workers
96 | )
97 | print("Training set size:", self.train_loader.dataset.__len__())
98 |
99 | # initialize model
100 | self.model = ReconstructionNet(args)
101 | if self.model_path != '':
102 | self._load_pretrain(args.model_path)
103 |
104 | # load model to gpu
105 | if not self.no_cuda:
106 | if len(self.gpu_ids) != 1: # multiple gpus
107 | self.model = torch.nn.DataParallel(self.model.cuda(self.first_gpu), self.gpu_ids)
108 | else:
109 | self.model = self.model.cuda(self.gpu_ids[0])
110 |
111 | # initialize optimizer
112 | self.parameter = self.model.parameters()
113 | self.optimizer = optim.Adam(self.parameter, lr=0.0001*16/args.batch_size, betas=(0.9, 0.999), weight_decay=1e-6)
114 |
115 |
116 | def run(self):
117 | self.train_hist = {
118 | 'loss': [],
119 | 'per_epoch_time': [],
120 | 'total_time': []
121 | }
122 | best_loss = 1000000000
123 | print('Training start!!')
124 | start_time = time.time()
125 | self.model.train()
126 | if self.model_path != '':
127 | start_epoch = self.model_path[-7:-4]
128 | if start_epoch[0] == '_':
129 | start_epoch = start_epoch[1:]
130 | start_epoch = int(start_epoch)
131 | else:
132 | start_epoch = 0
133 | for epoch in range(start_epoch, self.epochs):
134 | loss = self.train_epoch(epoch)
135 |
136 | # save snapeshot
137 | if (epoch + 1) % self.snapshot_interval == 0:
138 | self._snapshot(epoch + 1)
139 | if loss < best_loss:
140 | best_loss = loss
141 | self._snapshot('best')
142 |
143 | # save tensorboard
144 | if self.writer:
145 | self.writer.add_scalar('Train Loss', self.train_hist['loss'][-1], epoch)
146 | self.writer.add_scalar('Learning Rate', self._get_lr(), epoch)
147 |
148 | # finish all epoch
149 | self._snapshot(epoch + 1)
150 | if loss < best_loss:
151 | best_loss = loss
152 | self._snapshot('best')
153 | self.train_hist['total_time'].append(time.time() - start_time)
154 | print("Avg one epoch time: %.2f, total %d epochs time: %.2f" % (np.mean(self.train_hist['per_epoch_time']),
155 | self.epochs, self.train_hist['total_time'][0]))
156 | print("Training finish!... save training results")
157 |
158 |
159 | def train_epoch(self, epoch):
160 | epoch_start_time = time.time()
161 | loss_buf = []
162 | num_batch = int(len(self.train_loader.dataset) / self.batch_size)
163 | for iter, (pts, _) in enumerate(self.train_loader):
164 | if pts.size(0) == 1:
165 | continue
166 |
167 | if not self.no_cuda:
168 | pts = pts.cuda(self.first_gpu)
169 |
170 | # forward
171 | self.optimizer.zero_grad()
172 | output, _ = self.model(pts)
173 |
174 | # loss
175 | if len(self.gpu_ids) != 1: # multiple gpus
176 | loss = self.model.module.get_loss(pts, output)
177 | else:
178 | loss = self.model.get_loss(pts, output)
179 |
180 | # backward
181 | loss.backward()
182 | self.optimizer.step()
183 | loss_buf.append(loss.detach().cpu().numpy())
184 |
185 | # finish one epoch
186 | epoch_time = time.time() - epoch_start_time
187 | self.train_hist['per_epoch_time'].append(epoch_time)
188 | self.train_hist['loss'].append(np.mean(loss_buf))
189 | print(f'Epoch {epoch+1}: Loss {np.mean(loss_buf)}, time {epoch_time:.4f}s')
190 | return np.mean(loss_buf)
191 |
192 |
193 | def _snapshot(self, epoch):
194 | state_dict = self.model.state_dict()
195 | from collections import OrderedDict
196 | new_state_dict = OrderedDict()
197 | for key, val in state_dict.items():
198 | if key[:6] == 'module':
199 | name = key[7:] # remove 'module.'
200 | else:
201 | name = key
202 | new_state_dict[name] = val
203 | save_dir = os.path.join(self.save_dir, self.dataset_name)
204 | torch.save(new_state_dict, save_dir + "_" + str(epoch) + '.pkl')
205 | print(f"Save model to {save_dir}_{str(epoch)}.pkl")
206 |
207 |
208 | def _load_pretrain(self, pretrain):
209 | state_dict = torch.load(pretrain, map_location='cpu')
210 | from collections import OrderedDict
211 | new_state_dict = OrderedDict()
212 | for key, val in state_dict.items():
213 | if key[:6] == 'module':
214 | name = key[7:] # remove 'module.'
215 | else:
216 | name = key
217 | new_state_dict[name] = val
218 | self.model.load_state_dict(new_state_dict)
219 | print(f"Load model from {pretrain}")
220 |
221 |
222 | def _get_lr(self, group=0):
223 | return self.optimizer.param_groups[group]['lr']
224 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Point Cloud Segmentation for Classific Feature Learning
2 | ## Introduction
3 | This work aims to show whether learning a point cloud segmentation task is able to extract features performing well in classification. We do all experiments under the framework of DGCNN.
4 |
5 | Details for DGCNN see **Dynamic Graph CNN for Learning on Point Clouds** (https://arxiv.xilesou.top/pdf/1801.07829). We provide a PyTorch reimplementation for DGCNN in [antao97/dgcnn.pytorch](https://github.com/antao97/dgcnn.pytorch).
6 |
7 | DGCNN provides two type of networks, one for classification and one for segmentation. We use "DGCNN_Cls" to denote network for classification and "DGCNN_Seg" for segmentation. The network sturcture of DGCNN is
8 |
9 | 
10 |
11 |
12 | We also do experiments to see whether learning segmentation on meaningful point clouds assembled by some base point clouds can help to learning better features for base point clouds. To run the experments, we first seperate segmentation parts of each shape from ShapeNetPart dataset as new shapes and make this new dataset as ShapeNetPart Part dataset. Then, we train network on ShapeNetPart dataset as normal and test it on ShapeNetPart Part dataset.
13 |
14 | Some visualized point clouds in our ShapeNetPart Part dataset:
15 |
16 | 
17 | 
18 | 
19 |
20 | chair skateboard airplane
21 |
22 | Experimental results show that learning point cloud segmentation does help to extract features suitable for classification.
23 |
24 | **The key contributions of this work are as follows:**
25 |
26 | - Since the network provided by DGCNN for segmentation is supervised, we provide an revised DGCNN segmentation network with no category label.
27 | - When segmentation is trained and tested on intact point clouds, the trained model can help to extract better features.
28 | - When segmentation is trained on intact point clouds and tested on base point clouds, the trained model also can help to extract better features.
29 |
30 | If you find this work useful, please cite:
31 | ```
32 | @article{tao2020,
33 | Author = {An Tao},
34 | Title = {Point Cloud Segmentation for Classific Feature Learning},
35 | Journal = {https://github.com/antao97/PointCloudSegmentation},
36 | Year = {2020}
37 | }
38 | ```
39 |
40 |
41 | ## Requirements
42 | - Python 3.7
43 | - PyTorch 1.2
44 | - CUDA 10.0
45 | - Package: glob, h5py, tensorflow, tensorboard, tensorboardX and sklearn
46 |
47 |
48 | ## Download datasets
49 | Download the HDF5 format datasets (where each shape is sampled 2,048 points uniformly):
50 |
51 | - ShapeNetPart (338M) [[TsinghuaCloud]](https://cloud.tsinghua.edu.cn/f/c25d94e163454196a26b/) [[BaiduDisk]](https://pan.baidu.com/s/1yi4bMVBE2mV8NqVRtNLoqw)
52 | - ShapeNetPart Part (450M) [[TsinghuaCloud]](https://cloud.tsinghua.edu.cn/f/b6ee6e3b345744889e6b/) [[BaiduDisk]](https://pan.baidu.com/s/1uCotFveZ5R5ztRPIm1hfJQ)
53 | - ModelNet40 (194M) [[TsinghuaCloud]](https://cloud.tsinghua.edu.cn/f/b3d9fe3e2a514def8097/) [[BaiduDisk]](https://pan.baidu.com/s/1NQZgN8tvHVqQntxefcdVAg)
54 |
55 | You can find more details about the above datasets in this [repo](https://github.com/antao97/PointCloudDatasets).
56 |
57 |
58 | ## Experiment settings
59 | To evaluate the quality of extracted features, we use ShapeNetPart dataset to both train DGCNN and a linear SVM classifier. Specifically, we train the linear SVM classifier on ShapeNetPart dataset using the features (latent representations) obtained from the trained feature encoder.
60 |
61 | For transfer performance, we train the linear SVM classifier on ModelNet 40 dataset using the features (latent representations) obtained from the same network trained from the ShapeNetPart dataset.
62 |
63 | In this work we compare the performance for adopted training task among supervised segmentation, supervised segmentation without category label, supervised classification and unsupervised reconstruction. For supervised segmentation without category label, we discard the adding of categorical vector (mlp {64}) and directly repeat the 1024 dim feature into n x 1024. We do unsupervised reconstruction following the framework in this [repo](https://github.com/antao97/UnsupervisedPointCloudReconstruction) and use source points from sphere surface for decoder. We also change feature dimension into 1024. Except unsupervised reconstruction, we do all experiments under the framework of DGCNN.
64 |
65 | To train the network, run
66 | ```
67 | python main.py --exp_name --task --dataset_root --encoder --k <20 | 40> --feat_dims 1024 --batch_size <16 | 32> --dataset shapenetpart --gpu
68 | ```
69 | Use `--seg_no_label` if you want to run segmentation task without category label.
70 |
71 | You can download our already trained models from [[TsinghuaCloud]](https://cloud.tsinghua.edu.cn/d/d9e7a899582d432cbc11/) or [[BaiduDisk]](https://pan.baidu.com/s/1UigbY4jNts8LMZ6fqJXvxQ) and place them under `snapshot/`.
72 |
73 | Because this work is done before our PyTorch reimplementation for DGCNN, the training setting in this repo is slightly different with [antao97/dgcnn.pytorch](https://github.com/antao97/dgcnn.pytorch). You can set `self.epochs = 200` and `random_translate=False` in line 61 and 120 in `segmentation.py` to follow the settings in [antao97/dgcnn.pytorch](https://github.com/antao97/dgcnn.pytorch).
74 |
75 | Besides ShapeNetPart dataset, we also test the performace of linear SVM classifier on ShapeNetPart Part dataset, using the model trained on ShapeNetPart dataset with segmentation task.
76 |
77 | To evaluate the performance of a given trained model, run
78 | ```
79 | python main.py --eval --model_path --task --dataset_root --encoder --k <20 | 40> --feat_dims 1024 --dataset --gpu
80 | ```
81 |
82 | Use `--no_cuda` if you want to run in CPU.
83 |
84 | To use Tensorboard, run
85 | ```
86 | tensorboard --logdir tensorboard --bind_all
87 | ```
88 | You can find the Tensorboard records under `tensorboard/`.
89 |
90 |
91 | ## Classification accuracy of linear SVM classifier
92 | ### Results with best settings
93 | | Task | Info | Encoder | K | Batch Size | Epochs | ShapeNetPart | ModelNet40 |
94 | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
95 | | Segmentation | Supervised | DGCNN_Seg | 40 | 32 | 250 | 98.9% | 89.4% |
96 | | Segmentation | Supervised (no category label) | DGCNN_Seg | 40 | 32 | 250 | **99.9%** | 89.2% |
97 | | Classification | Supervised | DGCNN_Cls | 40 | 32 | 250 | 99.8% | 89.6% |
98 | | Reconstruction | Unsupervised | DGCNN_Cls | 20 | 16 | 250 | 98.7% | **89.8%** |
99 |
100 |
101 |
102 | 
103 | 
104 |
105 |
106 |
107 | ### Results with settings for segmentation task
108 | | Task | Info | Encoder | K | Batch Size | Epochs | ShapeNetPart | ModelNet40 |
109 | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
110 | | Segmentation | Supervised | DGCNN_Seg | 40 | 32 | 250 | 98.9% | **89.4%** |
111 | | Segmentation | Supervised (no category label) | DGCNN_Seg | 40 | 32 | 250 | **99.9%** | 89.2% |
112 | | Classification | Supervised | DGCNN_Seg | 40 | 32 | 250 | **99.9%** | 86.8% |
113 | | Reconstruction | Unsupervised | DGCNN_Seg | 40 | 32 | 290 | 98.8% | 89.2% |
114 |
115 |
116 |
117 | 
118 | 
119 |
120 |
121 |
122 | ### Results evaluated on ShapeNetPart Part dataset
123 | | Task | Info | Encoder | K | Batch Size | Epochs | Training Dataset | Eval Acc
124 | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
125 | | Segmentation | Supervised | DGCNN_Seg | 40 | 32 | 250 | ShapeNetPart | 85.0% |
126 | | Segmentation | Supervised (no category label) | DGCNN_Seg | 40 | 32 | 250 | ShapeNetPart | 84.0% |
127 | | Classification | Supervised | DGCNN_Cls | 40 | 32 | 250 | ShapeNetPart Part | **99.0%** | 4
128 | | Reconstruction | Unsupervised | DGCNN_Cls | 20 | 16 | 250 | ShapeNetPart Part | 87.5% |
129 |
130 |
131 |
132 | 
133 |
134 |
135 |
136 | ## Performance analysis
137 | Experimental results show that learning point cloud segmentation does help to extract features suitable for classification. However, simply adopting the training scheme from DGCNN for segmentation task is not suitable for transfer learning. We believe better results will get if using better training scheme.
138 |
139 |
--------------------------------------------------------------------------------
/classification.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- coding: utf-8 -*-
3 | """
4 | @Author: An Tao
5 | @Contact: ta19@mails.tsinghua.edu.cn
6 | @File: reconstruction.py
7 | @Time: 2020/1/2 10:26 AM
8 | """
9 |
10 | import os
11 | import sys
12 | import time
13 | import shutil
14 | import numpy as np
15 | import torch
16 | import torch.optim as optim
17 | from torch.optim.lr_scheduler import CosineAnnealingLR
18 | import sklearn.metrics as metrics
19 |
20 | from tensorboardX import SummaryWriter
21 |
22 | from model import ClassificationNet
23 | from dataset import Dataset
24 | from utils import Logger
25 |
26 |
27 | class Classification(object):
28 | def __init__(self, args):
29 | self.dataset_name = args.dataset
30 | if args.epochs != None:
31 | self.epochs = args.epochs
32 | else:
33 | self.epochs = 250
34 | self.batch_size = args.batch_size
35 | self.snapshot_interval = args.snapshot_interval
36 | self.no_cuda = args.no_cuda
37 | self.model_path = args.model_path
38 | self.no_scheduler = args.no_scheduler
39 |
40 | # create exp directory
41 | file = [f for f in args.model_path.split('/')]
42 | if args.exp_name != None:
43 | self.experiment_id = "Classify_" + args.exp_name
44 | elif file[-2] == 'models':
45 | self.experiment_id = file[-3]
46 | else:
47 | self.experiment_id = "Classify" + time.strftime('%m%d%H%M%S')
48 | snapshot_root = 'snapshot/%s' % self.experiment_id
49 | tensorboard_root = 'tensorboard/%s' % self.experiment_id
50 | self.save_dir = os.path.join(snapshot_root, 'models/')
51 | self.tboard_dir = tensorboard_root
52 |
53 | # check arguments
54 | if self.model_path == '':
55 | if not os.path.exists(self.save_dir):
56 | os.makedirs(self.save_dir)
57 | else:
58 | choose = input("Remove " + self.save_dir + " ? (y/n)")
59 | if choose == "y":
60 | shutil.rmtree(self.save_dir)
61 | os.makedirs(self.save_dir)
62 | else:
63 | sys.exit(0)
64 | if not os.path.exists(self.tboard_dir):
65 | os.makedirs(self.tboard_dir)
66 | else:
67 | shutil.rmtree(self.tboard_dir)
68 | os.makedirs(self.tboard_dir)
69 | sys.stdout = Logger(os.path.join(snapshot_root, 'log.txt'))
70 | self.writer = SummaryWriter(log_dir=self.tboard_dir)
71 |
72 | # print args
73 | print(str(args))
74 |
75 | # get gpu id
76 | gids = ''.join(args.gpu.split())
77 | self.gpu_ids = [int(gid) for gid in gids.split(',')]
78 | self.first_gpu = self.gpu_ids[0]
79 |
80 | # generate dataset
81 | self.train_dataset = Dataset(
82 | root=args.dataset_root,
83 | dataset_name=args.dataset,
84 | split='all',
85 | num_points=args.num_points,
86 | random_translate=True,
87 | random_rotate=args.use_rotate,
88 | random_jitter=args.use_jitter
89 | )
90 | self.train_loader = torch.utils.data.DataLoader(
91 | self.train_dataset,
92 | batch_size=args.batch_size,
93 | shuffle=True,
94 | num_workers=args.workers
95 | )
96 | print("Training set size:", self.train_loader.dataset.__len__())
97 |
98 | # initialize model
99 | self.model = ClassificationNet(args)
100 | if self.model_path != '':
101 | self._load_pretrain(args.model_path)
102 |
103 | # load model to gpu
104 | if not self.no_cuda:
105 | if len(self.gpu_ids) != 1: # multiple gpus
106 | self.model = torch.nn.DataParallel(self.model.cuda(self.first_gpu), self.gpu_ids)
107 | else:
108 | self.model = self.model.cuda(self.gpu_ids[0])
109 |
110 | # initialize optimizer
111 | self.parameter = self.model.parameters()
112 | if self.no_scheduler == False:
113 | self.optimizer = optim.SGD(self.parameter, lr=0.1, weight_decay=1e-4)
114 | self.scheduler = CosineAnnealingLR(self.optimizer, self.epochs, eta_min=1e-3)
115 | else:
116 | self.optimizer = optim.SGD(self.parameter, lr=0.01, weight_decay=1e-4)
117 |
118 |
119 | def run(self):
120 | self.train_hist = {
121 | 'loss': [],
122 | 'per_epoch_time': [],
123 | 'total_time': []
124 | }
125 | best_loss = 1000000000
126 | print('Training start!!')
127 | start_time = time.time()
128 | self.model.train()
129 | if self.model_path != '':
130 | start_epoch = self.model_path[-7:-4]
131 | if start_epoch[0] == '_':
132 | start_epoch = start_epoch[1:]
133 | start_epoch = int(start_epoch)
134 | else:
135 | start_epoch = 0
136 | for epoch in range(start_epoch, self.epochs):
137 | loss = self.train_epoch(epoch)
138 |
139 | # save snapeshot
140 | if (epoch + 1) % self.snapshot_interval == 0:
141 | self._snapshot(epoch + 1)
142 | if loss < best_loss:
143 | best_loss = loss
144 | self._snapshot('best')
145 |
146 | # save tensorboard
147 | if self.writer:
148 | self.writer.add_scalar('Train Loss', self.train_hist['loss'][-1], epoch)
149 | self.writer.add_scalar('Learning Rate', self._get_lr(), epoch)
150 |
151 | # finish all epoch
152 | self._snapshot(epoch + 1)
153 | if loss < best_loss:
154 | best_loss = loss
155 | self._snapshot('best')
156 | self.train_hist['total_time'].append(time.time() - start_time)
157 | print("Avg one epoch time: %.2f, total %d epochs time: %.2f" % (np.mean(self.train_hist['per_epoch_time']),
158 | self.epochs, self.train_hist['total_time'][0]))
159 | print("Training finish!... save training results")
160 |
161 |
162 | def train_epoch(self, epoch):
163 | epoch_start_time = time.time()
164 | loss_buf = []
165 | train_pred = []
166 | train_true = []
167 | num_batch = int(len(self.train_loader.dataset) / self.batch_size)
168 | for iter, (pts, label) in enumerate(self.train_loader):
169 | if pts.size(0) == 1:
170 | continue
171 | if not self.no_cuda:
172 | pts = pts.cuda(self.first_gpu)
173 | label = label.cuda(self.first_gpu)
174 |
175 | # forward
176 | self.optimizer.zero_grad()
177 | output, _ = self.model(pts)
178 |
179 | # loss
180 | if len(self.gpu_ids) != 1: # multiple gpus
181 | loss = self.model.module.get_loss(output, label)
182 | else:
183 | loss = self.model.get_loss(output, label)
184 |
185 | # backward
186 | loss.backward()
187 | self.optimizer.step()
188 | loss_buf.append(loss.detach().cpu().numpy())
189 |
190 | preds = output.max(dim=1)[1]
191 | train_true.append(label.view(-1).cpu().numpy())
192 | train_pred.append(preds.detach().cpu().numpy())
193 |
194 | # finish one epoch
195 | if self.no_scheduler == False:
196 | self.scheduler.step()
197 | epoch_time = time.time() - epoch_start_time
198 | self.train_hist['per_epoch_time'].append(epoch_time)
199 | self.train_hist['loss'].append(np.mean(loss_buf))
200 | train_true = np.concatenate(train_true)
201 | train_pred = np.concatenate(train_pred)
202 | print("Epoch %d: Loss %.6f, train acc %.6f, train avg acc %.6f, time %.4fs" % (epoch+1,
203 | np.mean(loss_buf),
204 | metrics.accuracy_score(
205 | train_true, train_pred),
206 | metrics.balanced_accuracy_score(
207 | train_true, train_pred),
208 | epoch_time))
209 | return np.mean(loss_buf)
210 |
211 |
212 | def _snapshot(self, epoch):
213 | state_dict = self.model.state_dict()
214 | from collections import OrderedDict
215 | new_state_dict = OrderedDict()
216 | for key, val in state_dict.items():
217 | if key[:6] == 'module':
218 | name = key[7:] # remove 'module.'
219 | else:
220 | name = key
221 | new_state_dict[name] = val
222 | save_dir = os.path.join(self.save_dir, self.dataset_name)
223 | torch.save(new_state_dict, save_dir + "_" + str(epoch) + '.pkl')
224 | print(f"Save model to {save_dir}_{str(epoch)}.pkl")
225 |
226 |
227 | def _load_pretrain(self, pretrain):
228 | state_dict = torch.load(pretrain, map_location='cpu')
229 | from collections import OrderedDict
230 | new_state_dict = OrderedDict()
231 | for key, val in state_dict.items():
232 | if key[:6] == 'module':
233 | name = key[7:] # remove 'module.'
234 | else:
235 | name = key
236 | new_state_dict[name] = val
237 | self.model.load_state_dict(new_state_dict)
238 | print(f"Load model from {pretrain}")
239 |
240 |
241 | def _get_lr(self, group=0):
242 | return self.optimizer.param_groups[group]['lr']
243 |
--------------------------------------------------------------------------------
/visualization.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- coding: utf-8 -*-
3 | """
4 | @Author: An Tao
5 | @Contact: ta19@mails.tsinghua.edu.cn
6 | @File: visualization.py
7 | @Time: 2020/1/2 10:26 AM
8 | """
9 |
10 | import os
11 | import time
12 | import numpy as np
13 | import torch
14 | import itertools
15 | import argparse
16 | from glob import glob
17 |
18 | from model import ReconstructionNet
19 |
20 | def standardize_bbox(pcl, points_per_object):
21 | pt_indices = np.random.choice(pcl.shape[0], points_per_object, replace=False)
22 | np.random.shuffle(pt_indices)
23 | pcl = pcl[pt_indices] # n by 3
24 | mins = np.amin(pcl, axis=0)
25 | maxs = np.amax(pcl, axis=0)
26 | center = ( mins + maxs ) / 2.
27 | scale = np.amax(maxs-mins)
28 | print("Center: {}, Scale: {}".format(center, scale))
29 | result = ((pcl - center)/scale).astype(np.float32) # [-0.5, 0.5]
30 | return result
31 |
32 | xml_head = \
33 | """
34 |
35 |
36 |
37 |
38 |
39 |
40 |
41 |
42 |
43 |
44 |
45 |
46 |
47 |
48 |
49 |
50 |
51 |
52 |
53 |
54 |
55 |
56 |
57 |
58 |
59 |
60 |
61 |
62 |
63 |
64 | """
65 |
66 | xml_ball_segment = \
67 | """
68 |
69 |
70 |
71 |
72 |
73 |
74 |
75 |
76 |
77 |
78 | """
79 |
80 | xml_tail = \
81 | """
82 |
83 |