nii convert to 2D image【python】

可以自己精简，我的label是二分类

import SimpleITK as sitk
import cv2
from PIL import Image
import numpy as np
import nibabel as nib  # nii格式一般都会用到这个包
import imageio  # 转换成图像
import os

import numpy as np
from scipy.ndimage import rotate
from scipy.ndimage import median_filter
import matplotlib.pyplot as plt



xy = 128
vol1 = int(xy/2)
vol2 = int(xy/4)
vol3 = int(vol2+16)

def preprocess(image):
    result = median_filter(image, size=3)

    """# 添加高斯噪声
    noise = np.random.normal(0, 25, size=image.shape)
    noise_img = image + noise.astype('uint8')

    # 双边滤波
    result = cv2.bilateralFilter(noise_img, 9, 75, 75)

    # 显示图像
    cv2.imshow('src', image)
    cv2.imshow('noise', noise_img)
    cv2.imshow('result', result)
    cv2.waitKey()
    cv2.destroyAllWindows()"""

    return result


def create_dirs(out_path, num):
    for i in range(1, num):  # 这里需要注意，i取不到6，因为range()是前闭后开的，即i的取值范围为1-5。

        dir = os.path.join(r'crop/test4/ct/')
        # 前者为路径，后者为待创建文件夹的名称。注意，批量创建文件夹时不能有重复名称的，因此可以对文件夹加上序号信息。

        isExists = os.path.exists(dir)
        if not isExists:
            os.mkdir(dir)

def mask2d(input_path, output_folder, idx):
    # 加载NIfTI文件
    img = nib.load(input_path)
    data = img