干货 | 基于PIE-Engine AI的UperNet-SwinTransformer模型上传实践——以光伏目标提取

2022-03-18 10:50 作者:航天宏图 0人读过 | 我要投稿

PIE-Engine AI 介绍

PIE-Engine AI 是航天宏图自主研发的一站式遥感图像智能解译服务云平台。该平台结合空间信息技术和人工智能技术，将深度学习理论与遥感解译实践进行融合，全方位提升遥感数据智能化处理和解译分析能力。可为用户提供实时、精准、高效的遥感图像在线智能解译服务，包括遥感地物语义分割、目标检测、变化监测等应用方向。

本文以 0.3m 分辨率影像光伏目标提取为例，介绍了如何将本地训练的 UperNet-SwinTransformer 模型集成至 PIE-Engine AI，并实现智能解译的一站式深度学习服务。

UperNet-SwinTransformer 模型简介
//

模型概述

SwinTransformer是微软亚洲研究院提出的新型视觉Transformer，它可以作为计算机视觉的通用骨干网络。视觉领域与自然语言领域之间存在巨大差异，这带来了使 Transformer 从自然语言领域适应视觉领域的挑战。

为了解决这些差异，SwinTransformer 提出了一个分层的 Transformer，其表示是通过移动窗口来计算的。通过将自注意力计算限制为不重叠的局部窗口，同时允许跨窗口连接，移位的窗口方案带来了更高的效率。这种分层体系结构具有在各种尺度上建模的灵活性，并且相对于图像大小具有线性计算复杂性。SwinTransformer 的这些品质使其可与多种视觉任务兼容。

本文以 SwinTransformer 作为 UperNet 的骨干网络，构建 UperNet-SwinTransformer 网络进行光伏目标提取。采用 0.3m 分辨率的光伏语义分割数据集进行训练，达到了 mIoU=96.07 的精度。

UperNet-SwinTransformer 模型上传代码实现
01

模型上传压缩包结构

模型上传压缩包分为两部分：模型描述文件和模型代码文件夹。

1、模型描述文件

模型描述文件 modelMeta.json 包含详细的模型信息描述，包括模型名称、权重文件名称、模型训练数据集的相关信息、网络结构相关信息、记录图片名称、评价指标等。

以下代码给出参数说明，各参数详细配置见示例文件。

{ "注释说明1":"以下为必填项", "name": "光伏提取模型", "weight": "pv.pth", "dataset": { "label":[{"name":"background","title":"背景","color":"rgb(0,0,0)","value":"0"}, {"name":"pv","title":"光伏","color":"rgb(255,0,0)","value":"1"}], "imageSize": "1024,1024", "bands": "3", "dataType": "Byte" }, "注释说明2":"以下为选填项", "labels": "0.3m光伏", "description": "该模型是基于pytorch 1.7深度学习框架,使用0.3m分辨率光伏数据集进行训练", "network": { "thumbnail": "model.jpg" }, "icon":"存放模型的原始及预测图片，图片名称需统一，详情查看REAMDE.md", "evaluate":{ "accuary": "98.14%", "description":{"miou": "miou"}, "best":{"miou": 0.9607} }}

←左右滑动查看更多→

2、模型代码文件夹

模型代码文件夹 model 主要包含 icon、weight 等文件夹和Predict.py、PredictModel.py 等脚本文件，相关命名及编写规范在示例文件进行了说明。

▶ icon 文件夹存放预测模型图片，命名规则参考 modelMeta.json，若无相关图片展示，则删除该文件夹。

▶ weight 文件夹存放权重文件，名称需与 modelMeta.json 中完全一致。

▶ model.jpg 上传网络结构图。

▶ Predict.py 智能解译脚本，主要包括输入输出两部分内容：

输入：

image_path(图片路径)、weight_path(权重文件路径)、gpu_num(使用gpu个数)；

输出:

(1)语义分割、变化检测：tif影像及对应矢量；

(2)目标识别：geojson目标框文件及对应矢量。

▶ PredictModel.py 模型发布脚本，主要包括输入输出两部分内容：

输入：

image_path(图片路径)、weight_path(权重文件路径)、gpu_num(使用gpu个数)；

输出:

(1)语义分割、变化检测：单波段图片二进制流；

(2)目标识别：geojson目标框二进制流。

影像智能解译 pipeline 及核心代码实现

用户需在 Predict.py 文件中编写影像预测代码，从而使上传的模型能完成影像智能解译任务。影像预测代码需实现以下功能：

以下代码给出影像预测主函数 prediction_image 的示例。

def prediction_image(**platform_arguments):

load_model = platform_arguments['load_model']

deal_data_img = platform_arguments['deal_data_img']

network_type = platform_arguments['network_type']

# 获取读取s3的gdal

s3gdal = inint_gdal()

# 加载模型

model = get_model(load_model)

# 获取预测图片列表

inputList = platform_arguments['image_path'].split(',')

# 记录日志

log_new = {"zipOutFile": '/' + '/'.join(outputSaveDir.split('/')[2:]) + "result.zip", "total_output_size": "", "task_list": []} total_output_size = 0

# 遍历图片，进行每张图片预测

for index, imagePath in enumerate(inputList):

process = float((index + 1) / len(inputList))

log_new["process"] = process

taskIdAndIndex = platform_arguments['estimate_id'] + ":" + str(index + 1) tmpfilename = os.path.basename(imagePath)

filename, extension = os.path.splitext(tmpfilename)

# 重命名，例如：1_xxx_pred

filename = str(index + 1) + "_" + filename + "_pred"

# 输出

outputFile = outputSaveDir + filename + ".tif"

pred_utils = pred_prepare(model,s3gdal,imagePath)

if network_type == 1 or network_type == 3:

# 滑窗预测图片，生成三波段展示tif及临时单波段tif pred_utils.pred_image_seg_change(outputFile, platform_arguments['load_size'],deal_data_img,platform_arguments['value_color'],platform_arguments['background_value'])

# 生成金字塔及转换为shp，并保存入数据库

staticMap = shp_tif_deal(outputFile, pred_utils, taskIdAndIndex, **platform_arguments) elif network_type == 2:

# 预测图片，生成 shp并生成geojson pred_utils.pred_image_detection(outputFile[:-4] + '.shp',platform_arguments['load_size'],platform_arguments['class_name_list'], deal_data_img,platform_arguments['value_title_map'],platform_arguments['name_title_map'],nms=0.3)

# 存入数据库

staticMap = shp_tif_deal(outputFile, pred_utils, taskIdAndIndex, **platform_arguments) else:

print('输入类型有误！') exit(1)

# 日志信息

task_list_inner = { "task_result_id": taskIdAndIndex, "output_shp_static": staticMap, "output_size": "", "input": imagePath, "output": [], }

# 记录日志信息create_log(task_list_inner,outputFile,total_output_size,log_new) print("[Process: {process}]".format(process=str(index + 1) + '/' + str(len(inputList))))

# 矢量进行压缩create_zip()

create_zip()

03

模型上传到 PIE-Engine AI 平台