医学影像分割代码

44 阅读4分钟

第一部分:图像预处理

import numpy as np
import cv2
import acl

# 定义图像处理函数
def imgProcess(imgPath):
    # 使用OpenCV读取图像文件
    image = cv2.imread(imgPath)
    # 将图像尺寸调整为352x352像素,使用区域插值法
    image = cv2.resize(image,(352, 352), interpolation = cv2.INTER_AREA)
    # 将图像从BGR颜色空间转换为YUV颜色空间
    yuvImg = cv2.cvtColor(image, cv2.COLOR_BGR2YUV)
    # 分离YUV图像的三个通道
    y, u, v = cv2.split(yuvImg)
    # 获取图像的高度和宽度
    height, width = yuvImg.shape[:2]
    # 将U通道尺寸缩小为原图的一半,使用线性插值法
    u = cv2.resize(u, (width // 2, height // 2), interpolation=cv2.INTER_LINEAR)
    # 将V通道尺寸缩小为原图的一半,使用线性插值法
    v = cv2.resize(v, (width // 2, height // 2), interpolation=cv2.INTER_LINEAR)
    # 创建一个空的UV合并矩阵,高度和宽度为原图的一半
    uv = np.zeros((height // 2, width), dtype=yuvImg.dtype)
    # 将U通道数据填充到UV矩阵的偶数列
    uv[:, 0::2] = u
    # 将V通道数据填充到UV矩阵的奇数列
    uv[:, 1::2] = v
    # 将Y通道和UV矩阵在垂直方向拼接
    image = np.concatenate((y, uv),  axis = 0)

# 定义源图像路径变量
source = ""
# 调用图像处理函数处理源图像
processData = imgProcess("source")
# 打印处理后的数据形状
print(processData.shape)

第二部分:模型定义和推理

# 定义ACL内存分配标志:优先分配大页内存
ACL_MEM_MALLOC_HUGE_FIRST = 0
# 定义ACL内存拷贝方向:从主机到设备
ACL_MEMCPY_HOST_TO_DEVICE = 1
# 定义ACL内存拷贝方向:从设备到主机
ACL_MEMCPY_DEVICE_TO_HOST = 2

# 定义OmModel类,用于管理离线模型
class OmModel:
    # 初始化方法,接收模型路径参数
    def __init__(self, modelPath):
        # 设置设备ID
        self.deviceId = 5
        # 初始化ACL运行时环境
        ret = acl.init()
        # 设置运行设备
        ret = acl.rt.set_device(self.deviceId)

        # 从文件加载模型,获取模型ID
        self.modelId, ret = acl.mdl.load_from_file(modelPath)
        # 创建模型描述符
        self.modelDesc = acl.mdl.create_desc()
        # 获取模型描述信息
        ret = acl.mdl.get_desc(self.modelDesc, self.modelId)

        # 准备输入数据集和输入数据
        self.inputDatasets, self.inputData = self.prepareDatasets("input")
        # 准备输出数据集和输出数据
        self.outputDatasets, self.outputData = self.prepareDatasets("output")

    # 准备数据集方法,接收IO类型参数
    def prepareDataset(self, ioType):
        # 判断是否为输入类型
        if ioType == "input":
            # 获取输入数量
            ioNum = acl.mdl.get_num_inputs(self.modelDesc)
            # 获取输入大小的函数引用
            aclMdlGetSizeByIndex = acl.mdl.get_input_size_by_index
        else:
            # 获取输出数量
            ioNum = acl.mdl.get_num_outputs(self.modelDesc)
            # 获取输出大小的函数引用
            aclMdlGetSizeByIndex = acl.mdl.get_output_size_by_index
        
        # 创建数据集
        dataset = acl.mdl.create_dataset()
        # 初始化数据列表
        datas = []

        # 遍历所有IO
        for i in range(ioNum):
            # 获取缓冲区大小
            bufferSize = aclMdlGetSizeByIndex(self.modelDesc, i)
            # 分配设备内存
            buffer, ret = acl.rt.malloc(bufferSize, ACL_MEM_MALLOC_HUGE_FIRST)
            # 创建数据缓冲区
            dataBuffer = acl.create_data_buffer(buffer, bufferSize)
            # 将数据缓冲区添加到数据集
            ret = acl.mdl.add_dataset_buffer(dataset, dataBuffer)
            # 将缓冲区信息添加到数据列表
            datas.append({"buffer": buffer, "data": dataBuffer, "size": bufferSize})
            # 返回数据集和数据列表
            return dataset, datas
        
    # 模型推理方法,接收输入数据
    def forward(self, inputs):
        # 获取输入数量
        inputNum = len(inputs)
        # 遍历所有输入
        for i in range(inputNum):
            # 将输入数据转换为字节格式
            bytesData = inputs[i].tobytes()
            # 将字节数据转换为指针
            bytesPtr = acl.util.bytes_to_ptr(bytesData)
            # 将数据从主机内存拷贝到设备内存
            ret = acl.rt.memcpy(self.inputData[i]["buffer"],
                                self.inputData[i]["size"],
                                bytesPtr,
                                len(bytesData),
                                ACL_MEMCPY_HOST_TO_DEVICE)
        # 执行模型推理
        ret = acl.mdl.execute(self.modelId, self.inputDatasets, self.outputDatasets)
        # 初始化推理结果列表
        inferenceResults = []

        # 遍历所有输出数据
        for i, item in enumerate(self.outputData):
            # 分配主机内存
            bufferHost, ret = acl.rt.malloc_host(self.outputData[i]["size"])
            # 将数据从设备内存拷贝到主机内存
            ret = acl.rt.memcpy(bufferHost,
                                self.outputData[i]["size"],
                                self.outputData[i]["buffer"],
                                self.outputData[i]["size"],
                                ACL_MEMCPY_DEVICE_TO_HOST
                                )  
            # 将指针转换为字节数据
            bytesOut = acl.util.ptr_to_bytes(bufferHost, self.outputData[i]["size"])
            # 从字节数据创建NumPy数组
            data = np.frombuffer(bytesOut, dtype = np.float32)
            # 将数据添加到推理结果列表
            inferenceResults.append(data)  
            # 返回推理结果
            return inferenceResults
        
    # 析构方法,用于资源清理
    def __del__(self):
        # 遍历输入和输出数据集
        for dataset in [self.inputDatasets, self.outputDatasets]:
            # 当数据集不为空时循环
            while dataset:
                # 从数据集弹出一个项
                item = dataset.pop()
                # 销毁数据缓冲区
                ret = acl.destroy_data_buffer(item["data"])
                # 释放设备内存
                ret = acl.rt.free(item["buffer"])
        # 销毁输入数据集
        ret = acl.mdl.destroy_dataset(self.inputDataset)
        # 销毁输出数据集
        ret = acl.mdl.destroy_dataset(self.outputDataset)
        # 销毁模型描述符
        ret = acl.destroy_desc(self.modelDesc)
        # 卸载模型
        ret = acl.mdl.unload(self.modelId)
        # 重置设备
        ret = acl.rt.reset_device(self.deviceId)
        # 最终化ACL运行时环境
        ret = acl.finalize()

第三部分:主程序流程

# 定义模型文件路径
pranetOmModelPath = "./PraNet-19_bs1.om"
# 创建OmModel实例
pranetOmModel = OmModel(pranetOmModelPath)

# 执行模型推理,获取输出结果
outputRes = pranetOmModel.forward([processData])
# 将输出结果重塑为指定形状
tensorRes = outputRes[0].reshape(1, 1, 352, 352)
# 打印重塑后的张量形状
print(tensorRes.shape)

第四部分:后处理函数

# 定义分割阈值
THRESHOLD = 0.5   
# 定义分割结果解码函数
def decodeSegMap(labelMask, savePath):
    # 获取分割结果
    segmentResult = labelMask[0][0]
    # 创建分割映射矩阵
    segMap = np.zeros((segmentResult.shape[0], segmentResult.shape[1]), dtype=np.float32)
    # 根据阈值将分割结果二值化
    segMap[segmentResult > THRESHOLD] = 255
    # 保存分割结果图像
    cv2.imwrite(savePath, segMap)
    return

# 定义对比输出函数
def enableContrastOutput(arr):
    # 打开第一张图像
    img1 = Image.open(arr[0])
    # 打开第二张图像
    img2 = Image.open(arr[1])
    # 将第二张图像调整为第一张图像的尺寸
    img2 = img2.resize(img1.size, Image.NEAREST)
 
    # 创建新的空白图像用于对比显示
    toImage = Image.new('RGB', (img1.width + img2.width + 35, img2.height), 'white')
    # 粘贴第一张图像到左侧
    toImage.paste(img1, (0, 0))
    # 粘贴第二张图像到右侧(留有空隙)
    toImage.paste(img2, (img1.width + 35, 0))
    # 保存对比图像
    toImage.save(arr[2], quality=100)