使用LabVIEW實現基於pytorch的DeepLabv3影像語義分割

virobotics發表於2022-11-04

前言

今天我們一起來看一下如何使用LabVIEW實現語義分割。

一、什麼是語義分割

影像語義分割（semantic segmentation），從字面意思上理解就是讓計算機根據影像的語義來進行分割，例如讓計算機在輸入下面左圖的情況下，能夠輸出右圖。語義在語音識別中指的是語音的意思，在影像領域，語義指的是影像的內容，對圖片意思的理解，比如下圖的語義就是一個人牽著四隻羊；分割的意思是從畫素的角度分割出圖片中的不同物件，對原圖中的每個畫素都進行標註，比如下圖中淺黃色代表人，藍綠色代表羊。語義分割任務就是將圖片中的不同類別，用不同的顏色標記出來，每一個類別使用一種顏色。常用於醫學影像，衛星影像，無人車駕駛，機器人等領域。在這裡插入圖片描述

如何做到將畫素點上色呢？

語義分割的輸出和影像分類網路類似，影像分類類別數是一個一維的one hot 矩陣。例如：三分類的[0,1,0]。語義分割任務最後的輸出特徵圖是一個三維結構，大小與原圖類似，其中通道數是類別數，每個通道所標記的畫素點，是該類別在影像中的位置，最後透過argmax 取每個通道有用畫素合成一張影像，用不同顏色表示其類別位置。語義分割任務其實也是分類任務中的一種，他不過是對每一個畫素點進行細分，找到每一個畫素點所述的類別。這就是語義分割任務啦~

在這裡插入圖片描述

二、什麼是deeplabv3

DeepLabv3是一種語義分割架構，它在DeepLabv2的基礎上進行了一些修改。為了處理在多個尺度上分割物件的問題，設計了在級聯或並行中採用多孔卷積的模組，透過採用多個多孔速率來捕獲多尺度上下文。此外，來自 DeepLabv2 的 Atrous Spatial Pyramid Pooling模組增加了編碼全域性上下文的影像級特徵，並進一步提高了效能。在這裡插入圖片描述

三、LabVIEW呼叫DeepLabv3實現影像語義分割

1、模型獲取及轉換

安裝pytorch和torchvision
獲取torchvision中的模型：deeplabv3_resnet101(我們獲取預訓練好的模型)：

 original_model = models.segmentation.deeplabv3_resnet101(pretrained=True)

轉onnx

 1 def get_pytorch_onnx_model(original_model):
 2     # define the directory for further converted model save
 3     onnx_model_path = dirname
 4     # define the name of further converted model
 5     onnx_model_name = "deeplabv3_resnet101.onnx"
 6 
 7     # create directory for further converted model
 8     os.makedirs(onnx_model_path, exist_ok=True)
 9 
10     # get full path to the converted model
11     full_model_path = os.path.join(onnx_model_path, onnx_model_name)
12 
13     # generate model input
14     generated_input = Variable(
15         torch.randn(1, 3, 448, 448)
16     )
17 
18     # model export into ONNX format
19     torch.onnx.export(
20         original_model,
21         generated_input,
22         full_model_path,
23         verbose=True,
24         input_names=["input"],
25         output_names=["output",'aux'],
26         opset_version=11
27     )
28 
29     return full_model_path

View Code

完整獲取及模型轉換python程式碼如下：

 1 import os
 2 import torch
 3 import torch.onnx
 4 from torch.autograd import Variable
 5 from torchvision import models
 6 import re
 7 
 8 dirname, filename = os.path.split(os.path.abspath(__file__))
 9 print(dirname)
10 
11 def get_pytorch_onnx_model(original_model):
12     # define the directory for further converted model save
13     onnx_model_path = dirname
14     # define the name of further converted model
15     onnx_model_name = "deeplabv3_resnet101.onnx"
16 
17     # create directory for further converted model
18     os.makedirs(onnx_model_path, exist_ok=True)
19 
20     # get full path to the converted model
21     full_model_path = os.path.join(onnx_model_path, onnx_model_name)
22 
23     # generate model input
24     generated_input = Variable(
25         torch.randn(1, 3, 448, 448)
26     )
27 
28     # model export into ONNX format
29     torch.onnx.export(
30         original_model,
31         generated_input,
32         full_model_path,
33         verbose=True,
34         input_names=["input"],
35         output_names=["output",'aux'],
36         opset_version=11
37     )
38 
39     return full_model_path
40 
41 
42 def main():
43     # initialize PyTorch ResNet-101 model
44     original_model = models.segmentation.deeplabv3_resnet101(pretrained=True)
45 
46     # get the path to the converted into ONNX PyTorch model
47     full_model_path = get_pytorch_onnx_model(original_model)
48     print("PyTorch ResNet-101 model was successfully converted: ", full_model_path)
49 
50 
51 if __name__ == "__main__":
52     main()

View Code

我們會發現，基於pytorch的DeepLabv3模型獲取和之前的mask rcnn模型大同小異。

2、關於deeplabv3_resnet101

我們使用的模型是：deeplabv3_resnet101，該模型返回兩個張量，與輸入張量相同，但有21個classes。輸出[“out”]包含語義掩碼，而輸出[“aux”]包含每畫素的輔助損失值。在推理模式中，輸出[‘aux]沒有用處。因此，輸出“out”形狀為（N、21、H、W）。我們在轉模型的時候設定H，W為448，N一般為1；

我們的模型是基於VOC2012資料集 VOC2012資料集分為20類，包括背景為21類，分別如下：