pokouqiege/utils/data_process/zyh/rgbd/pcd2rgbd.py

import os
import cv2
import numpy as np
import open3d as o3d
from glob import glob
from tifffile import imwrite as tif_write

# 相机内参矩阵
K = np.array([
    [1.30449e3, 0,         5.2602e2],
    [0,         1.30449e3, 1.07432e3],
    [0,         0,         1]
])
fx, fy = K[0, 0], K[1, 1]
cx, cy = K[0, 2], K[1, 2]

# 图像尺寸
height, width = 2000, 2000

def pointscloud2colorimg(points):
    color_image = np.zeros((height, width, 3), dtype=np.float32)
    for point in points:
        X, Y, Z, r, g, b = point
        if Z > 0:
            u = int((X * fx) / Z + cx)
            v = int((Y * fy) / Z + cy)
            if 0 <= u < width and 0 <= v < height:
                color_image[v, u] = [r * 255, g * 255, b * 255]
    return color_image

def pointscloud2depthmap(points):
    depth_map = np.zeros((height, width), dtype=np.float32)
    for point in points:
        X, Y, Z, *_ = point
        if Z > 0:
            u = int((X * fx) / Z + cx)
            v = int((Y * fy) / Z + cy)
            if 0 <= u < width and 0 <= v < height:
                depth_map[v, u] = Z
    return depth_map

def process_pcd_file(pcd_path, rgbd_dir):
    filename = os.path.splitext(os.path.basename(pcd_path))[0]
    print(f"Processing: {filename}")

    pcd = o3d.io.read_point_cloud(pcd_path)
    points = np.asarray(pcd.points)
    colors = np.asarray(pcd.colors)

    if points.shape[0] == 0 or colors.shape[0] == 0:
        print(f"Skipping empty or invalid PCD: {filename}")
        return

    points_colored = np.hstack((points, colors))

    color_img = pointscloud2colorimg(points_colored)
    depth_map = pointscloud2depthmap(points_colored)

    # RGB 转 uint8，Depth 保持 float32
    color_img_u8 = np.clip(color_img, 0, 255).astype(np.uint8)
    depth_map_mm = depth_map.astype(np.float32)

    # 合并为 4 通道 RGBD 图像（float32）
    rgbd = np.concatenate((color_img_u8, depth_map_mm[..., np.newaxis]), axis=2)
    tif_write(os.path.join(rgbd_dir, f"{filename}_rgbd.tiff"), rgbd, photometric='rgb')

    print(f"Saved RGBD TIFF for {filename}")

input_dir = r"G:\python_ws_g\data\pcd"
output_base = os.path.dirname(input_dir)

rgbd_dir = os.path.join(output_base, "rgbd_tiff")
os.makedirs(rgbd_dir, exist_ok=True)

# 获取 PCD 文件
pcd_files = glob(os.path.join(input_dir, "*.pcd"))
print(f"Found {len(pcd_files)} PCD files.")

# 批量处理
for pcd_path in pcd_files:
    process_pcd_file(pcd_path, rgbd_dir)