Python 重命名ply格式的点云特征, 修改点云特征变量类型

依赖 Python ply pcd 点云读取源码 中的ply.py

"""
rename_ply.py
重命名点云特征, 修改点云特征变量类型
github.com/DLW3D
2022年3月29日
"""
import sys
import numpy as np
from numpy import dtype
from ply import read_ply, write_ply

if __name__ == '__main__':
    # 获取命令行参数
    if len(sys.argv) == 2:
        file_path = sys.argv[1]
        save_path = sys.argv[1]
    elif len(sys.argv) == 3:
        file_path = sys.argv[1]
        save_path = sys.argv[2]
    else:
        print('Error: input args error')
        print('Sample: python rotate_pc.py in_pc.ply [out_pc.ply]')
        sys.exit(2)

    data = read_ply(file_path)
    header = list(data.dtype.names)
    typer = [data.dtype[i] for i in header]

    # 全局偏移
    # shift_x = data['x'].min() // 100000 * 100000
    # shift_y = data['y'].min() // 100000 * 100000
    # data['x'] -= shift_x
    # data['y'] -= shift_y
    # # 将shift记录在文件中
    # if bool(shift_x) or bool(shift_y):
    #     with open(save_path + '.shift.txt', 'w') as f:
    #         f.write('%d %d' % (shift_x, shift_y))
    #     print('shift:', shift_x, shift_y)

    # 修改变量名和类别
    name_map = {
        'scalar_Normal_change_rate_(0.4)': 'curvature',
        'Intensity': 'intensity',
        'ReturnNumber': 'return_number',
        'classification': 'label',
        'ScanAngleRank': 'angle',
    }
    name_type_map = {
        'intensity': dtype('uint16'),
        'return_number': dtype('uint8'),
        'angle': dtype('int8'),
        'label': dtype('int8'),
        'pred': dtype('int8'),
        'error': dtype('int8'),
    }

    for i, n in enumerate(header):
        if n.startswith('scalar_'):
            header[i] = n[7:]
            n = header[i]
        if n in name_map:
            header[i] = name_map[n]
        if n in name_type_map:
            typer[i] = name_type_map[n]

    # 构建新Array
    new_type = np.dtype(list(zip(header, typer)))
    new_data = np.empty(len(data), new_type)
    for new_name, old_name in zip(header, data.dtype.names):
        new_data[new_name] = data[old_name]

    # 修复曲率中的nan
    if 'curvature' in new_data.dtype.names and np.isnan(new_data['curvature']).any():
        print('修复曲率')
        new_data['curvature'][np.isnan(new_data['curvature'])] = new_data['curvature'][~np.isnan(new_data['curvature'])].max()

    # 保存
    write_ply(save_path, [new_data[h] for h in new_data.dtype.names], new_data.dtype.names)
    print('save in %s' % save_path)