折根妙妙屋

1
Input (N, 3)
2
    ↓ SA1 (N/4, 64)
3
    ↓ SA2 (N/16, 128)
4
    ↓ SA3 (N/64, 256)
5
    ↓ SA4 (1, 1024) —— Global Feature
6
    ↓ MLP → Softmax → Class Label

1
Input (N, 3)
2
    ↓ SA1 ────────┐
3
    ↓ SA2 ───────┐│
4
    ↓ SA3 ──────┐││
5
    ↓ SA4 ─────┐│││
6
         FP3 ←─┘│││
7
         FP2 ←──┘││
8
         FP1 ←───┘│
9
         FP0 ←────┘
10
    ↓ Point-wise Classifier → (N, NumClasses)

1
import os
2
import argparse
3
import numpy as np
4
import torch
5

6
CLASSES = ['airplane', 'bathtub', 'bed', 'bench', 'bookshelf', 'bottle', 'bowl',
7
           'car', 'chair', 'cone', 'cup', 'curtain', 'desk', 'door', 'dresser',
8
           'flower_pot', 'glass_box', 'guitar', 'keyboard', 'lamp', 'laptop',
9
           'mantel', 'monitor', 'night_stand', 'person', 'piano', 'plant',
10
           'radio', 'range_hood', 'sink', 'sofa', 'stairs', 'stool', 'table',
11
           'tent', 'toilet', 'tv_stand', 'vase', 'wardrobe', 'xbox']
12

13
def load_model(log_dir, use_normals=False):
14
    """从 log/classification/xxx 加载权重"""
15
    if 'ssg' in log_dir:
16
        from models.pointnet2_cls_ssg import get_model
17
        model = get_model(num_class=40, normal_channel=use_normals).cuda()
18
    else:
19
        from models.pointnet2_cls_msg import get_model
20
        model = get_model(num_class=40, normal_channel=use_normals).cuda()
21

22
    pth_path = os.path.join('log/classification', log_dir, 'best_model.pth')
23
    ckpt = torch.load(pth_path, map_location='cuda')
24
    model.load_state_dict(ckpt['model_state_dict'] if 'model_state_dict' in ckpt else ckpt)
25
    model.eval()
26
    return model
27

28
def preprocess(points, num_point=1024):
29
    """
30
    预处理：采样到固定点数 + 归一化到单位球
31
    ⚠️ 必须与训练时做完全相同的操作，否则半径参数失效
32
    """
33
    N = len(points)
34
    if N > num_point:
35
        idx = np.random.choice(N, num_point, replace=False)
36
    else:
37
        idx = np.random.choice(N, num_point, replace=True)
38
    points = points[idx]
39

40
    centroid = np.mean(points[:, :3], axis=0)
41
    points[:, :3] -= centroid
42
    max_dist = np.max(np.sqrt(np.sum(points[:, :3]**2, axis=1)))
43
    if max_dist > 1e-6:
44
        points[:, :3] /= max_dist
45
    return points.astype(np.float32)
46

47
def infer(model, points):
48
    pts = torch.from_numpy(points).unsqueeze(0).cuda()  # (1, 1024, 3)
49
    with torch.no_grad():
50
        pred, _ = model(pts)
51
        prob = torch.softmax(pred, dim=1)
52
        conf, pred_idx = torch.max(prob, dim=1)
53
    return pred_idx.item(), conf.item()
54

55
def main(args):
56
    model = load_model(args.weights, args.normal)
57

58
    if os.path.isfile(args.source):
59
        files = [args.source]
60
    else:
61
        files = [os.path.join(args.source, f) for f in os.listdir(args.source)
62
                 if f.endswith('.npy') or f.endswith('.txt')]
63
        files.sort()
64

65
    results = []
66
    for fpath in files:
67
        points = np.load(fpath) if fpath.endswith('.npy') else np.loadtxt(fpath)
68
        points = preprocess(points, num_point=1024)
69
        cls_idx, conf = infer(model, points)
70
        line = f"{os.path.basename(fpath):<30} -> {CLASSES[cls_idx]:<12} ({conf*100:.2f}%)"
71
        print(line)
72
        results.append(f"{os.path.basename(fpath)},{CLASSES[cls_idx]},{conf:.6f}")
73

74
    if args.save_txt:
75
        with open(args.save_txt, 'w') as f:
76
            f.write("filename,class,confidence\n")
77
            f.write("\n".join(results))
78
        print(f"\n结果已保存: {args.save_txt}")
79

80
if __name__ == '__main__':
81
    parser = argparse.ArgumentParser()
82
    parser.add_argument('--weights', type=str, required=True, help='如 pointnet2_cls_msg')
83
    parser.add_argument('--source', type=str, required=True, help='单个文件或文件夹')
84
    parser.add_argument('--normal', action='store_true')
85
    parser.add_argument('--save-txt', type=str, default='detect_results.csv')
86
    args = parser.parse_args()
