1
import os
2
import h5py # 用於讀取 .h5 文件
3
import numpy as np
4
from PIL import Image # 用於圖像處理和保存
5
from tqdm import tqdm # 用於顯示進度條
6
import glob # 用於查找文件
7

8
def export_images_from_h5(h5_dir, output_png_dir, target_resolution=(64, 64),
9
                          black_pixel_threshold_percentage=99.0): # 新增參數：黑色像素閾值
10
    """
11
    從指定目錄下的所有 .h5 文件中提取圖像數據，將其轉換為單通道灰度，
12
    調整尺寸為目標分辨率，然後保存為 PNG 圖片到一個共同的輸出目錄。
13
    同時會篩選掉黑色面積佔圖片達 99% 或以上的圖片。
14

15
    Args:
16
        h5_dir (str): 包含 .h5 文件的目錄路徑。
17
        output_png_dir (str): 圖片導出後保存的目錄路徑。
18
        target_resolution (tuple): 目標圖片分辨率，格式為 (寬度, 高度)，例如 (64, 64)。
19
        black_pixel_threshold_percentage (float): 圖片中黑色像素佔比達到或超過此值時，圖片將被跳過。
20
                                                    值範圍為 0 到 100。
21
    """
22
    print(f"--- 開始從 .h5 文件導出圖片 ---")
23
    print(f"源 .h5 文件目錄: {h5_dir}")
24
    print(f"圖片導出目錄: {output_png_dir}")
25
    print(f"目標圖片分辨率: {target_resolution[0]}x{target_resolution[1]} (單通道灰度)")
26
    print(f"黑色像素篩選閾值: >= {black_pixel_threshold_percentage:.2f}% 的黑色像素將被跳過")
27

28

29
    # 確保輸出目錄存在
30
    os.makedirs(output_png_dir, exist_ok=True)
31

32
    # 查找所有 .h5 文件
33
    h5_files = sorted(glob.glob(os.path.join(h5_dir, "*.h5")))
34

35
    if not h5_files:
36
        print(f"錯誤: 未在 '{h5_dir}' 中找到任何 .h5 文件！請檢查路徑或文件類型。")
37
        return
38

39
    print(f"找到 {len(h5_files)} 個 .h5 文件。")
40

41
    global_image_counter = 0 # 用於為所有導出的圖片生成唯一的文件名
42
    skipped_images_count = 0 # 記錄因黑色面積過大而被跳過的圖片數量
43

44
    # 設定黑色像素的閾值 (例如，像素值小於或等於 5 則視為黑色)
45
    # 您可以根據實際數據的特性調整這個值
46
    pixel_black_value_threshold = 5
47

48
    for h5_file_path in h5_files:
49
        h5_filename = os.path.basename(h5_file_path)
50
        print(f"\n--- 處理文件: {h5_filename} ---")
51

52
        try:
53
            with h5py.File(h5_file_path, 'r') as hf:
54
                # 假設圖片數據存儲在 'matrix' 鍵下
55
                if 'matrix' not in hf:
56
                    print(f"警告: 文件 '{h5_filename}' 中未找到 'matrix' 鍵，跳過此文件。")
57
                    continue
58

59
                images_data = hf['matrix'] # HDF5 dataset，通常為 (N, H, W, C)
60
                num_images_in_h5 = images_data.shape[0]
61

62
                print(f"文件 '{h5_filename}' 包含 {num_images_in_h5} 張圖片，原始形狀: {images_data.shape[1:]}")
63

64
                # 使用 tqdm 顯示進度
65
                for i in tqdm(range(num_images_in_h5), desc=f"導出 {h5_filename}"):
66
                    img_raw = images_data[i] # 獲取單張圖片數據 (H, W, C)
67

68
                    # 1. 將圖片數據轉換為 NumPy 陣列 (如果還不是的話)
69
                    # 並且處理數據類型，通常 PIL 期望 uint8 或 float
70

71
                    # 獲取原始通道數
72
                    original_channels = img_raw.shape[-1] if img_raw.ndim == 3 else 1
