Pytorch1.7复现PointNet++点云分割（含Open3D可视化）（文末有一个本身做的

  毕设必要，复现一下PointNet++的对象分类、零件分割和场景分割，找点灵感和思绪，做个踩坑记录。
  下载代码

https://github.com/yanx27/Pointnet_Pointnet2_pytorch
  我的运行情况是pytorch1.7+cuda11.0。
训练

  PointNet++代码能实现3D对象分类、对象零件分割和语义场景分割。
对象分类

  下载数据集ModelNet40，并存储在文件夹data/modelnet40_normal_resampled/。

1. ## e.g., pointnet2_ssg without normal features
2. python train_classification.py --model pointnet2_cls_ssg --log_dir pointnet2_cls_ssg
3. python test_classification.py --log_dir pointnet2_cls_ssg
5. ## e.g., pointnet2_ssg with normal features
6. python train_classification.py --model pointnet2_cls_ssg --use_normals --log_dir pointnet2_cls_ssg_normal
7. python test_classification.py --use_normals --log_dir pointnet2_cls_ssg_normal
9. ## e.g., pointnet2_ssg with uniform sampling
10. python train_classification.py --model pointnet2_cls_ssg --use_uniform_sample --log_dir pointnet2_cls_ssg_fps
11. python test_classification.py --use_uniform_sample --log_dir pointnet2_cls_ssg_fps

复制代码

• 主文件夹下运行代码python train_classification.py --model pointnet2_cls_ssg --log_dir pointnet2_cls_ssg时大概会报错：
  ImportError: cannot import name 'PointNetSetAbstraction'
  缘故原由是pointnet2_cls_ssg.py文件import时的工作目录时models文件夹，但是实际运行的工作目录时models的上级目录，因此必要在pointnet2_cls_ssg.py里把from pointnet2_utils import PointNetSetAbstraction改成from models.pointnet2_utils import PointNetSetAbstraction。
  

  参考README.md文件，分类不是我的主攻点，这里就略过了。
零件分割

  零件分割是将一个物体的各个零件分割出来，好比把椅子的椅子腿分出来。
  下载数据集ShapeNet，并存储在文件夹data/shapenetcore_partanno_segmentation_benchmark_v0_normal/。
  运行也很简单：

1. ## e.g., pointnet2_msg
2. python train_partseg.py --model pointnet2_part_seg_msg --normal --log_dir pointnet2_part_seg_msg
3. python test_partseg.py --normal --log_dir pointnet2_part_seg_msg

复制代码

  shapenet数据集txt文件格式：前三个点是xyz,点云的位置坐标，后三个点是点云的法向信息，末了一个点是这个点所属的小种别，即1表示所属50个小种别中的第一个。
  写个代码用open3d可视化shapenet数据集的txt文件（随机配色）：

1. import open3d as o3d
2. import numpy as np
3. '''
4. BASE_DIR = os.path.dirname(os.path.abspath(__file__))
5. ROOT_DIR = os.path.dirname(BASE_DIR)
6. sys.path.append(BASE_DIR)
7. sys.path.append(os.path.join(ROOT_DIR, 'data_utils'))
8. '''
10. txt_path = '/home/lin/CV_AI_learning/Pointnet_Pointnet2_pytorch-master/data/shapenetcore_partanno_segmentation_benchmark_v0_normal/02691156/1b3c6b2fbcf834cf62b600da24e0965.txt'
11. # 通过numpy读取txt点云
12. pcd = np.genfromtxt(txt_path, delimiter=" ")
14. pcd_vector = o3d.geometry.PointCloud()
15. # 加载点坐标
16. # txt点云前三个数值一般对应x、y、z坐标，可以通过open3d.geometry.PointCloud().points加载
17. # 如果有法线或颜色，那么可以分别通过open3d.geometry.PointCloud().normals或open3d.geometry.PointCloud().colors加载
18. pcd_vector.points = o3d.utility.Vector3dVector(pcd[:, :3])
19. pcd_vector.colors = o3d.utility.Vector3dVector(pcd[:, 3:6])
20. o3d.visualization.draw_geometries([pcd_vector])

