我自己的原文哦~???https://blog.51cto.com/whaosoft/11608027
一、使用Pytorch進行簡單的自定義圖像分類 ~ONNX 推理
???圖像分類是計算機視覺中的一項基本任務,涉及訓練模型將圖像分類為預定義類別。本文中,我們將探討如何使用 PyTorch 構建一個簡單的自定義對象分類模型,然后使用 ONNX 格式將其部署用于推理。
數據集準備
????在開始創建模型之前,準備一個標記數據集至關重要。收集要分類的不同對象類別的圖像,并根據其類別將它們組織到單獨的文件夾中。確保每個類別都有足夠數量的圖像,以避免過度擬合。
????準備如下樹所示的數據集
- data- Fruits (dataset name)- train- class 1- class 2- ...- class n- val- class 1- class 2- ...- class n- test- class 1- class 2- ...- class n
????我從 kaggle 獲取了水果數據集,作為示例鏈接在此處:
https://www.kaggle.com/datasets/shreyapmaher/fruits-dataset-images
????更改 main.py 中“train_dir”和“val_dir”中的路徑名
????如果需要,初始化數據加載器并添加增強功能。
構建模型
????首先導入必要的庫,包括 PyTorch。定義自定義對象分類模型的架構,通常使用卷積神經網絡 (CNN)。設計網絡層,包括卷積層和池化層
import torch.nn as nnclass CustomConvNet(nn.Module):def __init__(self, num_classes):super(CustomConvNet, self).__init__()self.num_classes = num_classesself.layer1 = self.conv_module(3, 16)self.layer2 = self.conv_module(16, 32)self.layer3 = self.conv_module(32, 64)self.layer4 = self.conv_module(64, 128)self.layer5 = self.conv_module(128, 256)self.gap = self.global_avg_pool(256, self.num_classes)def forward(self, x):out = self.layer1(x)out = self.layer2(out)out = self.layer3(out)out = self.layer4(out)out = self.layer5(out)out = self.gap(out)out = out.view(-1, self.num_classes)return outdef conv_module(self, in_num, out_num):return nn.Sequential(nn.Conv2d(in_num, out_num, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(out_num),nn.LeakyReLU(),nn.MaxPool2d(kernel_size=2, stride=2))def global_avg_pool(self, in_num, out_num):return nn.Sequential(nn.Conv2d(in_num, out_num, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(out_num),nn.LeakyReLU(),nn.AdaptiveAvgPool2d((1, 1)))
訓練模型并進行評估
????使用隨機梯度下降 (SGD) 或 Adam 優化器等技術實現訓練循環,包括前向和后向傳播、損失計算和梯度優化。通過跟蹤損失和準確度等指標來監控訓練過程。我們利用數據增強和正則化等技術來提高模型的泛化能力。
python main.py
for epoch in range(num_epochs):print("Epoch No -", epoch)model.train()running_loss = 0.0running_corrects = 0for inputs, labels in dataLoaders["train"]:# Feeding input and labels to deviceinputs = inputs.to(device, non_blocking=True)labels = labels.to(device, non_blocking=True)optimizer.zero_grad()with torch.set_grad_enabled(True):outputs = model(inputs)_, preds = torch.max(outputs,1)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()* inputs.size(0)#calculate accuracyrunning_corrects += torch.sum(preds == labels.data)#scheduler stepexp_lr_scheduler.step()# Calculate average loss and acc for a epochepoch_loss = running_loss/len(train_data)epoch_acc = running_corrects.double()/len(train_data)print('Loss:{} , Acc{}'.format(epoch_loss, epoch_acc))# Saving model every five epochif (epoch%5 == 0):save_model(model,epoch)
????運行預測
????確保你是否已經更改了權重文件名、訓練文件夾、predict.py 中的輸入圖像
python predict.py
導出為 ONNX 格式
????一旦您的模型經過訓練并在驗證集上表現良好,就可以將其導出為 ONNX 格式進行部署和推理。ONNX(開放神經網絡交換)是一種開放標準格式,允許不同深度學習框架之間的互操作性。PyTorch 提供了將模型導出為 ONNX 的工具。
????確保你是否更改了export.py中的權重文件名
python export.py
# Now we will save this model.
import torch.onnx
torch.onnx.export(model,img,"./savedModels/custommodel.onnx",export_params=True,opset_version=10,verbose=True, # Print verbose outputinput_names=['input'], # Names for input tensoroutput_names=['output'])
使用 ONNX 進行推理?
????加載已保存的 ONNX 模型并對新的未見過的圖像進行推理。使用 ONNX 運行時庫加載模型、提供輸入數據并獲得預測。測量推理時間并將其與 PyTorch 模型的推理時間進行比較,以評估通過 ONNX 優化實現的任何性能改進。
????確保你是否已經更改了權重文件名、訓練文件夾、predict.py 中的輸入圖像
python onnx_inference.py
# Load the ONNX model
onnx_model = onnx.load("./savedModels/custommodel.onnx")# Create an ONNX runtime session
ort_session = onnxruntime.InferenceSession("./savedModels/custommodel.onnx")inputs = {"input": trans_img.numpy()}
outputs = ort_session.run(None, inputs)
開發板商城 天皓智聯
????在本教程中,我們探索了使用 PyTorch 構建簡單自定義對象分類模型的過程。我們學習了如何訓練模型、評估其性能并將其導出為 ONNX 格式以進行推理。通過利用 ONNX 運行時,我們演示了如何高效地對新圖像進行推理。有了這些知識,您現在可以將自定義對象分類應用于各種實際應用程序并無縫部署模型。
二、讀懂 ONNX、TensorRT、OpenVINO部署框架
本文詳細介紹了深度學習模型部署過程中常用的幾個框架:ONNX、TensorRT 和 OpenVINO,包括它們的功能、優勢以及如何將 PyTorch 模型轉換為這些框架支持的格式,旨在提高模型在不同硬件平臺上的推理效率和性能。文章還討論了模型轉換過程中可能遇到的問題和相應的解決方案。
這一期主要會分幾個點展開:為什么我們做部署的時候要在 torch 上更進一步使用 ONNX,TensorRT,OpenVINO 等部署框架,在做 cv 模型部署的時候。我們怎么部署。在做 LLM 部署的時候,我們又會怎么做呢?
動靜轉換:Torch上更進一步
Torch
最核心的就是 torch 使用了動態圖組網。使用動態組網的好處是。可以使用更偏向 python 語法的格式對模型進行定義。下面就給大家一個常見的網絡:
import torch
import torch.nn as nn
import torch.nn.functional as Fclass SimpleNet(nn.Module):def __init__(self, input_size, hidden_size, num_classes):super(SimpleNet, self).__init__()self.fc1 = nn.Linear(input_size, hidden_size)self.fc2 = nn.Linear(hidden_size, hidden_size)self.fc3 = nn.Linear(hidden_size, num_classes)def forward(self, x):x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))x = self.fc3(x)return x# 示例使用
input_size = 784 # 例如,對于MNIST數據集
hidden_size = 128
num_classes = 10model = SimpleNet(input_size, hidden_size, num_classes)
print(model)
大家很容易開心的寫出這樣代碼,但問題是,使用動態圖模式。不可避免的帶來了一系列問題。對于一個動態圖來說,面臨了以下三點問題:
性能:
- 動態圖在每次執行時都需要重新構建計算圖,這可能導致額外的開銷。
- 靜態圖只需構建一次,然后可以重復高效執行。
- 優化難度:
- 動態圖難以進行全局優化,因為圖結構在運行時可能會改變。
- 靜態圖允許更多的編譯時優化,如內存分配優化、算子融合等。
- 內存使用:
- 動態圖可能需要更多的運行時內存,因為它需要保持 Python 解釋器和相關對象的活躍狀態。
- 靜態圖可以更有效地管理內存,尤其是在推理階段。
所以從 torch 開始,我們第一步要做的就是動轉靜。拿到靜態圖才能更好的做整體性能上的優化!
美好的愿景:ONNX
ONNX,全稱 Open Neural Network Exchange,是人工智能領域中一個引人入勝的故事。它的誕生源于一個美好的愿景:在紛繁復雜的深度學習世界中架起一座溝通的橋梁。
2017 年的硅谷,各大科技巨頭都在人工智能領域奮力拼搏。Facebook(現在的Meta)和 Microsoft 這兩個看似競爭對手的公司,卻因為一個共同的夢想走到了一起。他們希望打破AI框架之間的壁壘,讓不同平臺上訓練的模型能夠自由遷移。就這樣,ONNX 項目應運而生。
聽起來 ONNX 是不同模型間完美的橋梁,最后聚合到 ONNX 完成推理是很開心和能接受的事情。但是聽起來越完美的事情就面臨越多的問題,首先是對 ONNX 來說。ONNX 模型在某些情況下可能比原生框架的模型運行得慢。這主要是因為?ONNX 作為一個中間表示,可能無法充分利用特定硬件或框架的優化特性。想象一下,它就像是一個通用翻譯器,雖然能夠讓不同語言的人交流,但可能會損失一些語言中的微妙之處和效率。
除此之外,AI 領域發展的太快。ONNX 并不一定能很好的表示 torch 中各種各樣的算子,導致模型轉換成 ONNX 失敗。談回之前簡單的網絡,我們如何把它轉換成 ONNX 形式呢?請看:
# 將模型轉換為ONNX格式
import torch.onnx# 創建一個示例輸入張量
dummy_input = torch.randn(1, input_size)# 指定ONNX文件的輸出路徑
output_path = "simple_net.onnx"# 導出模型到ONNX
torch.onnx.export(model, # 要轉換的模型dummy_input, # 模型的輸入樣例output_path, # 輸出的ONNX文件路徑export_params=True, # 存儲訓練好的參數權重opset_versinotallow=11, # ONNX算子集版本do_constant_folding=True, # 是否執行常量折疊優化input_names=['input'], # 輸入節點的名稱output_names=['output'], # 輸出節點的名稱dynamic_axes={'input' : {0 : 'batch_size'}, # 批處理維度動態'output' : {0 : 'batch_size'}})print(f"Model has been converted to ONNX and saved as {output_path}")
廠家的秘方:OpenVINO、TensorRT
不同廠家都有自己的推理秘制配方:推理引擎。這種趨勢反映了 AI 領域的激烈競爭和快速創新。每家公司都希望在這場技術革命中占據有利地位,而自研推理引擎成為了關鍵戰略。
這種做法的核心原因在于硬件差異化和性能優化。不同公司擁有各自獨特的硬件架構,如英特爾的 CPU、NVIDIA 的 GPU 或谷歌的 TPU。為了充分發揮這些硬件的潛力,定制化的推理引擎成為必然選擇。這些引擎能夠針對特定硬件進行深度優化,實現最佳的性能和效率。這其中,我將為大家簡單介紹兩種。分別是 OpenVINO 和 TensorRT。
(一)OpenVINO
OpenVINO
讓我們先將目光投向 OpenVINO。它的故事始于英特爾的實驗室,在那里,一群充滿激情的工程師夢想著如何讓人工智能的力量觸手可及。2018 年,OpenVINO 正式誕生,其名字中的"VINO"代表"Visual Inference and Neural network Optimization",寓意著它要為視覺智能和神經網絡優化開辟一條康莊大道。
OpenVINO 可在英特爾?硬件上擴展計算機視覺和非視覺工作負載,從而最大限度地提高性能。它通過從邊緣到云的高性能,人工智能和深度學習推理來加速應用程序。
關于OpenVINO的模型轉換
import subprocess
import sysdef convert_onnx_to_openvino(onnx_model_path, output_dir):cmd = [sys.executable, # 使用當前Python解釋器"-m", "mo", # 調用model optimizer"--input_model", onnx_model_path,"--output_dir", output_dir,"--data_type", "FP32"]subprocess.run(cmd, check=True)print(f"Model has been converted to OpenVINO IR format and saved in {output_dir}")# 使用示例
onnx_model_path = "simple_net.onnx"
output_dir = "openvino_model"convert_onnx_to_openvino(onnx_model_path, output_dir)
這個轉換過程和 ONNX 很像,在 OpenVINO 具體執行流程里分為反序列化,輸入定義和前向執行幾方面。
(二)TensorRT
TensorRT
與此同時,在硅谷的另一端,NVIDIA 的工程師們也在編織著自己的 AI 夢想。2017 年,TensorRT 橫空出世,它的名字中的"RT"代表"Runtime",彰顯了它對高性能推理的執著追求。
TensorRT 就像是一位技藝精湛的魔法師,它能夠將龐大復雜的神經網絡模型變成小巧高效的推理引擎。它的法術可以讓模型在 NVIDIA 的 GPU 上飛馳,實現令人瞠目的低延遲和高吞吐量。想象一下,它就像是給AI裝上了火箭推進器,讓智能決策的速度突破音障。
TensorRT 可用于對超大規模數據中心,嵌入式平臺或自動駕駛平臺進行推理加速。TensorRT 現已能支持 TensorFlow,Caffe,Mxnet,Pytorch 等幾乎所有的深度學習框架,將 TensorRT 和 NVIDIA 的 GPU 結合起來,能在幾乎所有的框架中進行快速和高效的部署推理。但可惜,TensorRT 是一個閉源的庫。
關于TensorRT模型的轉換
我們一般會給 TensorRT 的模型叫為 engine,我們可以使用 trt 提供的命令行工具,trtexec進行轉換
trtexec --notallow=simple_net.onnx --saveEngine=simple_net.trt --explicitBatch
推理引擎:類似的執行流程
- OpenVINO模型部署分為兩個部分:模型優化器和推理引擎。
模型優化器將訓練好的模型轉換為推理引擎可以識別的中間表達 –IR 文件,并在轉換過程中對模型進行優化。推理引擎接受經過模型優化器轉換并優化的網絡模型,為 Intel 的各種計算設備提供高性能的神經網絡推理運算。
- TensorRT 模型部署也是分為兩個部分:build 和 deployment 。
build:這個階段主要完成模型轉換,將不同框架的模型轉換到 TensorRT。模型轉換時會完成前述優化過程中的層間融合,精度校準。這一步的輸出是一個針對特定 GPU 平臺和網絡模型的優化過的 TensorRT 模型,這個 TensorRT 模型可以序列化存儲到磁盤或內存中。存儲到磁盤中的文件稱之為 plan file。deployment:將上面一個步驟中的 plan 文件首先反序列化,并創建一個 runtime engine,然后就可以輸入數據(比如測試集或數據集之外的圖片),然后輸出分類向量結果或檢測結果。
寫在最后
模型部署以加速為最終目的,首先就會拋棄易用性。這里特指靜態圖,在固定的范圍內做極致的優化。除了模型上的優化,不同硬件廠商更會在貼近不同硬件上做各種底層上的優化。以獲得在特定芯片上極致的性能。請期待后續部署教程吧~
三、onnxruntime部署YOLOv8分割模型詳細教程
本文將詳細介紹如何使用onnxruntime框架來部署YOLOv8分割模型,為了方便理解,代碼采用Python實現。
0. 引言
我之前寫的文章《基于YOLOv8分割模型實現垃圾識別》介紹了如何使用??YOLOv8?
??分割模型來實現垃圾識別,主要是介紹如何用自定義的數據集來訓練??YOLOv8?
??分割模型。那么訓練好的模型該如何部署呢???YOLOv8?
?分割模型相比檢測模型多了一個實例分割的分支,部署的時候還需要做一些后處理操作才能得到分割結果。
本文將詳細介紹如何使用??onnxruntime?
??框架來部署??YOLOv8?
??分割模型,為了方便理解,代碼采用??Python?
?實現。
1. 準備工作
- 「安裝onnxruntime」?
?onnxruntime?
??分為??GPU?
??版本和??CPU?
??版本,均可以通過??pip?
?直接安裝:
pip install onnxruntime-gpu #安裝GPU版本 pip install onnxruntime #安裝CPU版本
「注意:」???GPU?
??版本和??CPU?
??版本建議只選其中一個安裝,否則默認會使用??CPU?
?版本。
- 「下載?
?YOLOv8?
?分割模型權重」??Ultralytics?
??官方提供了用??COCO?
??數據集訓練的模型權重,我們可以直接從官方網站??https://docs.ultralytics.com/tasks/segment/?
??下載使用,本文使用的模型為??yolov8m-seg.pt?
?。
- 「轉換onnx模型」調用下面的命令可以把?
?YOLOv8m-seg.pt?
??模型轉換為??onnx?
?格式的模型:
yolo task=segment mode=export model=yolov8m-seg.pt format=onnx
轉換成功后得到的模型為??yolov8m-seg.onnx?
?。
2. 模型部署
2.1 加載onnx模型
首先導入??onnxruntime?
??包,然后調用其??API?
?加載模型即可:
import onnxruntime as ort session = ort.InferenceSession("yolov8m-seg.onnx", providers=["CUDAExecutionProvider"])
因為我使用的是??GPU?
??版本的??onnxruntime?
??,所以??providers?
??參數設置的是??"CUDAExecutionProvider"?
??;如果是??CPU?
??版本,則需設置為??"CPUExecutionProvider"?
?。
模型加載成功后,我們可以查看一下模型的輸入、輸出層的屬性:
for input in session.get_inputs(): print("input name: ", input.name) print("input shape: ", input.shape) print("input type: ", input.type) for output in session.get_outputs(): print("output name: ", output.name) print("output shape: ", output.shape) print("output type: ", output.type)
結果如下:
input name: images
input shape: [1, 3, 640, 640]
input type: tensor(float)
output name: output0
output shape: [1, 116, 8400]
output type: tensor(float)
output name: output1
output shape: [1, 32, 160, 160]
output type: tensor(float)
從上面的打印信息可以知道,模型有一個尺寸為??[1, 3, 640, 640]?
??的輸入層和兩個尺寸分別為??[1, 116, 8400]?
??和??[1, 32, 160, 160]?
?的輸出層。
2.2 數據預處理
數據預處理采用??OpenCV?
??和??Numpy?
?實現,首先導入這兩個包
import cv2
import numpy as np
用??OpenCV?
??讀取圖片后,把數據按照??YOLOv8?
?的要求做預處理
??image?=?cv2.imread("soccer.jpg") image_height,?image_width,?_?=?image.shape input_tensor?=?prepare_input(image,?model_width,?model_height) print("input_tensor?shape:?",?input_tensor.shape)?
?
其中預處理函數??prepare_input?
?的實現如下:
def prepare_input(bgr_image, width, height): image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB) image = cv2.resize(image, (width, height)).astype(np.float32) image = image / 255.0 image = np.transpose(image, (2, 0, 1)) input_tensor = np.expand_dims(image, axis=0) return input_tensor
處理流程如下:
1. 把OpenCV讀取的BGR格式的圖片轉換為RGB格式;
2. 把圖片resize到模型輸入尺寸640x640;
3. 對像素值除以255做歸一化操作;
4. 把圖像數據的通道順序由HWC調整為CHW;
5. 擴展數據維度,將數據的維度調整為NCHW。
經過預處理后,輸入數據??input_tensor?
??的維度變為??[1, 3, 640, 640]?
?,與模型的輸入尺寸一致。
2.3 模型推理
輸入數據準備好以后,就可以送入模型進行推理:
outputs = session.run(None, {session.get_inputs()[0].name: input_tensor})
前面我們打印了模型的輸入輸出屬性,可以知道模型有兩個輸出分支,其中一個??output0?
??是目標檢測分支,另一個??output1?
?則是實例分割分支,這里打印一下它們的尺寸看一下
#squeeze函數是用于刪除shape中為1的維度,對output0做transpose操作是為了方便后續操作
output0 = np.squeeze(outputs[0]).transpose()
output1 = np.squeeze(outputs[1])
print("output0 shape:", output0.shape)
print("output1 shape:", output1.shape)
結果如下:
output0 shape: (8400, 116)
output1 shape: (32, 160, 160)
處理后目標檢測分支的維度為??[8400, 116]?
??,表示模型總共可以檢測出??8400?
??個目標(大部分是無效的目標),每個目標包含??116?
??個參數。剛接觸??YOLOv8?
??分割模型的時候可能會對??116?
??這個數字感到困惑,這里有必要解釋一下:每個目標的參數包含??4?
??個坐標屬性(??x,y,w,h?
??)、??80?
??個類別置信度和??32?
??個實例分割參數,所以總共是??116?
??個參數。實例分割分支的維度為??[32, 160, 160]?
??,其中第一個維度??32?
??與目標檢測分支中的??32?
??個實例分割參數對應,后面兩個維度則由模型輸入的寬和高除以??4?
??得到,本文所用的模型輸入寬和高都是??640?
??,所以這兩個維度都是??160?
?。
2.4 后處理
首先把目標檢測分支輸出的數據分為兩個部分,把實例分割相關的參數從中剝離。
boxes = output0[:, 0:84]
masks = output0[:, 84:]
print("boxes shape:", boxes.shape)
print("masks shape:", masks.shape)
boxes shape: (8400, 84)
masks shape: (8400, 32)
然后實例分割這部分數據??masks?
??要與模型的另外一個分支輸出的數據??output1?
??做矩陣乘法操作,在這之前要把??output1?
?的維度變換為二維。
output1 = output1.reshape(output1.shape[0], -1)
masks = masks @ output1
print("masks shape:", masks.shape)
masks shape: (8400, 25600)
做完矩陣乘法后,就得到了??8400?
??個目標對應的實例分割掩碼數據??masks?
??,可以把它與目標檢測的結果??boxes?
?拼接到一起。
detections = np.hstack([boxes, masks])
print("detections shape:", detections.shape)
detections shape: (8400, 25684)
到這里讀者應該就能理解清楚了,??YOLOv8?
??模型總共可以檢測出??8400?
??個目標,每個目標的參數包含??4?
??個坐標屬性(??x,y,w,h?
??)、??80?
??個類別置信度和一個??160x160=25600?
?大小的實例分割掩碼。
由于??YOLOv8?
??模型檢測出的??8400?
?個目標中有大量的無效目標,所以先要通過置信度過濾去除置信度低于閾值的目標,對于滿足置信度滿足要求的目標還需要通過非極大值抑制(NMS)操作去除重復的目標。
objects = []
for row in detections: prob = row[4:84].max() if prob < 0.5: continue class_id = row[4:84].argmax() label = COCO_CLASSES[class_id] xc, yc, w, h = row[:4] // 把x1, y1, x2, y2的坐標恢復到原始圖像坐標 x1 = (xc - w / 2) / model_width * image_width y1 = (yc - h / 2) / model_height * image_height x2 = (xc + w / 2) / model_width * image_width y2 = (yc + h / 2) / model_height * image_height // 獲取實例分割mask mask = get_mask(row[84:25684], (x1, y1, x2, y2), image_width, image_height) // 從mask中提取輪廓 polygon = get_polygon(mask, x1, y1) objects.append([x1, y1, x2, y2, label, prob, polygon, mask]) // NMS
objects.sort(key=lambda x: x[5], reverse=True)
results = []
while len(objects) > 0: results.append(objects[0]) objects = [object for object in objects if iou(object, objects[0]) < 0.5]
這里重點講一下獲取實例分割掩碼的過程。
前面說了每個目標對應的實例分割掩碼數據大小為??160x160?
??,但是這個尺寸是對應整幅圖的掩碼。對于單個目標來說,還要從這個??160x160?
??的掩碼中去截取屬于自己的掩碼,截取的范圍由目標的??box?
??決定。上面的代碼得到的??box?
??是相對于原始圖像大小,截取掩碼的時候需要把??box?
??的坐標轉換到相對于??160x160?
??的大小,截取完后再把這個掩碼的尺寸調整回相對于原始圖像大小。截取到??box?
??大小的數據后,還需要對數據做??sigmoid?
??操作把數值變換到??0?
??到??1?
??的范圍內,也就是求這個??box?
??范圍內的每個像素屬于這個目標的置信度。最后通過閾值操作,置信度大于??0.5?
?的像素被當做目標,否則被認為是背景。
具體實現的代碼如下:
def get_mask(row, box, img_width, img_height): mask = row.reshape(160, 160) x1, y1, x2, y2 = box // box坐標是相對于原始圖像大小,需轉換到相對于160*160的大小 mask_x1 = round(x1 / img_width * 160) mask_y1 = round(y1 / img_height * 160) mask_x2 = round(x2 / img_width * 160) mask_y2 = round(y2 / img_height * 160) mask = mask[mask_y1:mask_y2, mask_x1:mask_x2] mask = sigmoid(mask) // 把mask的尺寸調整到相對于原始圖像大小 mask = cv2.resize(mask, (round(x2 - x1), round(y2 - y1))) mask = (mask > 0.5).astype("uint8") * 255 return mask
這里需要注意的是,??160x160?
?是相對于模型輸入尺寸為??640x640?
?來的,如果模型輸入是其他尺寸,那么上面的代碼需要做相應的調整。
如果需要檢測的是下面這個圖片:
通過上面的代碼可以得到最左邊那個人的分割掩碼為
但是我們需要的并不是這樣一張圖片,而是需要用于表示這個目標的輪廓,這可以通過??OpenCV?
??的??findContours?
??函數來實現。??findContours?
??函數返回的是一個用于表示該目標的點集,然后我們可以在原始圖像中用??fillPoly?
?函數畫出該目標的分割結果。
全部目標的檢測與分割結果如下:
3. 一點其他的想法
從前面的部署過程可以知道,做后處理的時候需要對實例分割的數據做矩陣乘法、??sigmoid?
??激活、維度變換等操作,實際上這些操作也可以在導出模型的時候集成到??onnx?
?模型中去,這樣就可以簡化后處理操作。
首先需要修改??ultralytics?
??代碼倉庫中??ultralytics/nn/modules/head.py?
??文件的代碼,把??Segment?
??類??Forward?
?函數最后的代碼修改為:
if self.export: output1 = p.reshape(p.shape[0], p.shape[1], -1) boxes = x.permute(0, 2, 1) masks = torch.sigmoid(mc.permute(0, 2, 1) @ output1) out = torch.cat([boxes, masks], dim=2) return out
else: return (torch.cat([x[0], mc], 1), (x[1], mc, p))
然后修改??ultralytics/engine/exporter.py?
??文件中??torch.onnx.export?
??的參數,把模型的輸出數量改為??1?
?個。
代碼修改完成后,執行命令??pip install -e '.[dev]'?
??使之生效,然后再重新用??yolo?
??命令導出模型。用??netron?
??工具可以看到模型只有一個??shape?
??為??[1,8400,25684]?
?的輸出。
這樣在后處理的時候就可以直接去解析??box?
??和??mask?
??了,并且??mask?
??的數據不需要進行??sigmoid?
?激活。
參考資料
1.How to implement instance segmentation using YOLOv8 neural network
2.https://github.com/AndreyGermanov/yolov8_segmentation_python