图像形态学操作（cv2库实现）

#coding:utf-8
import SimpleITK as sitk
import numpy as np
import cv2
# 膨胀
def dilateion(image):
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5,5))
    # dilate = cv2.dilate(image, kernel, iterations=1)
    dilate = cv2.morphologyEx(image, cv2.MORPH_DILATE, kernel)
    return dilate
def erode(image):
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3))
    # erode = cv2.erode(image, kernel, iterations=1)
    erode = cv2.morphologyEx(image, cv2.MORPH_ERODE, kernel)
    return erode
# 形态学梯度
def edge(image):
    SE = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    img_grad = cv2.morphologyEx(image, cv2.MORPH_GRADIENT, SE)
    return img_grad
# 开运算
def openOpreation(image):
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    open = cv2.morphologyEx(image, cv2.MORPH_OPEN, kernel)
    return open
# 闭运算
def closeOperation(image):
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    close = cv2.morphologyEx(image, cv2.MORPH_CLOSE, kernel)
    return close
# 读取3D图像，对每一个slice进行形态学变化
def read_3D(image, model="closeOperation"):
    slices = image.shape[0]
    result = np.zeros(img_num.shape)
    for i in range(slices):
        sli = img_num[i:i + 1, ...]
        s = sli[0, ...]
        if model == "dilateion":
            slice = dilateion(s)
        elif model =="edge":
            slice = edge(s)
        elif model == "erode":
            slice = erode(s)
        elif model == "edge":
            slice = edge(s)
        elif model == "openOpreation":
            slice = openOpreation(s)
        elif model == "closeOperation":
            slice = closeOperation(s)
        result[i, ...] = slice
    return result
# 保存图像
def save(x, path):
    predict_seg = sitk.GetImageFromArray(x)
    sitk.WriteImage(predict_seg, path)
# 读取 nii文件
def read_nii(path):
    image = sitk.ReadImage(path)
    img_num = sitk.GetArrayFromImage(image)
    return img_num
if __name__ == "__main__":
    path = r"D:\myProject\HDC_vessel_seg\datasets\nii\image\vessel_12.nii"
    imgpath = r"D:\myProject\HDC_vessel_seg\datasets\nii\image\image_12.nii"
    img_num = read_nii(path)
    img = read_nii(imgpath)
    img_num = img_num[10:260,...]
    img = img[10:260,...]
    result = read_3D(img_num)
    name = "closeOperation_"
    save(result, path.replace("vessel_12", name + "vessel"))
    # save(img_num, path.replace("vessel_12", "pre_vessel"))
    # save(img, imgpath.replace("image_12", "pre_image"))