点云深度学习——点云配准网络DCP复现

---

前言

最近在学习点云深度学习，主要看了一下点云的配准网络，比如PointNetLK , DCP, DeepICP 等，打算先跑通代码，在进行原理的细致学习。下面记录一下DCP网络的实现。

---

一、效果展示

1.1 open3d中效果展示

下图为open3d测试训练好的模型，测试集为modelnet40中的样本

---

二、复现源码

2.1 参考链接

Github源码参考链接

2.2 复现流程

1）配置好pytorch 、cuda等相关环境
可以参照网上教程
2）依赖项配置

scipy安装

conda install scipy

h5py安装

conda install h5py

tqdm安装

conda install tqdm

tensorBoardX安装

conda install tensorboard
conda install tensorboardx

3）训练模型

参考readme training

4）测试模型

2.3遇到问题：

1）缺少nose模块

conda install nose
1

2)from-dcm 报错

错误原因：

from_dcm 替换为from_matrix

---

三、模型测试单个数据，并用open3d显示

3.1 单个数据测试代码

import numpy as np
import torch
import time
import os
from model import DCP
from util import transform_point_cloud, npmat2euler
import argparse
from scipy.spatial.transform import Rotation
from data import ModelNet40
import glob
import h5py
import open3d as o3d

def transform_input(pointcloud):
    """
    random rotation and transformation the input
    pointcloud: N*3
    """

    anglex = np.random.uniform() * np.pi / 4
    angley = np.random.uniform() * np.pi / 4
    anglez = np.random.uniform() * np.pi / 4
    
    
    # anglex = 0.04
    # angley = 0.04
    # anglez = 0.04

    print('angle: ',anglex,angley,anglez)
    
    cosx = np.cos(anglex)
    cosy = np.cos(angley)
    cosz = np.cos(anglez)
    sinx = np.sin(anglex)
    siny = np.sin(angley)
    sinz = np.sin(anglez)
    Rx = np.array([[1, 0, 0],
                   [0, cosx, -sinx],
                   [0, sinx, cosx]])
    Ry = np.array([[cosy, 0, siny],
                   [0, 1, 0],
                   [-siny, 0, cosy]])
    Rz = np.array([[cosz, -sinz, 0],
                   [sinz, cosz, 0],
                   [0, 0, 1]])
    R_ab = Rx.dot(Ry).dot(Rz)
    R_ba = R_ab.T
    translation_ab = np.array([np.random.uniform(-0.5, 0.5),
                               np.random.uniform(-0.5, 0.5),
                               np.random.uniform(-0.5, 0.5)])

    # translation_ab = np.array([0.01,0.05,0.05])
    print('trans: ',translation_ab)
    
    
    translation_ba = -R_ba.dot(translation_ab)

    pointcloud1 = pointcloud[:,:3].T

    rotation_ab = Rotation.from_euler('zyx', [anglez, angley, anglex])
    pointcloud2 = rotation_ab.apply(pointcloud1.T).T + np.expand_dims(translation_ab, axis=1)

    euler_ab = np.asarray([anglez, angley, anglex])
    euler_ba = -euler_ab[::-1]
    rotation_ba = Rotation.from_euler('zyx', euler_ba)

    pointcloud1 = np.random.permutation(pointcloud1.T)
    pointcloud2 = np.random.permutation(pointcloud2.T)

    return pointcloud1.astype('float32'), pointcloud2.astype('float32'), \
           rotation_ab,translation_ab, rotation_ba,translation_ba

def run_one_pointcloud(src,target,net):
    
    if len(src.shape)==2 and len(target.shape)==2: ##  (N,3)
    
        print("src/target shape:", src.shape,target.shape)
        
        src = np.expand_dims(src[:,:3],axis=0)
        src = np.transpose(src,[0,2,1])  ##  (1, 3, 1024)
        target = np.expand_dims(target[:,:3],axis=0)
        target = np.transpose(target,[0,2,1])  ##  (1, 3, 1024)
    
    net.eval()
 
    src = torch.from_numpy(src).cuda()
    target = torch.from_numpy(target).cuda()
    
    rotation_ab_pred, translation_ab_pred, \
    rotation_ba_pred, translation_ba_pred = net(src, target)
  
    target_pred = transform_point_cloud(src, rotation_ab_pred,
                                            translation_ab_pred)
    
    src_pred = transform_point_cloud(target, rotation_ba_pred,
                                               translation_ba_pred)

    # put on cpu and turn into numpy
    src_pred = src_pred.detach().cpu().numpy()
    src_pred = np.transpose(src_pred[0],[1,0])

    target_pred = target_pred.detach().cpu().numpy()
    target_pred = np.transpose(target_pred[0],[1,0])

    rotation_ab_pred = rotation_ab_pred.detach().cpu().numpy()
    translation_ab_pred = translation_ab_pred.detach().cpu().numpy()

    rotation_ba_pred = rotation_ba_pred.detach().cpu().numpy()
    translation_ba_pred = translation_ba_pred.detach().cpu().numpy()
    
    return src_pred,target_pred,rotation_ab_pred, translation_ab_pred,rotation_ba_pred, translation_ba_pred
    
