python视差图和原图生成点云

import os
import numpy as np
import cv2
import open3d as o3d
from path import Path
from tqdm import tqdm


def write_point_cloud(ply_filename, points):
    formatted_points = []
    for point in points:
        formatted_points.append("%f %f %f %d %d %d 0\n" % (point[0], point[1], point[2], point[3], point[4], point[5]))

    out_file = open(ply_filename, "w")
    out_file.write('''ply
    format ascii 1.0
    element vertex %d
    property float x
    property float y
    property float z
    property uchar blue
    property uchar green
    property uchar red
    property uchar alpha
    end_header
    %s
    ''' % (len(points), "".join(formatted_points)))
    out_file.close()


def depth_image_to_point_cloud(rgb, depth, scale, K, pose):
    u = range(0, rgb.shape[1])
    v = range(0, rgb.shape[0])

    u, v = np.meshgrid(u, v)
    u = u.astype(float)
    v = v.astype(float)

    Z = depth.astype(float) / scale
    X = (u - K[0, 2]) * Z / K[0, 0]
    Y = (v - K[1, 2]) * Z / K[1, 1]

    X = np.ravel(X)
    Y = np.ravel(Y)
    Z = np.ravel(Z)

    #这里是直通滤波
    valid = Z > 0.01
    valid = Z < 0.15
    X = X[valid]
    Y = Y[valid]
    Z = Z[valid]

    position = np.vstack((X, Y, Z, np.ones(len(X))))
    position = np.dot(pose, position)

    R = np.ravel(rgb[:, :, 0])[valid]
    G = np.ravel(rgb[:, :, 1])[valid]
    B = np.ravel(rgb[:, :, 2])[valid]

    points = np.transpose(np.vstack((position[0:3, :], R, G, B))).tolist()

    return points


# image_files: XXXXXX.png (RGB, 24-bit, PNG)
# depth_files: XXXXXX.png (16-bit, PNG)
# poses: camera-to-world, 4×4 matrix in homogeneous coordinates
def build_point_cloud(dataset_path, scale, view_ply_in_world_coordinate):
    K = np.fromfile(os.path.join(dataset_path, "k2.txt"), dtype=float, sep="\n ")
    K = np.reshape(K, newshape=(3, 3))
    image_files = sorted(Path(os.path.join(dataset_path, "images")).files('*.bmp'))
    shicha_files = sorted(Path(os.path.join(dataset_path, "shicha")).files('*.raw'))
    f = 1402.218791
    b = 120.635947
    img_w, img_h = 1280, 960
    bf_value = b*f

    if view_ply_in_world_coordinate:
        poses = np.fromfile(os.path.join(dataset_path, "poses.txt"), dtype=float, sep="\n ")
        poses = np.reshape(poses, newshape=(-1, 4, 4))
    else:
        poses = np.eye(4)
    mark=1
    for i ,j in zip(tqdm(range(0, len(image_files))),tqdm(range(0, len(shicha_files)))):
        image_file = image_files[i]
        shicha_file = shicha_files[j]

        rgb = cv2.imread(image_file)
        shicha= np.fromfile(shicha_file, dtype=np.uint16)
        
        depth = np.where(shicha > 2, bf_value / shicha, shicha)       # 计算深度图
        depth = depth.reshape(img_h, img_w)

        if view_ply_in_world_coordinate:
            current_points_3D = depth_image_to_point_cloud(rgb, depth, scale=scale, K=K, pose=poses[i])
        else:
            current_points_3D = depth_image_to_point_cloud(rgb, depth, scale=scale, K=K, pose=poses) 
        #这里是保存.ply文件的文件名，注释掉的是生成和原图一样的名称，我是用的数字命名的 
        #save_ply_name = os.path.basename(os.path.splitext(image_files[i])[0]) + ".ply"
        save_ply_name = os.path.basename(str(mark) )+ ".ply"
        save_ply_path = os.path.join(dataset_path, "point_clouds")
        mark=mark+1
        if not os.path.exists(save_ply_path):  # 判断是否存在文件夹如果不存在则创建为文件夹
            os.mkdir(save_ply_path)
        write_point_cloud(os.path.join(save_ply_path, save_ply_name), current_points_3D)


if __name__ == '__main__':
    dataset_folder = Path("D:\\build_code+data\\depth2Cloud-main\\datatest")
    scene = Path("")
    # 如果view_ply_in_world_coordinate为True,那么点云的坐标就是在world坐标系下的坐标，否则就是在当前帧下的坐标
    view_ply_in_world_coordinate = False
    # 深度图对应的尺度因子，即深度图中存储的值与真实深度（单位为m）的比例, depth_map_value / real depth = scale_factor
    # 不同数据集对应的尺度因子不同，比如TUM的scale_factor为5000， hololens的数据的scale_factor为1000, Apollo Scape数据的scale_factor为200
    scale_factor = 1000.0
    build_point_cloud(os.path.join(dataset_folder, scene), scale_factor, view_ply_in_world_coordinate)

文件夹的命名参考这个链接https://zhuanlan.zhihu.com/p/360320545

记得修改相机内参参数