SimpleITK使用sobel算子提取三维nii医学图像的梯度图

最近要做一个三维的剂量预测任务，准备采用多任务学习的方法，加入一个支路来预测梯度信息，所以需要先手动使用sobel算子提取一下梯度图，来作为此支路的label。但是通用的函数往往只能用于二维图像，例如cv2.Sobel()。

后来终于找到了一种可以方便的直接使用SimpleITK库提取nii图像的三维梯度图的方法，在此做一记录。代码实现如下：

import SimpleITK as sitk

path = '201_dose.nii.gz'
sitk_img = sitk.ReadImage(path)

# 保存文件信息
origin = sitk_img.GetOrigin()
spacing = sitk_img.GetSpacing()
direction = sitk_img.GetDirection()
# 转换数据类型
data = sitk.Cast(sitk_img, sitk.sitkFloat32)

# 提取梯度图
gra_filter = sitk.SobelEdgeDetectionImageFilter()
sobel_data = gra_filter.Execute(data)
sobel_data = sitk.Cast(sobel_data, sitk.sitkInt16)

# 保存结果
sobel_data.SetOrigin(origin)
sobel_data.SetSpacing(spacing)
sobel_data.SetDirection(direction)
sitk.WriteImage(sobel_data, 'sobel.nii.gz')