存储为jpg检查影像mask是否正确

def img_with_mask_save():
    # 自定义渐变透明度alpha的color map
    from skimage import exposure
    from matplotlib.colors import LinearSegmentedColormap
    # get colormap
    nalphas = 256
    color_array = plt.get_cmap('Reds')(range(nalphas))
    # change alpha values
    color_array[:, -1] = np.linspace(0, 1, nalphas)
    color_array[-1,:] = [1, 0, 0, 1]
    # create a colormap object
    Reds_alpha_objects = LinearSegmentedColormap.from_list(name='Reds_alpha', colors=color_array)
    # register this new colormap with matplotlib
    plt.register_cmap(cmap=Reds_alpha_objects)

    img_dir = os.path.join(os.getcwd(), 'img')
    mask_dir = os.path.join(os.getcwd(), 'mask')
    for nii in os.listdir(img_dir):
        if not nii.endswith('.nii'):
            continue
        img_path = os.path.join(img_dir, nii)
        mask_path = os.path.join(mask_dir, nii.split('_')[0]+'_mask.nii')

        img = sitk.ReadImage(img_path)       # type: sitk.Image
        img = sitk.GetArrayFromImage(img)    # type: np.ndarray
        img = exposure.equalize_hist(img)    # 直方图均衡化
        
        mask = sitk.ReadImage(mask_path)     # type: sitk.Image
        mask = sitk.GetArrayFromImage(mask)  # type: np.ndarray

        intensities = np.sum(mask, axis=(1, 2))
        slice_index = np.where(intensities == np.max(intensities))[0][0]  # ROI区域最大的一层的索引

        plt.imshow(img[slice_index, :, :], cmap='gray')
        plt.imshow(mask[slice_index, :, :], alpha=.3, cmap='Reds_alpha')
        # plt.colorbar()
        plt.tight_layout()
        output_path = os.path.join(os.getcwd(), 'jpg', nii.split('_')[0]+'.jpg')
        plt.savefig(output_path)
        print(output_path, 'ok')