复制代码

  GPU内存不敷减小一下batch_size。
  我这里训练了一下，接着代码的best_model.pth继续训练150轮，RTX3080单显卡训练一轮得六七分钟，150轮花了半天多的时间。
  网上的代码根本test一下分割的一些参数就竣事了，没有做可视化，参考这篇blog做了一下效果的可视化：PointNet++分割推测效果可视化。这篇blog首先用网络将输入图像的推测效果存为txt文件，然后用Matplotlib做可视化，过程有点复杂了，用open3d做可视化比力简便一点，代码如下：

1. import tqdm
2. import matplotlib
3. import torch
4. import os
5. import warnings
6. import numpy as np
7. import open3d as o3d
8. from torch.utils.data import Dataset
9. import pybullet as p
10. from models.pointnet2_part_seg_msg import get_model as pointnet2
11. import time
13. warnings.filterwarnings('ignore')
14. matplotlib.use("Agg")
15. def pc_normalize(pc):
16.     centroid = np.mean(pc, axis=0)
17.     pc = pc - centroid
18.     m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
19.     pc = pc / m
20.     return pc,centroid,m
22. def generate_pointcloud(color_image, depth_image,width=1280,height=720,fov=50,near=0.01,far=5):
23.     rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(color_image, depth_image,convert_rgb_to_intensity=False)
24.     intrinsic = o3d.camera.PinholeCameraIntrinsic(o3d.camera.PinholeCameraIntrinsicParameters.Kinect2DepthCameraDefault )
26.     aspect = width / height
28.     projection_matrix = p.computeProjectionMatrixFOV(fov, aspect, near, far)
29.     intrinsic.set_intrinsics(width=width, height=height, fx=projection_matrix[0]*width/2, fy=projection_matrix[5]*height/2, cx=width/2, cy=height/2)
30.     point_cloud = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, intrinsic)
31.    
32.     point_cloud.estimate_normals( search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.01, max_nn=30))
33.     return point_cloud
35. class PartNormalDataset(Dataset):
36.     def __init__(self, point_cloud, npoints=2500, normal_channel=False):
37.         self.npoints = npoints # 采样点数
38.         self.cat = {}
39.         self.normal_channel = normal_channel # 是否使用法向信息
41.         position_data = np.asarray(point_cloud.points)
42.         normal_data = np.asarray(point_cloud.normals)
43.         self.raw_pcd = np.hstack([position_data,normal_data]).astype(np.float32)
45.         self.cat = {'board':'12345678'}
46.         # 输出的是元组，('Airplane',123.txt)
48.         self.classes = {'board': 0}
50.         data = self.raw_pcd
52.         if not self.normal_channel:  # 判断是否使用法向信息
53.             self.point_set = data[:, 0:3]
54.         else:
55.             self.point_set = data[:, 0:6]
57.         self.point_set[:, 0:3],self.centroid,self.m = pc_normalize(self.point_set[:, 0:3]) # 做一个归一化
59.         choice = np.random.choice(self.point_set.shape[0], self.npoints, replace=True) # 对一个类别中的数据进行随机采样 返回索引，允许重复采样
60.         # resample
61.         self.point_set =  self.point_set[choice, :] # 根据索引采样
63.     def __getitem__(self, index):
65.         cat = list(self.cat.keys())[0]
66.         cls = self.classes[cat] # 将类名转换为索引
67.         cls = np.array([cls]).astype(np.int32)
69.         return self.point_set, cls, self.centroid, self.m # pointset是点云数据，cls十六个大类别，seg是一个数据中，不同点对应的小类别
71.     def __len__(self):
72.         return 1
76. class Generate_txt_and_3d_img:
77.     def __init__(self,num_classes,testDataLoader,model,visualize = False):
78.         self.testDataLoader = testDataLoader
79.         self.num_classes = num_classes
80.         self.heat_map = False # 控制是否输出heatmap
81.         self.visualize = visualize # 是否open3d可视化
82.         self.model = model
84.         self.generate_predict()
85.         self.o3d_draw_3d_img()
87.     def __getitem__(self, index):
88.         return self.predict_pcd_colored
90.     def generate_predict(self):
92.         for _, (points, label,centroid,m) in tqdm.tqdm(enumerate(self.testDataLoader),
93.                                                                       total=len(self.testDataLoader),smoothing=0.9):
95.             #点云数据、整个图像的标签、每个点的标签、  没有归一化的点云数据（带标签）torch.Size([1, 7, 2048])
96.             points = points.transpose(2, 1)
97.             #print('1',target.shape) # 1 torch.Size([1, 2048])
98.             xyz_feature_point = points[:, :6, :]
100.             model = self.model
102.             seg_pred, _ = model(points, self.to_categorical(label, 1))
103.             seg_pred = seg_pred.cpu().data.numpy()
105.             if self.heat_map:
106.                 out =  np.asarray(np.sum(seg_pred,axis=2))
107.                 seg_pred = ((out - np.min(out) / (np.max(out) - np.min(out))))
108.             else:
109.                 seg_pred = np.argmax(seg_pred, axis=-1)  # 获得网络的预测结果 b n c
111.             seg_pred = np.concatenate([np.asarray(xyz_feature_point), seg_pred[:, None, :]],
112.                     axis=1).transpose((0, 2, 1)).squeeze(0)
114.             self.predict_pcd = seg_pred
115.             self.centroid = centroid
116.             self.m = m
119.     def o3d_draw_3d_img(self):
121.         pcd = self.predict_pcd
122.         pcd_vector = o3d.geometry.PointCloud()
123.         # 加载点坐标
124.         pcd_vector.points = o3d.utility.Vector3dVector(self.m * pcd[:, :3] + self.centroid)
125.         # colors = np.random.randint(255, size=(2,3))/255
126.         colors = np.array([[0.8, 0.8, 0.8],[1,0,0]])
127.         pcd_vector.colors = o3d.utility.Vector3dVector(colors[list(map(int,pcd[:, 6])),:])
129.         if self.visualize:
130.             coord_mesh = o3d.geometry.TriangleMesh.create_coordinate_frame(size = 0.1, origin = [0,0,0])
131.             o3d.visualization.draw_geometries([pcd_vector,coord_mesh])
132.         self.predict_pcd_colored = pcd_vector
134.     def to_categorical(self,y, num_classes):
135.         """ 1-hot encodes a tensor """
136.         new_y = torch.eye(num_classes)[y.cpu().data.numpy(),]
137.         if (y.is_cuda):
138.             return new_y.cuda()
139.         return new_y
141. def load_models(model_dict={'PonintNet': [pointnet2(num_classes=2,normal_channel=True).eval(),r'./log/part_seg/pointnet2_part_seg_msg/checkpoints']}):
142.     model = list(model_dict.values())[0][0]
143.     checkpoints_dir = list(model_dict.values())[0][1]
144.     weight_dict = torch.load(os.path.join(checkpoints_dir,'best_model.pth'))
145.     model.load_state_dict(weight_dict['model_state_dict'])
146.     return model
148. class Open3dVisualizer():
150.         def __init__(self):
152.                 self.point_cloud = o3d.geometry.PointCloud()
153.                 self.o3d_started = False
155.                 self.vis = o3d.visualization.VisualizerWithKeyCallback()
156.                 self.vis.create_window()
158.         def __call__(self, points, colors):
160.                 self.update(points, colors)
162.                 return False
164.         def update(self, points, colors):
165.                 coord_mesh = o3d.geometry.TriangleMesh.create_coordinate_frame(size = 0.15, origin = [0,0,0])
166.                 self.point_cloud.points = points
167.                 self.point_cloud.colors = colors
168.                 # self.point_cloud.transform([[1,0,0,0],[0,-1,0,0],[0,0,-1,0],[0,0,0,1]])
169.                 # self.vis.clear_geometries()
170.                 # Add geometries if it is the first time
171.                 if not self.o3d_started:
172.                         self.vis.add_geometry(self.point_cloud)
173.                         self.vis.add_geometry(coord_mesh)
174.                         self.o3d_started = True
176.                 else:
177.                         self.vis.update_geometry(self.point_cloud)
178.                         self.vis.update_geometry(coord_mesh)
180.                 self.vis.poll_events()
181.                 self.vis.update_renderer()
183. if __name__ =='__main__':
184.    
185.     num_classes = 2 # 填写数据集的类别数 如果是s3dis这里就填13   shapenet这里就填50
186.    
187.     color_image = o3d.io.read_image('image/rgb1.jpg')
188.     depth_image = o3d.io.read_image('image/depth1.png')
189.    
190.     point_cloud = generate_pointcloud(color_image=color_image, depth_image=depth_image)
192.     TEST_DATASET = PartNormalDataset(point_cloud,npoints=30000, normal_channel=True)
193.     testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=1, shuffle=False, num_workers=0,drop_last=True)
194.     predict_pcd = Generate_txt_and_3d_img(num_classes,testDataLoader,load_models(),visualize = True)

