import os
import cv2
import numpy as np
import open3d as o3d
from glob import glob
from collections import defaultdict

# 相机内参矩阵
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)
    point_count = np.zeros((height, width), dtype=np.int32)  # 统计每个像素覆盖的点数

    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]
                point_count[v, u] += 1  # 该像素多了一个点

    return color_image, point_count


def process_pcd_file(pcd_path, color_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, point_count = pointscloud2colorimg(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)

    # 统计覆盖情况
    num_multi_cover = np.sum(point_count > 1)      # 有多少像素点被多个点覆盖
    total_points = np.sum(point_count)            # 总点数
    max_cover = np.max(point_count)               # 单个像素最多的点数

    print(f"{filename}: 总点数={total_points}, 被多个点覆盖的像素数={num_multi_cover}, 单像素最多点数={max_cover}")


# 设置输入输出路径
input_dir = r"\\192.168.50.222\share\zqy\dianyun0731"
output_base = r"G:\python_ws_g\data\color"
color_dir = os.path.join(output_base, "color_jpg")

os.makedirs(color_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)
