昇腾 AI 行业案例（五）：基于 DANet 和 Deeplabv3 模型的遥感图像分割

# 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)    # 将图片从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 = "./test_set/test_15.jpg"process_data = img_process(source)print(process_data.shape)

(3, 480, 640)(3, 480, 640)(3, 480, 640)(3, 480, 640)(3, 480, 640)(3, 480, 640)(3, 480, 640)

# atc --framework=5 --model=../models/DANet.onnx  --output=../models/DANet --soc_version=芯片型号 --insert_op_conf=./configure.cfg --log=error# atc --framework=5 --model=../models/Deeplabv3.onnx  --output=../models/Deeplabv3 --soc_version=芯片型号 --insert_op_conf=./configure.cfg --log=error# config 文件下载地址: https://gitee.com/ascend/mindxsdk-referenceapps/blob/master/contrib/RemoteSensingSegmentation/config/configure.cfg

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()

# 加载模型da_net_om_model_path = "./models/DANet.om"da_net_om_model = OmModel(da_net_om_model_path)deeplab_om_model_path = "./models/Deeplabv3.om"deeplab_om_model = OmModel(deeplab_om_model_path)# 推理output_res_DANet = da_net_om_model.forward([process_data])output_res_Deeplab = deeplab_om_model.forward([process_data])tensor_res_DANet = output_res_DANet[0].reshape(1, 8, 256, 256)  # 模型的输出shape是 (1, 8, 256, 256)tensor_res_Deeplab = output_res_Deeplab[0].reshape(1, 8, 256, 256)  # 模型的输出shape是 (1, 8, 256, 256)print(output_res_DANet.shape)print(output_res_Deeplab.shape)

def softmax(x):    """        Compute softmax values for each sets of scores in x.        Args:            matrix x        Returns:            softmax values for x    """    mx = np.max(x, axis=1, keepdims=True)    numerator = np.exp(x - mx)    denominator = np.sum(numerator, axis=1, keepdims=True)    return numerator / denominator
pred_DANet = softmax(tensor_res_DANet)pred_Deeplab = softmax(tensor_res_Deeplab)pred_final = pred_Deeplab + pred_DANet

def semantic_to_mask(mask, labels):    """        Turn Semantic diagram to label diagram        Args:            mask: semantic diagram            labels: 8 kinds of semantic tags        Returns:            the maximum probability prediction result after label replacement    """    x = np.argmax(mask, axis=1)    label_codes = np.array(labels)    x = np.uint16(label_codes[x.astype(np.uint8)])    return x

def decode_seg_map(label_mask, labels, label_colours):    """        The result is mapped to a picture        Args:            label_mask: the maximum probability prediction result after label replacement            labels: 8 kinds of semantic tags            label_colours: colors corresponding to 8 semantic tags        Returns:            the prediction result after color replacement    """    r = label_mask.copy()    g = label_mask.copy()    b = label_mask.copy()    for i, label in enumerate(labels):        r[label_mask == label] = label_colours[i, 0]        g[label_mask == label] = label_colours[i, 1]        b[label_mask == label] = label_colours[i, 2]    rgb = np.zeros((label_mask.shape[0], label_mask.shape[1], 3))    rgb[:, :, 0] = r / 255.0    rgb[:, :, 1] = g / 255.0    rgb[:, :, 2] = b / 255.0    return rgb

def final_result_create(final_res_pic, label_colours, save_path):    """        Result image post-processing        Args:            final_res_pic: the prediction result after color replacement            label_colours: colors corresponding to 8 semantic tags        Returns:            null        Output:            a result picture in directory result/temp_result/result.jpg    """
    tmp_rgb = np.array(label_colours) / 255.0    ic_map = mpl.colors.ListedColormap(tmp_rgb, name='my_color')    # Normalize the data to an interval of [0, 8]    norm = mpl.colors.Normalize(vmin=0, vmax=8)    # the size of result picture setting    plt.figure(figsize=(5, 4))    h = plt.imshow(final_res_pic, cmap=ic_map, norm=norm)    c_bar = plt.colorbar(mappable=h)    # Return evenly spaced numbers over a specified interval    v = np.linspace(0, 7, 8)    # Set the location and properties of the ticks    c_bar.set_ticks((v + 0.5))    # Set the label of the ticks    c_bar.set_ticklabels(['water', 'road', 'buildings', 'tree', 'grass', 'cultivated_land', 'soil', 'others'])    plt.savefig(save_path)    plt.close()

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])    img12 = Image.open(arr[1])    # create a new image, set the width and height    toImage = Image.new('RGB', (img1.width + img12.width + 35, img12.height), 'white')    # paste image1 to the new image, and set the position    toImage.paste(img1, (35, 80))    # paste image2 to the new image, and set the position    toImage.paste(img12, (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 matplotlib import pyplot as plt
img_dir = "./test_set/"save_path = "./result/"imgs_path = glob.glob(str(Path(img_dir).resolve()) + '/*.*')# 8 kinds of semantic tagslabels = [100, 200, 300, 400, 500, 600, 700, 800]# Colors corresponding to 8 semantic tagslabel_colours = np.array(    [[255, 127, 39], [255, 202, 24], [236, 28, 36], [14, 209, 69], [184, 61, 186],        [0, 168, 243], [88, 88, 88], [253, 236, 166]])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_DANet = da_net_om_model.forward([image])    output_res_Deeplab = deeplab_om_model.forward([image])    if output_res_DANet == [] or output_res_Deeplab == []:        continue    # reshape the matrix to (1, 8, 256, 256)    tensor_res_DANet = output_res_DANet[0].reshape(1, 8, 256, 256)    pred_DANet = softmax(tensor_res_DANet)
    # Get the result of the Deeplab_v3 model    # reshape the matrix to (1, 8, 256, 256)    tensor_res_Deeplab = output_res_Deeplab[0].reshape(1, 8, 256, 256)    pred_Deeplab = softmax(tensor_res_Deeplab)
    pred_final = pred_Deeplab + pred_DANet
    # Semantic diagram to label diagram    pred_final = semantic_to_mask(pred_final, labels).squeeze().astype(np.uint16)    # The result is mapped to a picture    rgb_res_pic = decode_seg_map(pred_final, labels, label_colours)    # Result image post-processing    final_result_create(rgb_res_pic, label_colours, segment_save_path)    # Enable comparison graph output if true    enable_contrast_output([img_path, segment_save_path, compare_save_path])    print('success!')