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, color_dir, depth_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) #xyz
    colors = np.asarray(pcd.colors) #rgb

    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)

    # 保存图像
    color_img_u8 = np.clip(color_img, 0, 255).astype(np.uint8)
    cv2.imwrite(os.path.join(color_dir, f"{filename}_color.jpg"), color_img_u8)

    depth_map_mm = (depth_map ).astype(np.float32)

    tif_write(os.path.join(depth_dir, f"{filename}_depth.tiff"), depth_map_mm)


    print(f"Saved color and depth images for {filename}")

# 设置输入输出路径
input_dir = r"G:\python_ws_g\data\2025-05-22-11-40-20_LaserData_Hi221518(1)"
output_base = r"G:\python_ws_g\data\pcd2color_result"
color_dir = os.path.join(output_base, "color_jpg")
depth_dir = os.path.join(output_base, "depth_tiff")

os.makedirs(color_dir, exist_ok=True)
os.makedirs(depth_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, color_dir, depth_dir)