if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='Point Cloud Registration')
    parser.add_argument('--exp_name', type=str, default='', metavar='N',
                        help='Name of the experiment')
    parser.add_argument('--model', type=str, default='dcp', metavar='N',
                        choices=['dcp'],
                        help='Model to use, [dcp]')
    parser.add_argument('--emb_nn', type=str, default='dgcnn', metavar='N',
                        choices=['pointnet', 'dgcnn'],
                        help='Embedding nn to use, [pointnet, dgcnn]')
    parser.add_argument('--pointer', type=str, default='transformer', metavar='N',
                        choices=['identity', 'transformer'],
                        help='Attention-based pointer generator to use, [identity, transformer]')
    parser.add_argument('--head', type=str, default='svd', metavar='N',
                        choices=['mlp', 'svd', ],
                        help='Head to use, [mlp, svd]')
    parser.add_argument('--emb_dims', type=int, default=512, metavar='N',
                        help='Dimension of embeddings')
    parser.add_argument('--n_blocks', type=int, default=1, metavar='N',
                        help='Num of blocks of encoder&decoder')
    parser.add_argument('--n_heads', type=int, default=4, metavar='N',
                        help='Num of heads in multiheadedattention')
    parser.add_argument('--ff_dims', type=int, default=1024, metavar='N',
                        help='Num of dimensions of fc in transformer')
    parser.add_argument('--dropout', type=float, default=0.0, metavar='N',
                        help='Dropout ratio in transformer')
    parser.add_argument('--batch_size', type=int, default=32, metavar='batch_size',
                        help='Size of batch)')
    parser.add_argument('--test_batch_size', type=int, default=1, metavar='batch_size',
                        help='Size of batch)')
    parser.add_argument('--epochs', type=int, default=250, metavar='N',
                        help='number of episode to train ')
    parser.add_argument('--use_sgd', action='store_true', default=False,
                        help='Use SGD')
    parser.add_argument('--lr', type=float, default=0.001, metavar='LR',
                        help='learning rate (default: 0.001, 0.1 if using sgd)')
    parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
                        help='SGD momentum (default: 0.9)')
    parser.add_argument('--no_cuda', action='store_true', default=False,
                        help='enables CUDA training')
    parser.add_argument('--seed', type=int, default=1234, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--eval', action='store_true', default=False,
                        help='evaluate the model')
    parser.add_argument('--cycle', type=bool, default=False, metavar='N',
                        help='Whether to use cycle consistency')
    parser.add_argument('--model_path', type=str,
                        default= 'pretrained/dcp_v2.t7',
                        metavar='N',
                        help='Pretrained model path')
    args = parser.parse_args()
    torch.backends.cudnn.deterministic = True
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)


    # net prepared
    net = DCP(args).cuda()
    net.load_state_dict(torch.load( args.model_path), strict=False)


    f = h5py.File('data/modelnet40_ply_hdf5_2048/ply_data_train2.h5','r')
    data = f['data'][:].astype('float32')  # (2048, 2048, 3) 
    f.close()

    # index = np.random.randint(data.shape[0])
    index=0

    point1 = data[index,:,:]
    _,point2,_,_,_,_ = transform_input(point1)


    # src1=o3d.io.read_point_cloud("/home/pride/3d_registration/dcp-master/0_modelnet_src.ply")
    # point1=np.asarray(src1.points)
    # print(point1)
    # _, point2, _, _, _, _ = transform_input(point1)

    src,target = point1,point2


    ## run
    src_pred, target_pred,r_ab,t_ab,r_ba,t_ba, = run_one_pointcloud(src, target,net)

    print("#############  src -> target :\n", r_ab, t_ab)
    print("#############  src <- target :\n", r_ba, t_ba)

    #np->open3d
    src_cloud=o3d.geometry.PointCloud()
    src_cloud.points=o3d.utility.Vector3dVector(point1)
    tgt_cloud = o3d.geometry.PointCloud()
    tgt_cloud.points = o3d.utility.Vector3dVector(point2)
    trans_cloud = o3d.geometry.PointCloud()
    trans_cloud.points = o3d.utility.Vector3dVector(src_pred)

    # view
    src_cloud.paint_uniform_color([1,0,0])
    tgt_cloud.paint_uniform_color([0, 1, 0])
    trans_cloud.paint_uniform_color([0, 0, 1])
    o3d.visualization.draw_geometries([src_cloud,tgt_cloud,trans_cloud],width=800)

     

        
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
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
207
208
209
210
211
212
213
214
215
216

3.2 问题

1）用的都是modelnet的数据集，没有用过自己数据集测试

2）读取的是h5模型，后续会更改到利用open3d读取模型

---

over!!!