复制代码

  把之前的代码改成了针对单个点云推测的可视化，点云由GRB图像和深度图像天生，如果想直接输入点云本身稍微改下代码就可以了，目前仅针对shapenet数据集格式的数据。这里要留意如果训练的时间选择了--normal，那么normal_channel要改成True。
  看下训练效果，用modelnet40里的一个chair文件举行推测。
               可以看到这个椅子大抵是分成了四块，但是椅子靠背、腿分割地挺好的，就是扶手有一部门分割到了坐垫那里了，究竟训练时间不长。modelnet40数据集只是用来分类，并没有分割的标注，以是这里可视化了一下shapenet里标注好的椅子点云，看看椅子各个部位的分割（并非上面的椅子）。
           这里就比力显而易见地看出椅子分为靠背、扶手、坐垫、腿四个部门。
  开端观察到效果以后可以开始尝试本身制作数据集举行训练了，可以参考我写的这篇文章：《CloudCompare制作ShapeNet格式点云数据集》。
场景分割

  零件分割网络可以很容易地扩展到语义场景分割，点标志成为语义对象类而不是目标零件标志。
  在 Stanford 3D语义分析数据集上举行实行。该数据集包罗来自Matterport扫描仪的6个区域的3D扫描，包罗271个房间。扫描中的每个点都用13个种别(椅子、桌子、地板、墙壁等加上杂物)中的一个语义标签举行表明。
  先把文件下载过来： S3DIS ，存到文件夹data/s3dis/Stanford3dDataset_v1.2_Aligned_Version/.
  处理惩罚数据，数据会存到data/stanford_indoor3d/。

