Pytorch1.7复现PointNet++点云分割（含Open3D可视化）（文末有一个自己做的书缝识别项目代码）

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

下载代码

https://github.com/yanx27/Pointnet_Pointnet2_pytorch

  我的运行环境是pytorch1.7+cuda11.0。

训练

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

对象分类

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

## e.g., pointnet2_ssg without normal features
python train_classification.py --model pointnet2_cls_ssg --log_dir pointnet2_cls_ssg
python test_classification.py --log_dir pointnet2_cls_ssg
## e.g., pointnet2_ssg with normal features
python train_classification.py --model pointnet2_cls_ssg --use_normals --log_dir pointnet2_cls_ssg_normal
python test_classification.py --use_normals --log_dir pointnet2_cls_ssg_normal
## e.g., pointnet2_ssg with uniform sampling
python train_classification.py --model pointnet2_cls_ssg --use_uniform_sample --log_dir pointnet2_cls_ssg_fps
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/。

    运行也很简单：

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

    shapenet数据集txt文件格式：前三个点是xyz,点云的位置坐标，后三个点是点云的法向信息，最后一个点是这个点所属的小类别，即1表示所属50个小类别中的第一个。

    写个代码用open3d可视化shapenet数据集的txt文件（随机配色）：

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

    GPU内存不够减小一下batch_size。

    我这里训练了一下，接着代码的best_model.pth继续训练150轮，RTX3080单显卡训练一轮得六七分钟，150轮花了半天多的时间。

    网上的代码基本test一下分割的一些参数就结束了，没有做可视化，参考这篇blog做了一下结果的可视化：PointNet++分割预测结果可视化。这篇blog首先用网络将输入图像的预测结果存为txt文件，然后用Matplotlib做可视化，过程有点复杂了，用open3d做可视化比较简洁一点，代码如下：

  import tqdm
  import matplotlib
  import torch
  import os
  import warnings
  import numpy as np
  import open3d as o3d
  from torch.utils.data import Dataset
  import pybullet as p
  from models.pointnet2_part_seg_msg import get_model as pointnet2
  import time
  warnings.filterwarnings('ignore')
  matplotlib.use("Agg")
  def pc_normalize(pc):
      centroid = np.mean(pc, axis=0)
      pc = pc - centroid
      m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
      pc = pc / m
      return pc,centroid,m
  def generate_pointcloud(color_image, depth_image,width=1280,height=720,fov=50,near=0.01,far=5):
      rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(color_image, depth_image,convert_rgb_to_intensity=False)
      intrinsic = o3d.camera.PinholeCameraIntrinsic(o3d.camera.PinholeCameraIntrinsicParameters.Kinect2DepthCameraDefault )
      aspect = width / height
      projection_matrix = p.computeProjectionMatrixFOV(fov, aspect, near, far)
      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)
      point_cloud = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, intrinsic)
      
      point_cloud.estimate_normals( search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.01, max_nn=30))
      return point_cloud
  class PartNormalDataset(Dataset):
      def __init__(self, point_cloud, npoints=2500, normal_channel=False):
          self.npoints = npoints # 采样点数
          self.cat = {}
          self.normal_channel = normal_channel # 是否使用法向信息
          position_data = np.asarray(point_cloud.points)
          normal_data = np.asarray(point_cloud.normals)
          self.raw_pcd = np.hstack([position_data,normal_data]).astype(np.float32)
          self.cat = {'board':'12345678'}
          # 输出的是元组，('Airplane',123.txt)
          self.classes = {'board': 0} 
          data = self.raw_pcd
          if not self.normal_channel:  # 判断是否使用法向信息
              self.point_set = data[:, 0:3]
          else:
              self.point_set = data[:, 0:6]
          self.point_set[:, 0:3],self.centroid,self.m = pc_normalize(self.point_set[:, 0:3]) # 做一个归一化
          choice = np.random.choice(self.point_set.shape[0], self.npoints, replace=True) # 对一个类别中的数据进行随机采样 返回索引，允许重复采样
          # resample
          self.point_set =  self.point_set[choice, :] # 根据索引采样
      def __getitem__(self, index):
          cat = list(self.cat.keys())[0]
          cls = self.classes[cat] # 将类名转换为索引
          cls = np.array([cls]).astype(np.int32)
          return self.point_set, cls, self.centroid, self.m # pointset是点云数据，cls十六个大类别，seg是一个数据中，不同点对应的小类别
      def __len__(self):
          return 1
  class Generate_txt_and_3d_img:
      def __init__(self,num_classes,testDataLoader,model,visualize = False):
          self.testDataLoader = testDataLoader
          self.num_classes = num_classes
          self.heat_map = False # 控制是否输出heatmap
          self.visualize = visualize # 是否open3d可视化
          self.model = model
          self.generate_predict()
          self.o3d_draw_3d_img()
      def __getitem__(self, index):
          return self.predict_pcd_colored
      def generate_predict(self):
          for _, (points, label,centroid,m) in tqdm.tqdm(enumerate(self.testDataLoader),
                         total=len(self.testDataLoader),smoothing=0.9):
              #点云数据、整个图像的标签、每个点的标签、  没有归一化的点云数据（带标签）torch.Size([1, 7, 2048])
              points = points.transpose(2, 1)
              #print('1',target.shape) # 1 torch.Size([1, 2048])
              xyz_feature_point = points[:, :6, :]
              model = self.model
              seg_pred, _ = model(points, self.to_categorical(label, 1))
              seg_pred = seg_pred.cpu().data.numpy()
              if self.heat_map:
                  out =  np.asarray(np.sum(seg_pred,axis=2))
                  seg_pred = ((out - np.min(out) / (np.max(out) - np.min(out))))
              else:
                  seg_pred = np.argmax(seg_pred, axis=-1)  # 获得网络的预测结果 b n c
              seg_pred = np.concatenate([np.asarray(xyz_feature_point), seg_pred[:, None, :]],
                      axis=1).transpose((0, 2, 1)).squeeze(0) 
              self.predict_pcd = seg_pred
              self.centroid = centroid
              self.m = m
      def o3d_draw_3d_img(self):
          pcd = self.predict_pcd
          pcd_vector = o3d.geometry.PointCloud()
          # 加载点坐标
          pcd_vector.points = o3d.utility.Vector3dVector(self.m * pcd[:, :3] + self.centroid)
          # colors = np.random.randint(255, size=(2,3))/255
          colors = np.array([[0.8, 0.8, 0.8],[1,0,0]])
          pcd_vector.colors = o3d.utility.Vector3dVector(colors[list(map(int,pcd[:, 6])),:])
          if self.visualize:
              coord_mesh = o3d.geometry.TriangleMesh.create_coordinate_frame(size = 0.1, origin = [0,0,0])
              o3d.visualization.draw_geometries([pcd_vector,coord_mesh])
          self.predict_pcd_colored = pcd_vector
      def to_categorical(self,y, num_classes):
          """ 1-hot encodes a tensor """
          new_y = torch.eye(num_classes)[y.cpu().data.numpy(),]
          if (y.is_cuda):
              return new_y.cuda()
          return new_y
  def load_models(model_dict={'PonintNet': [pointnet2(num_classes=2,normal_channel=True).eval(),r'./log/part_seg/pointnet2_part_seg_msg/checkpoints']}):
      model = list(model_dict.values())[0][0]
      checkpoints_dir = list(model_dict.values())[0][1]
      weight_dict = torch.load(os.path.join(checkpoints_dir,'best_model.pth'))
      model.load_state_dict(weight_dict['model_state_dict'])
      return model
  class Open3dVisualizer():
  	def __init__(self):
  		self.point_cloud = o3d.geometry.PointCloud()
  		self.o3d_started = False
  		self.vis = o3d.visualization.VisualizerWithKeyCallback()
  		self.vis.create_window()
  	def __call__(self, points, colors):
  		self.update(points, colors)
  		return False
  	def update(self, points, colors):
  		coord_mesh = o3d.geometry.TriangleMesh.create_coordinate_frame(size = 0.15, origin = [0,0,0])
  		self.point_cloud.points = points
  		self.point_cloud.colors = colors
  		# self.point_cloud.transform([[1,0,0,0],[0,-1,0,0],[0,0,-1,0],[0,0,0,1]])
  		# self.vis.clear_geometries()
  		# Add geometries if it is the first time
  		if not self.o3d_started:
  			self.vis.add_geometry(self.point_cloud)
  			self.vis.add_geometry(coord_mesh)
  			self.o3d_started = True
  		else:
  			self.vis.update_geometry(self.point_cloud)
  			self.vis.update_geometry(coord_mesh)
  		self.vis.poll_events()
  		self.vis.update_renderer()
  if __name__ =='__main__':
      
      num_classes = 2 # 填写数据集的类别数 如果是s3dis这里就填13   shapenet这里就填50
      
      color_image = o3d.io.read_image('image/rgb1.jpg')
      depth_image = o3d.io.read_image('image/depth1.png')
      
      point_cloud = generate_pointcloud(color_image=color_image, depth_image=depth_image)
      TEST_DATASET = PartNormalDataset(point_cloud,npoints=30000, normal_channel=True)
      testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=1, shuffle=False, num_workers=0,drop_last=True)
      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/。

  cd data_utils
  python collect_indoor3d_data.py
  

    运行：

  ## Check model in ./models 
  ## e.g., pointnet2_ssg
  python train_semseg.py --model pointnet2_sem_seg --test_area 5 --log_dir pointnet2_sem_seg
  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()变量显示，代码如下：

  import copy
  import numpy as np
  import open3d as o3d
  import os
  objFilePath = 'log/sem_seg/pointnet2_sem_seg/visual/Area_5_office_8_gt.obj'
  with open(objFilePath) as file:
      points = []
      while 1:
          line = file.readline()
          if not line:
              break
          strs = line.split(" ")
          if strs[0] == "v":
              points.append(np.array(strs[1:7],dtype=float))
          if strs[0] == "vt":
              break
  # points原本为列表，需要转变为矩阵，方便处理          
  pcd = np.array(points)
  pcd_vector = o3d.geometry.PointCloud()
  pcd_vector.points = o3d.utility.Vector3dVector(pcd[:, :3])
  pcd_vector.colors = o3d.utility.Vector3dVector(pcd[:,3:6])
  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。