Visualization.py

import numpy as np
import open3d as o3d
from skimage import io


# region 功能函数
def read_img(filename, dtype=np.uint8):
    """ 读取图片为矩阵, RGBA/RGB/GRAY """
    arr = io.imread(filename)
    arr = arr.astype(dtype)
    return arr

def z_val_to_pts_3d(z_val: np.ndarray, x_pitch=1., y_pitch=1., voxel_size=None, mask_arr=None, feature_arr=None):
    """
    z 矩阵转点集
    :param z_val:
    :param x_pitch: X 间隙，纵向
    :param y_pitch: y 间隙，横向
    :param voxel_size: 减采样间隔
    :param mask_arr: 掩膜
    :param feature_arr: 特征矩阵 z_val.shape * n
    :return:
    """
    if voxel_size is None:
        voxel_size = [x_pitch, y_pitch]
    if isinstance(voxel_size, float) or isinstance(voxel_size, int):
        voxel_size = [voxel_size, voxel_size]
    x_scale = int(round(voxel_size[0] / x_pitch))
    y_scale = int(round(voxel_size[1] / y_pitch))
    z_val = z_val[::x_scale, ::y_scale]
    x_pitch, y_pitch = x_pitch * x_scale, y_pitch * y_scale
    # 输出点集
    x_arr, y_arr = np.meshgrid(range(z_val.shape[0]), range(z_val.shape[1]), indexing='ij')
    x_arr, y_arr = x_arr * x_pitch, y_arr * y_pitch
    x_arr, y_arr, z_arr = x_arr.reshape(-1, 1), y_arr.reshape(-1, 1), z_val.reshape(-1, 1)
    pts_3d = np.hstack((x_arr, y_arr, z_arr))
    if mask_arr is not None:
        mask_arr = mask_arr[::x_scale, ::y_scale]
        mask_arr = mask_arr.flatten()
        pts_3d = pts_3d[mask_arr > 0, :]
    # 添加 feature_arr 列
    if feature_arr is not None:
        feature_arr = feature_arr[::x_scale, ::y_scale]
        if len(feature_arr.shape) > 2:
            feature_arr = feature_arr.reshape(-1, feature_arr.shape[-1])
        else:
            feature_arr = feature_arr.reshape(-1, 1)
        pts_3d = np.hstack((pts_3d, feature_arr))
    return pts_3d

def ndarray_to_o3d_cloud(points):
    """
    ndarray 格式的点集转化为 O3D 的点云
    :param points: ndarray [[xyz]]
    :return: o3d 点云
    """
    o3d_point_cloud = o3d.utility.Vector3dVector(points)
    o3d_point_cloud = o3d.cpu.pybind.geometry.PointCloud(o3d_point_cloud)
    return o3d_point_cloud

def draw_in_o3d(pts, is_2d=False, color_arr=None, estimate_normals=False):
    """
    绘制点集
    :param pts: 点集
    :param is_2d: 是否 2d
    :param color_arr: 点集颜色
    :param estimate_normals: 是否 估计法向量
    :return:
    """
    if is_2d:
        tmp_z = np.zeros((len(pts), 1))
        pts = np.hstack((pts, tmp_z))
    pts = ndarray_to_o3d_cloud(pts)
    if color_arr is not None:
        pts.colors = o3d.utility.Vector3dVector(color_arr)
    if estimate_normals:
        pts.estimate_normals()
    o3d.visualization.draw_geometries([pts])

# endregion

#文件路径
tif_path = r"D:\point-cloud-projects\Seam\13_3.tif"

if __name__ == '__main__':

    #加载tif文件
    height_map = read_img(tif_path, np.float32)

    # 创建点云对象
    pts_3d = z_val_to_pts_3d(height_map, 0.1, 0.025)
    draw_in_o3d(pts_3d, estimate_normals=True)