本文首發(fā)于公眾號【DeepDriving】，歡迎關(guān)注。
0. 引言
我之前寫的文章《基于YOLOv8分割模型實現(xiàn)垃圾識別》介紹了如何使用YOLOv8分割模型來實現(xiàn)垃圾識別，主要是介紹如何用自定義的數(shù)據(jù)集來訓(xùn)練YOLOv8分割模型。那么訓(xùn)練好的模型該如何部署呢？YOLOv8分割模型相比檢測模型多了一個實例分割的分支，部署的時候還需要做一些后處理操作才能得到分割結(jié)果。
本文將詳細(xì)介紹如何使用onnxruntime框架來部署YOLOv8分割模型，為了方便理解，代碼采用Python實現(xiàn)。
1. 準(zhǔn)備工作
-
安裝onnxruntime
onnxruntime分為GPU版本和CPU版本，均可以通過pip直接安裝：pip install onnxruntime-gpu #安裝GPU版本 pip install onnxruntime #安裝CPU版本注意:
GPU版本和CPU版本建議只選其中一個安裝，否則默認(rèn)會使用CPU版本。 -
下載
YOLOv8分割模型權(quán)重Ultralytics官方提供了用COCO數(shù)據(jù)集訓(xùn)練的模型權(quán)重，我們可以直接從官方網(wǎng)站https://docs.ultralytics.com/tasks/segment/下載使用，本文使用的模型為yolov8m-seg.pt。
-
轉(zhuǎn)換onnx模型
調(diào)用下面的命令可以把
YOLOv8m-seg.pt模型轉(zhuǎn)換為onnx格式的模型：yolo task=segment mode=export model=yolov8m-seg.pt format=onnx轉(zhuǎn)換成功后得到的模型為
yolov8m-seg.onnx。
2. 模型部署
2.1 加載onnx模型
首先導(dǎo)入onnxruntime包，然后調(diào)用其API加載模型即可:
import onnxruntime as ort
session = ort.InferenceSession("yolov8m-seg.onnx", providers=["CUDAExecutionProvider"])
因為我使用的是GPU版本的onnxruntime，所以providers參數(shù)設(shè)置的是"CUDAExecutionProvider"；如果是CPU版本，則需設(shè)置為"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)
結(jié)果如下：
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 數(shù)據(jù)預(yù)處理
數(shù)據(jù)預(yù)處理采用OpenCV和Numpy實現(xiàn)，首先導(dǎo)入這兩個包
import cv2
import numpy as np
用OpenCV讀取圖片后，把數(shù)據(jù)按照YOLOv8的要求做預(yù)處理
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)
其中預(yù)處理函數(shù)prepare_input的實現(xiàn)如下：
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格式的圖片轉(zhuǎn)換為RGB格式；
2. 把圖片resize到模型輸入尺寸640x640；
3. 對像素值除以255做歸一化操作；
4. 把圖像數(shù)據(jù)的通道順序由HWC調(diào)整為CHW；
5. 擴(kuò)展數(shù)據(jù)維度，將數(shù)據(jù)的維度調(diào)整為NCHW。
經(jīng)過預(yù)處理后，輸入數(shù)據(jù)input_tensor的維度變?yōu)?code>[1, 3, 640, 640]，與模型的輸入尺寸一致。
2.3 模型推理
輸入數(shù)據(jù)準(zhǔn)備好以后，就可以送入模型進(jìn)行推理：
outputs = session.run(None, {session.get_inputs()[0].name: input_tensor})
前面我們打印了模型的輸入輸出屬性，可以知道模型有兩個輸出分支，其中一個output0是目標(biāo)檢測分支，另一個output1則是實例分割分支，這里打印一下它們的尺寸看一下
#squeeze函數(shù)是用于刪除shape中為1的維度，對output0做transpose操作是為了方便后續(xù)操作
output0 = np.squeeze(outputs[0]).transpose()
output1 = np.squeeze(outputs[1])
print("output0 shape:", output0.shape)
print("output1 shape:", output1.shape)
結(jié)果如下：
output0 shape: (8400, 116)
output1 shape: (32, 160, 160)
處理后目標(biāo)檢測分支的維度為[8400, 116]，表示模型總共可以檢測出8400個目標(biāo)（大部分是無效的目標(biāo)），每個目標(biāo)包含116個參數(shù)。剛接觸YOLOv8分割模型的時候可能會對116這個數(shù)字感到困惑，這里有必要解釋一下：每個目標(biāo)的參數(shù)包含4個坐標(biāo)屬性（x,y,w,h）、80個類別置信度和32個實例分割參數(shù)，所以總共是116個參數(shù)。實例分割分支的維度為[32, 160, 160]，其中第一個維度32與目標(biāo)檢測分支中的32個實例分割參數(shù)對應(yīng)，后面兩個維度則由模型輸入的寬和高除以4得到，本文所用的模型輸入寬和高都是640，所以這兩個維度都是160。
2.4 后處理
首先把目標(biāo)檢測分支輸出的數(shù)據(jù)分為兩個部分，把實例分割相關(guān)的參數(shù)從中剝離。
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)
然后實例分割這部分?jǐn)?shù)據(jù)masks要與模型的另外一個分支輸出的數(shù)據(jù)output1做矩陣乘法操作，在這之前要把output1的維度變換為二維。
output1 = output1.reshape(output1.shape[0], -1)
masks = masks @ output1
print("masks shape:", masks.shape)
masks shape: (8400, 25600)
做完矩陣乘法后，就得到了8400個目標(biāo)對應(yīng)的實例分割掩碼數(shù)據(jù)masks，可以把它與目標(biāo)檢測的結(jié)果boxes拼接到一起。
detections = np.hstack([boxes, masks])
print("detections shape:", detections.shape)
detections shape: (8400, 25684)
到這里讀者應(yīng)該就能理解清楚了，YOLOv8模型總共可以檢測出8400個目標(biāo)，每個目標(biāo)的參數(shù)包含4個坐標(biāo)屬性（x,y,w,h）、80個類別置信度和一個160x160=25600大小的實例分割掩碼。
由于YOLOv8模型檢測出的8400個目標(biāo)中有大量的無效目標(biāo)，所以先要通過置信度過濾去除置信度低于閾值的目標(biāo)，對于滿足置信度滿足要求的目標(biāo)還需要通過非極大值抑制（NMS）操作去除重復(fù)的目標(biāo)。
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的坐標(biāo)恢復(fù)到原始圖像坐標(biāo)
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]
這里重點講一下獲取實例分割掩碼的過程。
前面說了每個目標(biāo)對應(yīng)的實例分割掩碼數(shù)據(jù)大小為160x160，但是這個尺寸是對應(yīng)整幅圖的掩碼。對于單個目標(biāo)來說，還要從這個160x160的掩碼中去截取屬于自己的掩碼，截取的范圍由目標(biāo)的box決定。上面的代碼得到的box是相對于原始圖像大小，截取掩碼的時候需要把box的坐標(biāo)轉(zhuǎn)換到相對于160x160的大小，截取完后再把這個掩碼的尺寸調(diào)整回相對于原始圖像大小。截取到box大小的數(shù)據(jù)后，還需要對數(shù)據(jù)做sigmoid操作把數(shù)值變換到0到1的范圍內(nèi)，也就是求這個box范圍內(nèi)的每個像素屬于這個目標(biāo)的置信度。最后通過閾值操作，置信度大于0.5的像素被當(dāng)做目標(biāo)，否則被認(rèn)為是背景。
具體實現(xiàn)的代碼如下：
def get_mask(row, box, img_width, img_height):
mask = row.reshape(160, 160)
x1, y1, x2, y2 = box
// box坐標(biāo)是相對于原始圖像大小，需轉(zhuǎn)換到相對于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的尺寸調(diào)整到相對于原始圖像大小
mask = cv2.resize(mask, (round(x2 - x1), round(y2 - y1)))
mask = (mask > 0.5).astype("uint8") * 255
return mask
這里需要注意的是，160x160是相對于模型輸入尺寸為640x640來的，如果模型輸入是其他尺寸，那么上面的代碼需要做相應(yīng)的調(diào)整。
如果需要檢測的是下面這個圖片：
通過上面的代碼可以得到最左邊那個人的分割掩碼為
但是我們需要的并不是這樣一張圖片，而是需要用于表示這個目標(biāo)的輪廓，這可以通過OpenCV的findContours函數(shù)來實現(xiàn)。findContours函數(shù)返回的是一個用于表示該目標(biāo)的點集，然后我們可以在原始圖像中用fillPoly函數(shù)畫出該目標(biāo)的分割結(jié)果。
全部目標(biāo)的檢測與分割結(jié)果如下：
3. 一點其他的想法
從前面的部署過程可以知道，做后處理的時候需要對實例分割的數(shù)據(jù)做矩陣乘法、sigmoid激活、維度變換等操作，實際上這些操作也可以在導(dǎo)出模型的時候集成到onnx模型中去，這樣就可以簡化后處理操作。
首先需要修改ultralytics代碼倉庫中ultralytics/nn/modules/head.py文件的代碼，把Segment類Forward函數(shù)最后的代碼修改為：
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的參數(shù)，把模型的輸出數(shù)量改為1個。
代碼修改完成后，執(zhí)行命令pip install -e '.[dev]'使之生效，然后再重新用yolo命令導(dǎo)出模型。用netron工具可以看到模型只有一個shape為[1,8400,25684]的輸出。
文章來源：http://www.zghlxwxcb.cn/news/detail-856879.html
這樣在后處理的時候就可以直接去解析box和mask了，并且mask的數(shù)據(jù)不需要進(jìn)行sigmoid激活。文章來源地址http://www.zghlxwxcb.cn/news/detail-856879.html
4. 參考資料
- How to implement instance segmentation using YOLOv8 neural network
- https://github.com/AndreyGermanov/yolov8_segmentation_python
到了這里，關(guān)于AI模型部署 | onnxruntime部署YOLOv8分割模型詳細(xì)教程的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！