87
    main(args)

1
# 单文件推理
2
python3 detect_cls.py --weights pointnet2_cls_msg \
3
    --source data/modelnet40_normal_resampled/airplane/airplane_0001.npy
4

5
# 批量推理并保存 CSV
6
python3 detect_cls.py --weights pointnet2_cls_msg \
7
    --source data/modelnet40_normal_resampled/airplane/ \
8
    --save-txt airplane_results.csv

1
import torch
2
import numpy as np
3
from models.pointnet2_cls_msg import get_model
4

5
model = get_model(num_class=40, normal_channel=False).cuda()
6
# ... 加载权重 ...
7

8
dummy_input = torch.randn(1, 1024, 3).cuda()
9

10
# 热身：GPU 缓存预热
11
for _ in range(10):
12
    with torch.no_grad():
13
        _ = model(dummy_input)
14

15
# 正式计时
16
starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
17
times = []
18
for _ in range(100):
19
    starter.record()
20
    with torch.no_grad():
21
        pred = model(dummy_input)
22
    ender.record()
23
    torch.cuda.synchronize()
24
    times.append(starter.elapsed_time(ender))
25

26
print(f"平均延迟: {np.mean(times):.2f} ms")
27
print(f"等效 FPS: {1000/np.mean(times):.1f}")

1
import numpy as np
2
import torch
3
import open3d as o3d
4

5
CLASSES = ['airplane', 'bathtub', 'bed', 'bench', 'bookshelf', 'bottle', 'bowl',
6
           'car', 'chair', 'cone', 'cup', 'curtain', 'desk', 'door', 'dresser',
7
           'flower_pot', 'glass_box', 'guitar', 'keyboard', 'lamp', 'laptop',
8
           'mantel', 'monitor', 'night_stand', 'person', 'piano', 'plant',
9
           'radio', 'range_hood', 'sink', 'sofa', 'stairs', 'stool', 'table',
10
           'tent', 'toilet', 'tv_stand', 'vase', 'wardrobe', 'xbox']
11

12
def visualize_prediction(points_path, model, use_normals=False):
13
    points = np.load(points_path).astype(np.float32)
14
    if points.shape[1] == 6 and not use_normals:
15
        points = points[:, :3]
16

17
    # 预处理
18
    pts = torch.from_numpy(points).unsqueeze(0).cuda()
19
    with torch.no_grad():
20
        pred, _ = model(pts)
21
        pred_idx = torch.argmax(pred, dim=1).item()
22
        conf = torch.softmax(pred, dim=1)[0][pred_idx].item()
23

24
    # Open3D 可视化
25
    pcd = o3d.geometry.PointCloud()
26
    pcd.points = o3d.utility.Vector3dVector(points[:, :3])
27

28
    # 按高度着色
29
    colors = np.zeros((points.shape[0], 3))
30
    z_norm = (points[:, 2] - points[:, 2].min()) / (points[:, 2].max() - points[:, 2].min() + 1e-6)
31
    colors[:, 0] = z_norm
32
    colors[:, 2] = 1 - z_norm
33
    pcd.colors = o3d.utility.Vector3dVector(colors)
34

35
    vis = o3d.visualization.Visualizer()
36
    vis.create_window(window_name=f"Pred: {CLASSES[pred_idx]} ({conf*100:.2f}%)", width=800, height=600)
37
    vis.add_geometry(pcd)
38
    vis.get_render_option().point_size = 4.0
39
    vis.run()
40
    vis.destroy_window()

1
cd ~/projects/pointcloud/Pointnet_Pointnet2_pytorch
2

3
# 语义分割（S3DIS 或自定义数据集）
4
python3 train_semseg.py \
5
    --model pointnet2_sem_seg \
6
    --log_dir pointnet2_sem_seg \
7
    --batch_size 16 \
8
    --epoch 50
9

10
# 部件分割（ShapeNet Part，带法向量）
11
python3 train_partseg.py \
12
    --model pointnet2_part_seg_msg \
13
    --normal \
14
    --log_dir pointnet2_part_seg_msg \
15
    --batch_size 16 \
16
    --epoch 50

1
python3 test_semseg.py --log_dir pointnet2_sem_seg --batch_size 16

1
IoU = 预测正确的点数 / (预测为该类的点数 + 真值为该类的点数 - 预测正确的点数)
2
    = 交集 / 并集

1
from models.pointnet2_sem_seg import get_model
2
model = get_model(num_class=4).cuda()  # 4 类：背景/车厢壁/车钩/煤炭

1
import os
2
import argparse
3
import numpy as np
4
import torch
5
from scipy.spatial import cKDTree
6

7
# 分割类别定义（按需修改）
8
SEG_CLASSES = {
9
    0: 'background',
10
    1: 'carriage_wall',
11
    2: 'coupler',
12
    3: 'coal'
13
}
14

