Why is the segmentation success rate very low when I test with your weight?

[SWQN2]

THANK YOU for your great work. But when I used the pretrained weights you provided on the nuScenes dataset for testing, it didn’t achieve a high success rate. Could you please tell me if there are some issues with my code? Here is my code：

import os
import sys
import numpy as np
import torch
import matplotlib.colors as mcolors
import open3d as o3d

确保 PTv3 和项目根目录在 sys.path 中

script_dir = os.path.dirname(os.path.abspath(file))
project_root = os.path.abspath(os.path.join(script_dir, ".."))
ptv3_path = os.path.join(script_dir, "PTv3")
sys.path.append(ptv3_path)
sys.path.append(project_root)

导入模型和所需模块

from PTv3.point_transformer_v3 import PointTransformerV3, Point
from nuscenes.nuscenes import NuScenes
from nuscenes.utils.data_classes import LidarPointCloud

定义标签映射字典

def get_learning_map(ignore_index=-1):
return {
0: ignore_index, 1: ignore_index, 2: 6, 3: 6, 4: 6, 5: ignore_index,
6: 6, 7: ignore_index, 8: ignore_index, 9: 0, 10: ignore_index,
11: ignore_index, 12: 7, 13: ignore_index, 14: 1, 15: 2, 16: 2, 17: 3,
18: 4, 19: ignore_index, 20: ignore_index, 21: 5, 22: 8, 23: 9,
24: 10, 25: 11, 26: 12, 27: 13, 28: 14, 29: ignore_index, 30: 15,
31: ignore_index
}

初始化 NuScenes 数据集路径

nusc = NuScenes(version='v1.0-mini', dataroot='/home/lt/anynode/scripts/data/nuscenes', verbose=True)

加载一个样本并提取点云数据

sample = nusc.sample[0]
lidar_token = sample['data']['LIDAR_TOP']
lidar_data = nusc.get('sample_data', lidar_token)
lidar_filepath = os.path.join(nusc.dataroot, lidar_data['filename'])
pointcloud = LidarPointCloud.from_file(lidar_filepath)
points = pointcloud.points[:3, :].T # 提取 x, y, z 坐标
reflectivity = pointcloud.points[3, :].T / 255.0 # 与训练时一致，缩放至 [0, 1] 范围

加载语义分割标签

lidarseg_filepath = os.path.join(nusc.dataroot, nusc.get('lidarseg', lidar_token)['filename'])
segmentation_labels = np.fromfile(lidarseg_filepath, dtype=np.uint8)

使用 learning_map 将实际标签映射到模型训练的 16 个类别标签空间

learning_map = get_learning_map(ignore_index=-1)
mapped_labels = np.vectorize(learning_map.get)(segmentation_labels)

准备与训练一致的 grid_size 并确保输入处理与训练时一致

grid_size = 0.05 # 与训练配置文件一致
coords = torch.tensor(points, dtype=torch.float32).cuda()
reflectivity = torch.tensor(reflectivity[:, None], dtype=torch.float32).cuda()
feat = torch.cat([coords, reflectivity], dim=1)

准备模型的输入

input_data = Point({
'feat': feat,
'coord': coords,
'grid_coord': torch.div(coords, grid_size).int().cuda(),
'offset': torch.tensor([0, coords.shape[0]]).cuda(),
'batch': torch.zeros(coords.shape[0], dtype=torch.long).cuda(),
})

定义模型

model = PointTransformerV3(
in_channels=4,
order=('z', 'z-trans', 'hilbert', 'hilbert-trans'),
stride=(2, 2, 2, 2),
enc_depths=(2, 2, 2, 6, 2),
enc_channels=(32, 64, 128, 256, 512),
enc_num_head=(2, 4, 8, 16, 32),
enc_patch_size=(1024, 1024, 1024, 1024, 1024),
dec_depths=(2, 2, 2, 2),
dec_channels=(64, 64, 128, 256),
dec_num_head=(4, 4, 8, 16),
dec_patch_size=(1024, 1024, 1024, 1024),
mlp_ratio=4,
qkv_bias=True,
qk_scale=None,
attn_drop=0.0,
proj_drop=0.0,
drop_path=0.3,
pre_norm=True,
shuffle_orders=True,
enable_rpe=False, # 与训练一致
enable_flash=True,
upcast_attention=False,
upcast_softmax=False,
cls_mode=False
).cuda()

定义函数用于去除多余的前缀并加载模型权重

def load_weights(model, checkpoint_path):
checkpoint = torch.load(checkpoint_path, map_location='cuda')
state_dict = checkpoint['state_dict']

# 去除 'module.' 前缀
new_state_dict = {}
for k, v in state_dict.items():
    name = k.replace('module.', '')  # 去掉 'module.' 前缀
    new_state_dict[name] = v

# 加载去掉前缀后的 state_dict
model.load_state_dict(new_state_dict, strict=False)

加载模型权重

checkpoint_path = "/home/lt/anynode/scripts/PTv3/model_best.pth"
load_weights(model, checkpoint_path)

前向传播，获取每个点的预测语义标签

model.eval()
with torch.no_grad():
output = model(input_data)
predicted_labels = torch.argmax(output.feat, dim=1).cpu().numpy()

对预测标签进行映射，使得预测标签也落在 0-16 类别范围内

mapped_predicted_labels = np.vectorize(lambda x: learning_map.get(x, -1))(predicted_labels)

确保忽略 index 处理与训练一致

ignore_index = -1
valid_mask = mapped_labels != ignore_index
correct_predictions = (mapped_predicted_labels[valid_mask] == mapped_labels[valid_mask]).sum()
total_points = valid_mask.sum()
accuracy = correct_predictions / total_points

打印标签对比和预测成功率

print("实际的语义标签:", mapped_labels[valid_mask])
print("预测的语义标签:", mapped_predicted_labels[valid_mask])
print(f"预测成功率: {accuracy * 100:.2f}%")

可视化设置

semantic_colors = list(mcolors.TABLEAU_COLORS.values()) # 使用 TABLEAU 颜色集
unique_labels = np.unique(mapped_predicted_labels)
max_label = unique_labels.max()

扩展颜色映射，以应对大于 16 的类别

color_map = [semantic_colors[i % len(semantic_colors)] for i in range(max_label + 1)]

根据预测标签给每个点分配颜色

point_colors = np.array([mcolors.to_rgb(color_map[label]) for label in mapped_predicted_labels])

创建 Open3D 点云并添加颜色

pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points)
pcd.colors = o3d.utility.Vector3dVector(point_colors)

显示点云

o3d.visualization.draw_geometries([pcd], window_name="Point Cloud Segmentation", width=800, height=600)

[1Hun0ter1]

Have you solved the problem?

[Gofinge]

Hi, the released model weight is trained with v1.5.1. There exists some change in our data pipeline, so align the codebase version for using these model weights.