基于华为atlas的unet分割模型探索

Unet���型使用官方基于kaggle Carvana Image Masking Challenge数据集训练的模型。

模型输入为572*572*3，输出为572*572*2。分割目标分别为，0：背景，1：汽车。

Pytorch的pth模型转化onnx模型：

import torchfrom unet import UNetmodel = UNet(n_channels=3, n_classes=2, bilinear=False)model = model.to(memory_format=torch.channels_last)state_dict = torch.load("unet_carvana_scale1.0_epoch2.pth", map_location="cpu")#del state_dict['mask_values']model.load_state_dict(state_dict)dummy_input = torch.randn(1, 3, 572, 572)torch.onnx.export(model, dummy_input, "unet.onnx", verbose=True)

模型输入输出节点分析：

使用工具Netron查看模型结构，确定模型输入节点名称为input.1，输出节点名称为/outc/conv/Conv

onnx模型转化atlas模型：

atc --model=./unet.onnx --framework=5 --output=unet --soc_version=Ascend310P3  --input_shape="input.1:1,3,572,572" --output_type="/outc/conv/Conv:0:FP32" --out_nodes="/outc/conv/Conv:0"

推理代码实现：

import base64import jsonimport osimport timeimport numpy as npimport cv2import MxpiDataType_pb2 as mxpi_datafrom StreamManagerApi import InProtobufVectorfrom StreamManagerApi import MxProtobufInfrom StreamManagerApi import StreamManagerApidef check_dir(dir):    if not os.path.exists(dir):        os.makedirs(dir, exist_ok=True)class SDKInferWrapper:    def __init__(self): # 完成初始化        self._stream_name = None        self._stream_mgr_api = StreamManagerApi()        if self._stream_mgr_api.InitManager() != 0:            raise RuntimeError("Failed to init stream manager.")        pipeline_name = './nested_unet.pipeline'        self.load_pipeline(pipeline_name)        self.width = 572        self.height = 572    def load_pipeline(self, pipeline_path):        with open(pipeline_path, 'r') as f:            pipeline = json.load(f)        self._stream_name = list(pipeline.keys())[0].encode() # 'unet_pytorch'        if self._stream_mgr_api.CreateMultipleStreams(                json.dumps(pipeline).encode()) != 0:            raise RuntimeError("Failed to create stream.")    def do_infer(self, img_bgr):        # preprocess        image = cv2.resize(img_bgr, (self.width, self.height))        image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)        image = image.astype('float32') / 255.0        image = image.transpose(2, 0, 1)        tensor_pkg_list = mxpi_data.MxpiTensorPackageList()        tensor_pkg = tensor_pkg_list.tensorPackageVec.add()        tensor_vec = tensor_pkg.tensorVec.add()        tensor_vec.deviceId = 0        tensor_vec.memType = 0        for dim in [1, *image.shape]:            tensor_vec.tensorShape.append(dim) # tensorshape属性为[1,3,572,572]        input_data = image.tobytes()        tensor_vec.dataStr = input_data        tensor_vec.tensorDataSize = len(input_data)        protobuf_vec = InProtobufVector()        protobuf = MxProtobufIn()        protobuf.key = b'appsrc0'        protobuf.type = b'MxTools.MxpiTensorPackageList'        protobuf.protobuf = tensor_pkg_list.SerializeToString()        protobuf_vec.push_back(protobuf)        unique_id = self._stream_mgr_api.SendProtobuf(            self._stream_name, 0, protobuf_vec)        if unique_id = 0] = 1.0 / (1 + np.exp(-x[x >= 0]))    y[x  
运行代码：
 
set -e. /usr/local/Ascend/ascend-toolkit/set_env.sh# Simple log helper functionsinfo() { echo -e "3[1;34m[INFO ][MxStream] 3[1;37m" ; }warn() { echo >&2 -e "3[1;31m[WARN ][MxStream] 3[1;37m" ; }#export MX_SDK_HOME=/home/work/mxVisionexport LD_LIBRARY_PATH=${MX_SDK_HOME}/lib:${MX_SDK_HOME}/opensource/lib:${MX_SDK_HOME}/opensource/lib64:/usr/local/Ascend/ascend-toolkit/latest/acllib/lib64:${LD_LIBRARY_PATH}export GST_PLUGIN_SCANNER=${MX_SDK_HOME}/opensource/libexec/gstreamer-1.0/gst-plugin-scannerexport GST_PLUGIN_PATH=${MX_SDK_HOME}/opensource/lib/gstreamer-1.0:${MX_SDK_HOME}/lib/plugins#to set PYTHONPATH, import the StreamManagerApi.pyexport PYTHONPATH=$PYTHONPATH:${MX_SDK_HOME}/pythonpython3 unet.pyexit 0
  
运行效果：
   
个人思考：
 
华为atlas的参考案例细节不到位，步骤缺失较多，摸索困难，代码写法较差，信创化道路任重而道远。
  
参考资料：
 
GitHub - milesial/Pytorch-UNet: PyTorch implementation of the U-Net for image semantic segmentation with high quality images
 
https://gitee.com/ascend/samples/tree/master/python/level2_simple_inference/3_segmentation/unet++