import numpy as np
import matplotlib.pyplot as plt

def generate_scanline_image(points, image_shape):
    """
    将点云投影数据转换为扫描线图像。
    
    参数：
    - points: (N, 3) 的点云投影数据，包含 [x, y, value]，分别为横坐标、纵坐标和深度值/颜色值。
    - image_shape: 输出图像的形状 (height, width)。
    
    返回：
    - scanline_image: 生成的扫描线图像。
    """
    # 初始化图像
    scanline_image = np.zeros(image_shape)
    
    # 将点按 y 坐标分组（行）
    unique_rows = np.unique(points[:, 1])
    for y in unique_rows:
        # 获取当前行的所有点
        row_points = points[points[:, 1] == y]
        
        # 按 x 坐标排序
        row_points = row_points[np.argsort(row_points[:, 0])]
        
        # 提取 x 和 value 值
        x_coords = row_points[:, 0].astype(int)
        values = row_points[:, 2]
        
        # 插值填充到整行
        full_x = np.arange(x_coords.min(), x_coords.max() + 1)
        full_values = np.interp(full_x, x_coords, values)
        
        # 将插值后的值写入图像
        scanline_image[int(y), full_x] = full_values
    
    return scanline_image

# 示例用法
if __name__ == "__main__":
    # 生成示例点云投影数据 [x, y, value]
    np.random.seed(42)
    points = np.array([
        [x, y, np.sin(x / 10) * 100] for y in range(100, 200, 10) for x in np.random.randint(0, 500, 30)
    ])
    
    # 图像形状
    image_shape = (300, 500)
    
    # 生成扫描线图像
    scanline_image = generate_scanline_image(points, image_shape)
    
    # 显示图像
    plt.figure(figsize=(10, 6))
    plt.imshow(scanline_image, cmap='viridis')
    plt.title("Generated Scanline Image")
    plt.colorbar(label="Value")
    plt.show()

import cv2
import numpy as np

def patchmatch_stereo(left_img, right_img, patch_size=5, max_disparity=64, iterations=5):
    """
    实现一个简单版本的 PatchMatch Stereo 算法
    
    :param left_img: 左图像 (灰度图像)
    :param right_img: 右图像 (灰度图像)
    :param patch_size: 匹配窗口的大小
    :param max_disparity: 最大视差范围
    :param iterations: PatchMatch 的迭代次数
    :return: 视差图
    """

    # 获取图像的高度和宽度
    H, W = left_img.shape
    disparity_map = np.zeros_like(left_img, dtype=np.float32)
    
    # 初始化视差图（随机初始化视差值）
    for y in range(H):
        for x in range(W):
            disparity_map[y, x] = np.random.randint(0, max_disparity)
    
    # 匹配操作
    for it in range(iterations):
        # 对每个像素进行 PatchMatch
        for y in range(H):
            for x in range(W):
                best_disparity = disparity_map[y, x]
                best_cost = float('inf')
                
                # 遍历视差范围
                for d in range(-max_disparity, max_disparity + 1):
                    if 0 <= x + d < W:
                        # 获取左图和右图对应窗口的块
                        left_patch = left_img[y:y+patch_size, x:x+patch_size]
                        right_patch = right_img[y:y+patch_size, (x + d): (x + d + patch_size)]
                        
                        # 计算块之间的差异（可以使用其他度量，比如 SSIM 或其他）
                        cost = np.sum((left_patch - right_patch) ** 2)
                        
                        if cost < best_cost:
                            best_cost = cost
                            best_disparity = d
                
                disparity_map[y, x] = best_disparity
        
        # 传播优化：更新邻域像素的视差
        for y in range(H):
            for x in range(W):
                # 可以进行视差传播（例如水平、垂直方向传播）
                # 基于已匹配的邻域值做优化
                # 简单的传播方法可以使用简单平均
                if x > 0:
                    disparity_map[y, x] = (disparity_map[y, x] + disparity_map[y, x-1]) / 2
                if y > 0:
                    disparity_map[y, x] = (disparity_map[y, x] + disparity_map[y-1, x]) / 2

    return disparity_map

# 载入图像
left_img = cv2.imread('left_image.png', cv2.IMREAD_GRAYSCALE)
right_img = cv2.imread('right_image.png', cv2.IMREAD_GRAYSCALE)

# 确保图像大小一致
assert left_img.shape == right_img.shape, "左右图像大小不一致"

# 运行 PatchMatch Stereo
disparity_map = patchmatch_stereo(left_img, right_img)

# 显示视差图
cv2.imshow('Disparity Map', disparity_map / np.max(disparity_map))  # 归一化显示
cv2.waitKey(0)
cv2.destroyAllWindows()