昇腾 AI 行业案例（六）：基于 PraNet 的医疗影像分割

# DO NOT CHANGE THIS CELL# activate this cell by selecting it with the mouse or arrow keys then use the keyboard shortcut [Shift+Enter] to executeprint('This is just a simple print statement')

This is just a simple print statement

import numpy as npimport cv2

def img_process(img_path):    image = cv2.imread(img_path)    # resize成模型的输入shape    image = cv2.resize(image, (352, 352), interpolation=cv2.INTER_AREA)    # 将图片从BGR颜色空间转换为YUV颜色空间    yuv_image = cv2.cvtColor(image, cv2.COLOR_BGR2YUV)    # 获取图片的尺寸    height, width = yuv_image.shape[:2]    # 分离Y, U, V分量    y, u, v = cv2.split(yuv_image)    # 将U和V分量下采样到4:2:0    u = cv2.resize(u, (width // 2, height // 2), interpolation=cv2.INTER_LINEAR)    v = cv2.resize(v, (width // 2, height // 2), interpolation=cv2.INTER_LINEAR)    # 交错UV分量，形成NV12格式    uv = np.zeros((height // 2, width), dtype=yuv_image.dtype)    uv[:, 0::2] = u    uv[:, 1::2] = v    # 将Y分量和交错的UV分量合并    image = np.concatenate((y, uv), axis=0)    return image

source = "./kvasir_test_pics/cju2zrojo9kcd0878ld2epejq.jpg"process_data = img_process(source)print(process_data.shape)

(528, 352)

# atc --model=./PraNet/PraNet-19.onnx --output=./PraNet-19_bs1 --framework=5 --input_shape="actual_input_1:1,3,352,352" --soc_version=Ascendxxx --input_format=NCHW --output_type=FP32 --insert_op_conf=./pranet.aippconfig# config 文件下载地址: https://gitee.com/ascend/mindxsdk-referenceapps/blob/master/contrib/PraNetSegmentation/model/pranet.aippconfig

import acl
ACL_MEM_MALLOC_HUGE_FIRST = 0ACL_MEMCPY_HOST_TO_DEVICE = 1ACL_MEMCPY_DEVICE_TO_HOST = 2
class OmModel:    def __init__(self, model_path):        # 初始化函数        self.device_id = 5
        # step1: 初始化        ret = acl.init()        # 指定运算的Device        ret = acl.rt.set_device(self.device_id)
        # step2: 加载模型，本示例为pfld模型        # 加载离线模型文件，返回标识模型的ID        self.model_id, ret = acl.mdl.load_from_file(model_path)        # 创建空白模型描述信息，获取模型描述信息的指针地址        self.model_desc = acl.mdl.create_desc()        # 通过模型的ID，将模型的描述信息填充到model_desc        ret = acl.mdl.get_desc(self.model_desc, self.model_id)
        # step3：创建输入输出数据集        # 创建输入数据集        self.input_dataset, self.input_data = self.prepare_dataset('input')        # 创建输出数据集        self.output_dataset, self.output_data = self.prepare_dataset('output')
    def prepare_dataset(self, io_type):        # 准备数据集        if io_type == "input":            # 获得模型输入的个数            io_num = acl.mdl.get_num_inputs(self.model_desc)            acl_mdl_get_size_by_index = acl.mdl.get_input_size_by_index        else:            # 获得模型输出的个数            io_num = acl.mdl.get_num_outputs(self.model_desc)            acl_mdl_get_size_by_index = acl.mdl.get_output_size_by_index        # 创建aclmdlDataset类型的数据，描述模型推理的输入。        dataset = acl.mdl.create_dataset()        datas = []        for i in range(io_num):            # 获取所需的buffer内存大小            buffer_size = acl_mdl_get_size_by_index(self.model_desc, i)            # 申请buffer内存            buffer, ret = acl.rt.malloc(buffer_size, ACL_MEM_MALLOC_HUGE_FIRST)            # 从内存创建buffer数据            data_buffer = acl.create_data_buffer(buffer, buffer_size)            # 将buffer数据添加到数据集            _, ret = acl.mdl.add_dataset_buffer(dataset, data_buffer)            datas.append({"buffer": buffer, "data": data_buffer, "size": buffer_size})        return dataset, datas
    def forward(self, inputs):        # 执行推理任务        # 遍历所有输入，拷贝到对应的buffer内存中        input_num = len(inputs)        for i in range(input_num):            bytes_data = inputs[i].tobytes()            bytes_ptr = acl.util.bytes_to_ptr(bytes_data)            # 将图片数据从Host传输到Device。            ret = acl.rt.memcpy(self.input_data[i]["buffer"],  # 目标地址 device                                self.input_data[i]["size"],  # 目标地址大小                                bytes_ptr,  # 源地址 host                                len(bytes_data),  # 源地址大小                                ACL_MEMCPY_HOST_TO_DEVICE)  # 模式:从host到device        # 执行模型推理。        ret = acl.mdl.execute(self.model_id, self.input_dataset, self.output_dataset)        # 处理模型推理的输出数据，输出top5置信度的类别编号。        inference_result = []        for i, item in enumerate(self.output_data):            buffer_host, ret = acl.rt.malloc_host(self.output_data[i]["size"])            # 将推理输出数据从Device传输到Host。            ret = acl.rt.memcpy(buffer_host,  # 目标地址 host                                self.output_data[i]["size"],  # 目标地址大小                                self.output_data[i]["buffer"],  # 源地址 device                                self.output_data[i]["size"],  # 源地址大小                                ACL_MEMCPY_DEVICE_TO_HOST)  # 模式：从device到host            # 从内存地址获取bytes对象            bytes_out = acl.util.ptr_to_bytes(buffer_host, self.output_data[i]["size"])            # 按照float32格式将数据转为numpy数组            data = np.frombuffer(bytes_out, dtype=np.float32)            inference_result.append(data)        return inference_result
    def __del__(self):        # 析构函数 按照初始化资源的相反顺序释放资源。        # 销毁输入输出数据集        for dataset in [self.input_data, self.output_data]:            while dataset:                item = dataset.pop()                ret = acl.destroy_data_buffer(item["data"])  # 销毁buffer数据                ret = acl.rt.free(item["buffer"])  # 释放buffer内存        ret = acl.mdl.destroy_dataset(self.input_dataset)  # 销毁输入数据集        ret = acl.mdl.destroy_dataset(self.output_dataset)  # 销毁输出数据集        # 销毁模型描述        ret = acl.mdl.destroy_desc(self.model_desc)        # 卸载模型        ret = acl.mdl.unload(self.model_id)        # 释放device        ret = acl.rt.reset_device(self.device_id)        # acl去初始化        ret = acl.finalize()

