[小技巧]把.nii.gz文件从3D图像变成2D图像

问题

很多时候3D网络由于参数量大，不好训练，一般转换成2D图像使用2D网络进行训练，当然也会有比如切片上下文信息丢失的问题，抛开不谈。（2D->2.5D）在这里插入代码片

代码

下面是读取.nii.gz文件，然后进行窗口窗位变换，得到变换后的数值，然后在保存为2D图片。

import SimpleITK as sitk
import os
import nibabel as nib
import numpy as np
import matplotlib.pyplot as plt
import pydicom

def load_nifti(file_path):
    img = nib.load(file_path)
    img_data = img.get_fdata()
    return img_data

def adjust_window(image, window_center, window_width):
    image_min = window_center - window_width / 2.0
    image_max = window_center + window_width / 2.0
    adjusted_image = np.clip(image, image_min, image_max)
    adjusted_image = (adjusted_image - image_min) / (image_max - image_min)
    return adjusted_image

def save_slices_as_png(image_path, output_folder):
    # Load the NIfTI image
    image = sitk.ReadImage(image_path)
    # mask = sitk.ReadImage(mask_path)

    # 获取数组表示
    image_array = sitk.GetArrayFromImage(image)
    # mask_array = sitk.GetArrayFromImage(mask)

    # Extract file name without extension
    file_name = os.path.basename(image_path).split(".")[0]

    # Create output folder if it doesn't exist
    os.makedirs(output_folder, exist_ok=True)

    # Iterate through slices and save as .png
    for i in range(image_array.shape[0]):
        # Get the ith slice
        slice_data = image_array[i, :, :]

        # Adjust CT window
        adjusted_slice = adjust_window(slice_data, window_center, window_width)

        # Save as .png
        png_path = os.path.join(output_folder, f"{file_name}_{i}.png")
        plt.imsave(png_path, adjusted_slice, cmap='gray')

if __name__ == "__main__":

    data_folder = r"D:\Download\data\data"

    # Set the CT window parameters
    window_center = 30
    window_width = 400

    # Process image folder
    image_folder = os.path.join(data_folder, "images")
    mask_folder = os.path.join(data_folder, "masks")
    for image_file in os.listdir(mask_folder):
        image_path = os.path.join(image_folder, image_file)
        save_slices_as_png(image_path, r"D:/Download/data/data/2D/images")

    # Process mask folder
    
    for mask_file in os.listdir(mask_folder):
        mask_path = os.path.join(mask_folder, mask_file)
        save_slices_as_png(mask_path, r"D:/Download/data/data/2D/masks")

可以看到很多图片对应的mask文件是全黑的，说明其中没有病灶区域，为了除去这些图片，可以用下面这段代码。

import os
import cv2
import numpy as np

def is_black_mask(mask_array):
    return np.all(mask_array == 0)

def remove_images_with_black_masks(image_folder, mask_folder):
    for image_file in os.listdir(image_folder):
        image_path = os.path.join(image_folder, image_file)
        mask_path = os.path.join(mask_folder, image_file)

        # 读取标签
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

        # 检查标签是否全是黑色
        if is_black_mask(mask):
            # 删除对应的图像和标签
            os.remove(image_path)
            os.remove(mask_path)

if __name__ == "__main__":
    image_folder = "D:/Download/data/data/2D/images"
    mask_folder = "D:/Download/data/data/2D/masks"

    remove_images_with_black_masks(image_folder, mask_folder)