expansion.py 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. import os
  2. import numpy as np
  3. import imageio.v3 as iio
  4. from scipy.ndimage import grey_dilation
  5. # 输入输出路径
  6. input_dir = r"G:\python_ws_g\data\pcd2color_result\depth_tiff"
  7. output_dir = r"G:\python_ws_g\data\pcd2color_result\dilated_tiff_filled"
  8. os.makedirs(output_dir, exist_ok=True)
  9. # 获取 .tiff 文件
  10. image_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.tiff')]
  11. # 定义膨胀核(结构元素)
  12. kernel = np.ones((3, 3), dtype=bool) # 小一点更精细填补
  13. def fill_empty_by_dilation(depth, max_iter=50):
  14. # 判断有效像素(非0 且 非NaN)
  15. mask_valid = (depth > 0) & np.isfinite(depth)
  16. filled = depth.copy()
  17. for i in range(max_iter):
  18. prev = filled.copy()
  19. # 膨胀图像
  20. dilated = grey_dilation(filled, footprint=kernel)
  21. # 只用膨胀值去填空(原来为0的地方)
  22. filled = np.where(mask_valid, filled, dilated)
  23. # 更新 mask(哪些位置现在也有效了)
  24. new_mask_valid = (filled > 0) & np.isfinite(filled)
  25. if np.array_equal(new_mask_valid, mask_valid):
  26. print(f"→ 第{i+1}次迭代:无更多像素可填补,停止")
  27. break
  28. mask_valid = new_mask_valid
  29. return filled
  30. for filename in image_files:
  31. try:
  32. img_path = os.path.join(input_dir, filename)
  33. depth = iio.imread(img_path) # 支持 float32, uint16
  34. if depth is None:
  35. print(f"读取失败: {filename}")
  36. continue
  37. result = fill_empty_by_dilation(depth)
  38. # 保存结果
  39. save_path = os.path.join(output_dir, f"filled_{filename}")
  40. iio.imwrite(save_path, result.astype(depth.dtype))
  41. print(f"填补完成: {filename} → {save_path}")
  42. except Exception as e:
  43. print(f"处理失败: {filename} → {e}")