import os
import numpy as np
import imageio.v3 as iio
from scipy.ndimage import grey_dilation

# 输入输出路径
input_dir = r"G:\python_ws_g\data\pcd2color_result\depth_tiff"
output_dir = r"G:\python_ws_g\data\pcd2color_result\dilated_tiff_filled"
os.makedirs(output_dir, exist_ok=True)

# 获取 .tiff 文件
image_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.tiff')]

# 定义膨胀核（结构元素）
kernel = np.ones((3, 3), dtype=bool)  # 小一点更精细填补

def fill_empty_by_dilation(depth, max_iter=50):
    # 判断有效像素（非0 且 非NaN）
    mask_valid = (depth > 0) & np.isfinite(depth)

    filled = depth.copy()
    for i in range(max_iter):
        prev = filled.copy()

        # 膨胀图像
        dilated = grey_dilation(filled, footprint=kernel)

        # 只用膨胀值去填空（原来为0的地方）
        filled = np.where(mask_valid, filled, dilated)

        # 更新 mask（哪些位置现在也有效了）
        new_mask_valid = (filled > 0) & np.isfinite(filled)
        if np.array_equal(new_mask_valid, mask_valid):
            print(f"→ 第{i+1}次迭代：无更多像素可填补，停止")
            break
        mask_valid = new_mask_valid

    return filled

for filename in image_files:
    try:
        img_path = os.path.join(input_dir, filename)
        depth = iio.imread(img_path)  # 支持 float32, uint16

        if depth is None:
            print(f"读取失败: {filename}")
            continue

        result = fill_empty_by_dilation(depth)

        # 保存结果
        save_path = os.path.join(output_dir, f"filled_{filename}")
        iio.imwrite(save_path, result.astype(depth.dtype))

        print(f"填补完成: {filename} → {save_path}")

    except Exception as e:
        print(f"处理失败: {filename} → {e}")