1. cd data_utils
2. python collect_indoor3d_data.py

复制代码

  运行：

1. ## Check model in ./models
2. ## e.g., pointnet2_ssg
3. python train_semseg.py --model pointnet2_sem_seg --test_area 5 --log_dir pointnet2_sem_seg
4. python test_semseg.py --log_dir pointnet2_sem_seg --test_area 5 --visual

复制代码

  上面的利用走完以后会在log/sem_seg/pointnet2_sem_seg/visual/天生推测效果的obj文件，可以用open3d举行可视化，就是不能用o3d.io.read_triangle_mesh函数来可视化obj文件，因为这里天生的obj文件还带了颜色信息用来表示语义信息，以是得读取成列表数据然后界说成o3d.geometry.PointCloud()变量表现，代码如下：

1. import copy
2. import numpy as np
3. import open3d as o3d
4. import os
6. objFilePath = 'log/sem_seg/pointnet2_sem_seg/visual/Area_5_office_8_gt.obj'
8. with open(objFilePath) as file:
9.     points = []
10.     while 1:
11.         line = file.readline()
12.         if not line:
13.             break
14.         strs = line.split(" ")
15.         if strs[0] == "v":
16.             points.append(np.array(strs[1:7],dtype=float))
17.         if strs[0] == "vt":
18.             break
19. # points原本为列表，需要转变为矩阵，方便处理         
20. pcd = np.array(points)
22. pcd_vector = o3d.geometry.PointCloud()
23. pcd_vector.points = o3d.utility.Vector3dVector(pcd[:, :3])
24. pcd_vector.colors = o3d.utility.Vector3dVector(pcd[:,3:6])
25. o3d.visualization.draw_geometries([pcd_vector])

复制代码

  康康Area_5里office_8的效果：
  原图：
               ground truth：

  predict：
         

---

     OK，大抵算是把这个PointNet++复现完了，侧重做了下点云分割，给毕设做准备。总的来讲，训练和推测的过程并不难，为了康康效果，可视化的部门倒是花了挺长时间。零件分割和场景分割本质上讲着实是一回事，就是在代码内里这两个分割用了差别的模型来训练。之后筹划本身制作数据集来训练一下，先拿零件分割的模型来做，究竟场景分割做成S3DIS情势的数据集有点贫苦。总之跟着这篇blog走肯定是能跑通PointNet++的。
  

---

     添加一个本身做的用pointnet++做的书缝辨认项目，GitHub内里有数据集和代码：https://github.com/struggler176393/Pointnet_book_seam。

来源：https://blog.csdn.net/astruggler/article/details/128354761
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！