15
def farthest_point_sample(points, npoint):
16
    """
17
    FPS 最远点采样：保证采样点空间覆盖均匀，
18
    避免随机采样导致角落遗漏。
19
    """
20
    N, D = points.shape
21
    xyz = points[:, :3]
22
    centroids = np.zeros(npoint)
23
    distance = np.ones(N) * 1e10
24
    farthest = np.random.randint(0, N)
25

26
    for i in range(npoint):
27
        centroids[i] = farthest
28
        centroid = xyz[farthest, :]
29
        dist = np.sum((xyz - centroid) ** 2, -1)
30
        mask = dist < distance
31
        distance[mask] = dist[mask]
32
        farthest = np.argmax(distance, -1)
33

34
    return points[centroids.astype(np.int32)]
35

36

37
def infer_semseg(model, raw_points, num_point=4096, num_class=4):
38
    """
39
    输入: raw_points (N, 3)  任意点数，如 10000
40
    输出: labels (N,)         每个原始点的类别标签
41
    """
42
    # 1. FPS 降采样到网络训练时的固定输入尺寸
43
    sampled = farthest_point_sample(raw_points, num_point)
44

45
    # 2. 归一化（必须与训练时完全一致！）
46
    centroid = np.mean(sampled, axis=0)
47
    sampled -= centroid
48
    max_dist = np.max(np.sqrt(np.sum(sampled**2, axis=1)))
49
    sampled /= (max_dist + 1e-8)
50

51
    # 3. 网络推理
52
    pts = torch.from_numpy(sampled).unsqueeze(0).float().cuda()  # (1, 4096, 3)
53
    model.eval()
54
    with torch.no_grad():
55
        pred = model(pts)  # (1, 4096, num_class)
56
        sampled_labels = torch.argmax(pred, dim=2).squeeze(0).cpu().numpy()  # (4096,)
57

58
    # 4. 插值回原始点数：每个原始点找采样点中最近的一个，复制标签
59
    tree = cKDTree(sampled[:, :3])
60
    _, nearest_idx = tree.query(raw_points[:, :3], k=1)
61
    final_labels = sampled_labels[nearest_idx]
62

63
    return final_labels
64

65

66
def main(args):
67
    # 加载模型
68
    from models.pointnet2_sem_seg import get_model
69
    model = get_model(num_class=args.num_class).cuda()
70

71
    checkpoint = torch.load(args.weights, map_location='cuda')
72
    if 'model_state_dict' in checkpoint:
73
        model.load_state_dict(checkpoint['model_state_dict'])
74
    else:
75
        model.load_state_dict(checkpoint)
76

77
    print(f"✅ 已加载语义分割模型: {args.weights}")
78
    print(f"   类别数: {args.num_class}, 采样点数: {args.num_point}\n")
79

80
    # 解析输入源
81
    if os.path.isfile(args.source):
82
        files = [args.source]
83
    else:
84
        files = [os.path.join(args.source, f) for f in os.listdir(args.source)
85
                 if f.endswith('.npy') or f.endswith('.txt')]
86
        files.sort()
87

88
    os.makedirs(args.save_dir, exist_ok=True)
89

90
    # 批量推理
91
    for fpath in files:
92
        # 加载原始点云
93
        if fpath.endswith('.npy'):
94
            raw_points = np.load(fpath).astype(np.float32)
95
        else:
96
            raw_points = np.loadtxt(fpath).astype(np.float32)
97

98
        # 推理
99
        labels = infer_semseg(model, raw_points, args.num_point, args.num_class)
100

101
        # 保存逐点标签
102
        fname = os.path.basename(fpath)
103
        out_name = fname.replace('.npy', '_pred.npy').replace('.txt', '_pred.npy')
104
        np.save(os.path.join(args.save_dir, out_name), labels)
105

106
        # 统计各类点数
107
        unique, counts = np.unique(labels, return_counts=True)
108
        stats = {SEG_CLASSES.get(u, f'class_{u}'): c for u, c in zip(unique, counts)}
109
        print(f"{fname:<30} -> {stats}")
110

111
    print(f"\n📄 所有逐点标签已保存至: {args.save_dir}/")
112

113

114
if __name__ == '__main__':
115
    parser = argparse.ArgumentParser(description='PointNet++ 语义分割推理（逐点分类）')
116
    parser.add_argument('--weights', type=str, required=True,
117
                        help='权重路径，如 log/sem_seg/pointnet2_sem_seg/best_model.pth')
118
    parser.add_argument('--source', type=str, required=True,
119
                        help='输入：单个 .npy/.txt 文件，或文件夹')
120
    parser.add_argument('--save_dir', type=str, default='seg_results',
121
                        help='输出标签保存目录')
122
    parser.add_argument('--num_point', type=int, default=4096,
123
                        help='网络固定输入点数（必须与训练时一致）')
124
    parser.add_argument('--num_class', type=int, default=4,
125
                        help='分割类别数')
126
    args = parser.parse_args()
127
    main(args)