省流：
（1）训练自己模型需要部分代码修改请看博主上一篇博文：Windows11和Ubuntu用PointNet++训练自己的数据集（部件分割模型）
（2）此外由于自己的需求，所以只做物体部件分割（partseg）任务，没有复现分类（classification）和语义分割（semseg）。
（3）本文功能：利用已经训练好的部分分割模型，去分割一个新的luntai点云（1个物体luntai有2个部件top和bottom），用open3d显示分割效果，再将检测结果保存为txt文档，文档格式是n*6即n个点，一个点包括x、y、z、r、g、b信息，该点属于luntai的哪个部件就对应同一颜色。方便利用cloudCompare查看。

代码参考这位老哥：PointNet++分割结果可视化程序。但是这位老哥是可视化的shapenet数据集（16种物体，共50个部件。）训练的模型，有些参数没有改到位，并且老哥也没有保存检测结果到txt文档。博主在这里对老哥的代码再详细改进一下。

先上本项目visualization.py可视化代码，后面小节有改动讲解

#可视化代码参考csdn的
# https://blog.csdn.net/m0_63108038/article/details/139250573?ops_request_misc=%257B%2522request%255Fid%2522%253A%2522214c5efa414336fe0cf2a194844a950f%2522%252C%2522scm%2522%253A%252220140713.130102334..%2522%257D&request_id=214c5efa414336fe0cf2a194844a950f&biz_id=0&utm_medium=distribute.pc_search_result.none-task-blog-2~all~baidu_landing_v2~default-1-139250573-null-null.142^v102^pc_search_result_base6&utm_term=pointnet%2B%2B%E9%83%A8%E4%BB%B6%E5%88%86%E5%89%B2%E5%8F%AF%E8%A7%86%E5%8C%96&spm=1018.2226.3001.4187
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 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


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 = {'Airplane': '02691156'}
        #自己改动  1/6
        self.cat = {'luntai': '01234567'}
        # 输出的是元组，('Airplane',123.txt)
        # 下面self.classe中的数字是重点，数字代表所选择的类别，比如0是飞机，2是帽子等等
        #self.classes = {'Airplane': 0}

        # 自己改动  2/6
        self.classes = {'luntai': 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
            # 下面这行注意改一下，用shapenet的话填16
            #seg_pred, _ = model(points, self.to_categorical(label, 16))
            #自己改动  3/6
            seg_pred, _ = model(points, self.to_categorical(label, 1))#只有一个物体即luntai
            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)
            #print("seg_pred shape:", seg_pred.shape)   (30000, 4)
            self.predict_pcd = seg_pred
            self.centroid = centroid# tensor([[  45.2890, -159.3051,    9.7356]])  平移
            self.m = m#一个浮点数   缩放

            print("seg_pred shape:", seg_pred.shape)#30000 4
            print("seg_pred:", seg_pred)
            print("m:",m)#一个浮点数   缩放
            print("centroid:", centroid)  # tensor([[  45.2890, -159.3051,    9.7356]])  平移

    def o3d_draw_3d_img(self):

        pcd = self.predict_pcd
        pcd_vector = o3d.geometry.PointCloud()
        print("pcd shape:", pcd.shape)
        # 加载点坐标
        pcd_vector.points = o3d.utility.Vector3dVector(self.m * pcd[:, :3] + self.centroid)  #pcd_vector 除了点，还有颜色     这里括号的计算是影射回原来的点 即114行的缩放和平移
        print("pcd_vector.points", pcd_vector.points)#pcd_vector.points std::vector<Eigen::Vector3d> with 30000 elements.


        #自己加的代码1/3：点坐标
        Points = np.asarray(pcd_vector.points)
        print("Points shape:", Points.shape)  # 输出形状，例如 (30000, 3)
        print("First 5 points:\n", Points[:5])  # 打印前 5 个点的坐标


        # colors = np.random.randint(255, size=(2,3))/255
        # colors这地方改了一下，注意设置多一些，避免分割多个部分颜色不够
        colors = np.array([[1, 0, 0], [0, 0, 0], [0.3, 0.3, 0.3], [0, 1, 0], [0, 0, 1], [0.5, 0.5, 1],[0,1,1],[0.7,0.7,0.7],[0,0,0],[1,1,1]])#颜色个数大于等于  部件总数，这里槽底槽顶两个部件，两个颜色就好
        #print("pcd shape:", pcd.shape)   #(30000, 4)
        print("pcd:", pcd)
        #自己改动 4/6
        pcd_vector.colors = o3d.utility.Vector3dVector(colors[list(map(int, pcd[:, 3])), :])

        #自己加的代码2/3：颜色
        Colors=np.asarray(pcd_vector.colors)
        print("Colors shape:", Colors.shape)  # 输出形状，例如 (30000, 3)
        print("First 5 colors:\n", Colors[:5])  # 打印前 5 个点的颜色

        # 自己加的代码3/3：打印点的坐标和分类颜色
        combined_data = np.hstack((Points, Colors))# 合并坐标和颜色为一个 (N, 6) 的数组  xyzrgb
        # 指定输出文件名
        output_file = "D:\\YinParker\\Desktop\\pointnet2_out.txt"
        # 将数据保存到文件
        # 每行格式为：x \t y \t z \t r \t g \t b
        np.savetxt(output_file, combined_data, delimiter="\t", fmt="%.6f")


        #print("pcd_vector shape:", pcd_vector.points)
        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 """
        print(y)
        print(num_classes)
        print(y.cpu().data.numpy())
        new_y = torch.eye(num_classes)[y.cpu().data.numpy(),]
        if (y.is_cuda):
            return new_y.cuda()
        return new_y

#自己改动 5/6   改50改为2  True改为False
def load_models(model_dict={'PonintNet': [pointnet2(num_classes=2, normal_channel=False).eval(),
                                          r'D:\\YinParker\\Desktop\\AllFileFolder\\A_Other_projects_in_laboratory\\pointNet2/Pointnet_Pointnet2_pytorch-master/log/part_seg/2025-03-12_09-43/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 = 50  # 填写数据集的类别数 如果是s3dis这里就填13   shapenet这里就填50

    #自己改 6/6
    num_classes = 2 #1个物体 的2个部件

    # color_image = o3d.io.read_image('image/rgb1.jpg')
    # depth_image = o3d.io.read_image('image/depth1.png')
    #txt_path = 'D:\\YinParker\\Desktop\\AllFileFolder\\A_Other_projects_in_laboratory\\pointNet2/Pointnet_Pointnet2_pytorch-master/data/shapenetcore_partanno_segmentation_benchmark_v0_normal/01234567/2.txt'
    txt_path = 'D:\\YinParker\\Desktop\\1.txt'
    # 通过numpy读取txt点云
    pcd = np.genfromtxt(txt_path, delimiter="\t")  #注意以啥进行分割 我是分隔符
    print("pcd 的形状:", pcd.shape)  # 应该是 (n, 3)   (14737, 3)

    point_cloud = o3d.geometry.PointCloud()
    # 加载点坐标
    # txt点云前三个数值一般对应x、y、z坐标，可以通过open3d.geometry.PointCloud().points加载
    # 如果有法线或颜色，那么可以分别通过open3d.geometry.PointCloud().normals或open3d.geometry.PointCloud().colors加载
    point_cloud.points = o3d.utility.Vector3dVector(pcd[:, :3])
    # point_cloud.normals = o3d.utility.Vector3dVector(pcd[:, 3:6])

point_cloud.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.03, max_nn=30))
#print(np.asarray(point_cloud.points))
#print(np.asarray(point_cloud.normals))
#TEST_DATASET = PartNormalDataset(point_cloud, 30000, True)
TEST_DATASET = PartNormalDataset(point_cloud, 30000, False)
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)

一、新建.py

新建一个visualization.py放到models文件夹下。如下图。

因为你如果把visualization.py放到其他文件夹，就需要把visualization.py的from pointnet2_part_seg_msg import get_model as pointnet2改为from models.pointnet2_part_seg_msg import get_model as pointnet2，但这样会造成问题，因为要导入的models文件夹下的pointnet2_part_seg_msg.py
里有from pointnet2_utils import PointNetSetAbstractionMsg,PointNetSetAbstraction,PointNetFeaturePropagation会报错找不到 pointnet2_utils，虽然 pointnet2_utils.py就在models文件夹里。所以省事就把visualization.py放到models文件夹下了。

二、改动代码

改动1/6：

程序第41行左右，改成自己的物件，后面的‘0123456’是自己训练时该物件用的编号，只要看过上一篇博客或者训练过自己的模型，应该都懂。

改动2/6：

程序47行左右。说实话，这里我不太理解，因为之前训练时的train_partseg.py里是seg_classes = {'luntai': [0, 1]}，这里别写。我是一个物体就写0就可以，即使这个物体有两个部件也写0。如果你的程序报错，可以试试改改这里。

改动3/6：

程序103行左右，一个物体即luntai就写1，几个物体就写几。

改动4/6：

程序146行左右，之前老哥传入的数据是带向量的，所以写的6。我的不带向量，只有xyz所以就是3

改动5/6：

程序178行左右，几个部件就写几，我是2个部件，且没有法向量所以False，最后是你训练的模型参数best_model.pth的目录，写到所在文件夹就行。

改动6/6：

程序229行左右：改成你的部件数量、你的待检测的点云路径

改动7/6

也不是改动，只是说明：在代码142行可以设置颜色。我是两个部件，所以代码就选前两个颜色了，红色（1，0，0）和黑色（0，0，0）。这里的rgb值是0-1 而不是 0~255 。所以在后面将txt文件导入到cloudCompare时要选择好。

2025年3月31日11:34:45 补充

253行左右的这个参数是模型采样的点的个数，也是最后输出的txt的点个数（意思说不会输出原txt文本点的个数，除非设置为原文本点个数），该参数根据自己的实际点云点数多少确定。

保存检测结果

加代码1/3：得到点

代码134行左右：pcd_vector包括两部分数据：点数据和颜色，要分别得到这两部分数据各自存到一个变量中

加代码2/3：得到每个点对应颜色

代码148行左右：

加代码3/3:保存txt

代码153行左右：数据的格式是n×6 每行的六个数即x y z r g b，中间用\t即tab隔开。

结果演示

open3d：

txt文档

cloudCompare

导入数据时RGB选择0到1的

三、训练自己的语义分割模型，请看下一篇文章

若对你有用的话，请帮忙 点个赞或者收藏