# 加载模型pranet_om_model_path = "./PraNet-19_bs1.om"pranet_om_model = OmModel(pranet_om_model_path)# 推理output_res = pranet_om_model.forward([process_data])tensor_res = output_res[0].reshape(1, 1, 352, 352)  # 模型的输出shape是 (1, 1, 352, 352)print(tensor_res.shape)

(1, 1, 352, 352)

THRESHOLD_VALUE = 0.5  # 二值化置信度阈值，根据算法效果调试得到
def decode_seg_map(label_mask, save_path):    # 获取2维的分类结果    segment_result = label_mask[0][0]    # 新建一个0矩阵，表示二值分割像素矩阵    seg_map = np.zeros((segment_result.shape[0], segment_result.shape[1]), dtype=np.float32)    # 二值化    seg_map[segment_result > THRESHOLD_VALUE] = 255    # 保存二值化结果    cv2.imwrite(save_path, seg_map)    return

def enable_contrast_output(arr):    """        Enable comparison graph output        Args:            arr: arr[0] is the img one, arr[1] is the img two, arr[2] is the output directory of the comparison graph result        Returns:            null        Output:            a comparison graph result in arr[2]    """    img1 = Image.open(arr[0])    img2 = Image.open(arr[1])    img2 = img2.resize(img1.size, Image.NEAREST)    # create a new image, set the width and height    toImage = Image.new('RGB', (img1.width + img2.width + 35, img2.height), 'white')    # paste image1 to the new image, and set the position    toImage.paste(img1, (0, 0))    # paste image2 to the new image, and set the position    toImage.paste(img2, (img1.width + 35, 0))    # save the result image, and the quality is 100    toImage.save(arr[2], quality=100)

import osimport globfrom pathlib import Pathfrom PIL import Image
img_dir = "./kvasir_test_pics/"save_path = "./result/"imgs_path = glob.glob(str(Path(img_dir).resolve()) + '/*.*')if os.path.exists(img_dir) != 1:    print("The test image " + str(img_dir) + " does not exist.")    exit()for img_path in imgs_path:    img_name = img_path.split("/")[-1]    segment_save_path = save_path + "seg_" + img_name    compare_save_path = save_path + "compare_" + img_name
    image = img_process(img_path)    # Get the result of the DANet model    output_res_PraNet = pranet_om_model.forward([image])    if output_res_PraNet == []:        continue    # reshape the matrix to (1, 1, 352, 352)    tensor_res_PraNet = output_res_PraNet[0].reshape(1, 1, 352, 352)    # The result is mapped to a picture    decode_seg_map(tensor_res_PraNet, segment_save_path)    # Enable comparison graph output if true    enable_contrast_output([img_path, segment_save_path, compare_save_path])    print('success!')

success!success!success!