import os
import numpy as np
import open3d as o3d
from tifffile import tifffile

# 相机内参矩阵
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 pointscloud2depthmap(points):
    """将点云投影为深度图"""
    point_image = np.zeros((height, width, 3), 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:
                point_image[v, u, :] = point
    return point_image

def process_pcd_folder(folder_path, output_folder, mode='f32'):
    """
    批量处理PCD文件，生成深度图并保存为TIFF格式。
    mode: 'f32' | 'int8' | 'all32'
    """
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    for filename in os.listdir(folder_path):
        if filename.lower().endswith('.pcd'):
            full_path = os.path.join(folder_path, filename)
            print(f"Processing: {full_path}")

            try:
                # 读取点云
                pcd = o3d.io.read_point_cloud(full_path)
                points = np.asarray(pcd.points)

                # 投影为图像
                point_image = pointscloud2depthmap(points)

                if mode == 'f32':
                    depth_map = point_image[:, :, 2].astype(np.float32)
                elif mode == 'int8':
                    depth_map = point_image[:, :, 2].astype(np.int8)
                elif mode == 'all32':
                    depth_map = point_image.astype(np.float32)
                else:
                    raise ValueError(f"Unsupported mode: {mode}")

                # 输出路径
                output_name = os.path.splitext(filename)[0] + ".tiff"
                output_path = os.path.join(output_folder, output_name)

                # 保存TIFF
                tifffile.imwrite(output_path, depth_map)
                print(f"Saved depth map to: {output_path}")
            except Exception as e:
                print(f"Error processing {filename}: {e}")

def main():
    input_folder = r"G:\python_ws_g\data\2025-05-22-11-40-20_LaserData_Hi221518(1)"

    # 定义各输出路径
    output_f32 = r"G:\python_ws_g\data\process_pcd_folder_f32"
    output_int8 = r"G:\python_ws_g\data\process_pcd_folder_int8"
    output_all32 = r"G:\python_ws_g\data\process_pcd_folder_all32"

    # 执行处理
    process_pcd_folder(input_folder, output_f32, mode='f32')
    process_pcd_folder(input_folder, output_int8, mode='int8')
    process_pcd_folder(input_folder, output_all32, mode='all32')

if __name__ == '__main__':
    main()