import cv2
import numpy as np
import imageio
import open3d as o3d
from tifffile import tifffile
# 相机内参矩阵 [fx, 0, cx; 0, fy, cy; 0, 0, 1]
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]

def pointscloud2depthmap(points):
    # 初始化一个空的目标数组
    point_image = np.zeros((height, width, 3), dtype=np.float32)
    # 遍历点云中的每个点，进行投影并填充目标数组
    for point in points:
        X, Y, Z ,r,g,b= point
        if Z > 0:  # 确保Z值有效
            # 计算2D图像坐标
            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 pointscloud2colorimg(points):
    # 初始化一个空的目标数组
    point_image = np.zeros((height, width, 3), dtype=np.float32)
    color_image = np.zeros((height, width, 3), dtype=np.float32)
    # 遍历点云中的每个点，进行投影并填充目标数组
    for point in points:
        X, Y, Z ,r,g,b= point
        if Z > 0:  # 确保Z值有效
            # 计算2D图像坐标
            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
                color_image[v,u]=[r*255,g*255,b*255]

    return color_image



# 加载PCD文件
pcd = o3d.io.read_point_cloud(r"F:\test_pointcloud\color.pcd")

# 打印点的数量
print("Number of points:", len(pcd.points))

# 获取点云数据
points = np.asarray(pcd.points)

colors=np.asarray(pcd.colors)
points = np.hstack((points, colors))  # (N, 6)

# 打印前5个点的坐标
print("First 5 points:\n", points[:5])
print(f'color:{colors}')

# 目标深度图尺寸
height, width = 2000, 2000

point_image=pointscloud2colorimg(points)
# 打印结果以验证
print("Shape of the projected point cloud:", point_image.shape)
print("First few pixels (if any):", point_image[:5, :5, :])

# 提取 Z 值作为深度图
depth_map = point_image[:, :, 2]

# 打印深度图的一些信息以验证
print("Depth map shape:", depth_map.shape)
print("Depth map dtype:", depth_map.dtype)
print("Min depth value:", np.min(depth_map))
print("Max depth value:", np.max(depth_map))

# 保存为 TIFF 文件
color_img_path='color_img.jpg'

cv2.imwrite(color_img_path,point_image)