73

74
                    # 2. 轉換為單通道灰度圖
75
                    if original_channels >= 1: # 確保至少有一個通道
76
                        gray_img_np = img_raw[:, :, 0] # 假設第一個通道是適合用於灰度的
77

78
                        # 處理數據類型和範圍，以適應 PIL 的 'L' (灰度) 模式
79
                        if gray_img_np.max() - gray_img_np.min() > 1e-8: # 避免除以零
80
                            gray_img_np = (gray_img_np - gray_img_np.min()) / \
81
                                          (gray_img_np.max() - gray_img_np.min()) * 255.0
82
                        else: # 如果所有像素值都相同
83
                            gray_img_np = np.zeros_like(gray_img_np)
84

85
                        gray_img_np = np.clip(gray_img_np, 0, 255).astype(np.uint8) # 確保在 0-255 範圍並轉為 uint8
86

87
                        # 在調整尺寸和保存之前，檢查黑色面積比例
88
                        total_pixels = gray_img_np.size
89
                        black_pixels = np.sum(gray_img_np <= pixel_black_value_threshold)
90
                        black_pixel_ratio = (black_pixels / total_pixels) * 100
91

92
                        if black_pixel_ratio >= black_pixel_threshold_percentage:
93
                            skipped_images_count += 1
94
                            tqdm.write(f"跳過圖片: {h5_filename} 索引 {i}，黑色像素佔比 {black_pixel_ratio:.2f}%")
95
                            continue # 跳過當前圖片，不保存
96

97
                        img_pil = Image.fromarray(gray_img_np, 'L') # 'L' 表示 8-bit 灰度圖片
98
                    else:
99
                        print(f"\n警告: 圖片 {h5_filename} 索引 {i} 沒有足夠的通道來提取灰度圖，跳過。")
100
                        continue
101

102
                    # 3. 調整圖片尺寸
103
                    if img_pil.size != target_resolution:
104
                        img_pil = img_pil.resize(target_resolution, Image.LANCZOS) # 使用 Lanczos 高品質縮放
105

106
                    # 4. 保存為 PNG
107
                    output_filename = f"image_{global_image_counter:08d}.png" # 生成唯一文件名
108
                    output_file_path = os.path.join(output_png_dir, output_filename)
109
                    img_pil.save(output_file_path)
110

111
                    global_image_counter += 1
112

113
        except Exception as e:
114
            print(f"\n錯誤: 處理文件 '{h5_filename}' 時發生錯誤: {e}")
115
            continue # 繼續處理下一個 .h5 文件
116

117
    print(f"\n--- 所有 .h5 文件處理完成 ---")
118
    print(f"總共成功導出了 {global_image_counter} 張圖片到 '{output_png_dir}'。")
119
    print(f"總共跳過了 {skipped_images_count} 張圖片 (因黑色像素過多)。")
120

121
if __name__ == "__main__":
122
    # >>> 請在這裡修改為包含你的 TCIR .h5 文件的目錄路徑 <<<
123
    # 例如: h5_data_directory = "/home/chy/hbx/InverseBench_original/raw_tcir_h5"
124
    h5_data_directory = "./"
125

126
    # >>> 請在這裡修改為你希望導出圖片的目標目錄路徑 <<<
127
    # 例如: exported_png_directory = "/home/hbx/InverseBench_original/processed_tcir_png_64x64"
128
    exported_png_directory = "./tcir_png_64x64"
129

130
    # 目標分辨率，與你的訓練配置 (64x64) 相符
131
    target_res = (64, 64)
132

133
    # 設置黑色像素佔比的閾值 (例如 99%)
134
    black_pixel_percentage = 99.0
135

136
    # 檢查源目錄是否存在
137
    if not os.path.isdir(h5_data_directory):
138
        print(f"錯誤: 源目錄 '{h5_data_directory}' 不存在。請檢查路徑是否正確。")
139
    else:
140
        export_images_from_h5(h5_data_directory, exported_png_directory,
141
                              target_res, black_pixel_threshold_percentage=black_pixel_percentage)