前面我们得知YOLOv8
不但可以实现目标检测任务,还包揽了分类、分割、姿态估计等计算机视觉任务。在上一篇博文中,博主已经介绍了YOLOv8
如何实现分类,在这篇博文里,博主将介绍其如何将实例分割给收入囊中。
YOLOv8实例分割架构图
如下图所示,YOLOv8
采用了一种分割头与检测头相结合的方式来进行实例分割,在这个过程中,其会输出目标检测框与实例分割蒙版。
(先前博主以为这个是语义分割的,但后经人指正才发觉是实例分割,这也同时解答了我一些疑惑)
输出结果图像如下:
经典语义分割模型结构
为了让我们更好的理解语义分割模型,我们以最经典的语义分割模型UNet
为例,可以看到其最终的结果要与原图像大小相同,但最终的维度(n)
会有差别,这与我们确定使用的mask
的数量有关。
YOLOv8实例分割模型结构
YOLOv8的实例分割YOLOv8`的目标检测模型结构即为接近,区别在于在最后的目标检测头基础上添加了实例分割头,同时其最终的实例分割头也是具有三种尺度的:
下图中对各个模块进行了编号,大家可以与yaml
的模型文件进行对照
YOLOv8检测头(可忽略)
那么我们看下这个分割头到底是如何定义的,分割头继承了检测头:
检测头代码如下:我们可以看到ultralytics
更新了检测头(加入了YOLOv10
,博主这里将该方法删掉了,因为用不到),其创新点为混合匹配机制,故在检测头中多出了forward_end2end
:
YOLOv8分割头
分割头代码如下:
输入到分割头的图像存储在list
中,共有三个不同尺度,这与YOLOv8
目标检测是相同的
上述第一个操作便是Proto
操作 ,传入的是第一尺度的输出特征图,Proto
的功能是针对x[0]
进行卷积,将原来80x80大小的feature通过上采样变为160x160,这个图像是基础蒙版(mask)。
随后将x输入到cv4
模块(期内包含3个模块组成))(即图像中的三个不同尺度的操作),cv4结构如下:
其过程如下图所示:
得到的即为mask
随后进入检测的前向传播过程,因为YOLOv8
本身就是做的检测,因此这个结果还是要进入检测头:
cv2
中包含三个模块,最终的输出大小不变,通道数均为64
,即为(64,80,80)(64,40,40)(64,20,20)
cv3
中也是三个模块,图像大小依旧不变,通道维度变为80
,即(80,80,80)(80,40,40)(80,20,20)
最终将其使用torch.ca
t进行拼接,得到(144,80,80)(144,40,40)(144,20,20)
随后便是推理的后处理过程,即对输出的这三个尺度的图像进行解码:
self.no = nc + self.reg_max * 4
,其中reg_max
是根据YOLOv8
不同模型大小设定的,即 scale 4/8/12/16/20 for n/s/m/l/x)
,此处reg_max=16
self.anchors为torch.Size([2, 8400])
, self.strides为torch.Size([1, 8400])
shape
为torch.Size([1, 144, 80, 80])
144=64+80,这个64是预测的box的值,最后还要进行转换
根据x_cat
进行拆分,得到预测的box
与cls
,box
即为(1,64,8400)
,cls
为(1,80,8400)
随后通过DEL
模块对box
进行分解:
得到的box
即为(1,4,8400)
DE
L中的Conv2d
没有梯度,即参数不会更新,这个模块作用便是将64
分解为4*16
,进而得到4*1
最后将 dbox
和 cls
(类别)返回
其维度为(84,8400)
,84=80+4
,8400
代表预测的目标个数
最终返回数据:
其中mc
为(1,32,8400)x是一个元组,x[0]为(1,84,8400)x[1】为列表,包含(1,144,80,80)(1,144,40,40),(1,144,20,20),p为(1,32,160,160)
p为基础蒙版
返回的数据:
其中(1,32,8400)
即为预测的mask
非极大值抑制
在推理过程中,博主使用的图像大小为(3,480,640)
,所有最后得到的数据维度为(1,116,6300)
其中,116=84(80+4)+32
,这是因为YOLOv8
中不仅要完成语义分割还要实现目标检测,其中(1,32,6300)
是用于语义分割的。
而6300=60*80+30*40+15*20
下面的分解代码证实了这一点,即mask
的数量(nm
)为116-80-4=32
,mask
开始坐标为80+4=84
设定输出:
得到的output为38
个值,其中38=4+class_score+class+32
(保存检测与分割结果)
随后筛选出的大于阈值的类别,得到36
个,即(1,36,116)
,这里的36
指的是符合的个数,是从6300
中筛选出的。
将box
,类别 以及分割mask
分开:
box
为(36,4)
,cls
为(36,80)
,mask
为 (36,32)
随后再从类别中选出最大的
得到的conf
为分值,j
为坐标(代表类别),维度均为(36,1)
,并将这些数据再次拼接到一起,得到(36,38)
,其中36
为目标个数,38
为4+1+1+32
,即 box+conf+cls_id+mask
返回的 i
为 tensor([20, 24, 3, 32, 34], device='cuda:0')
,这里给出的i
是36
个中经过筛选后的检测框编号,最终将x中的目标筛选出存储到output
中,可以看到output
是一个列表,存放的是每个batch
的结果,由于在预测时只输入一张图像,故里面只有一个数据,筛选出的结果为(5,38)
,即有5
个目标。
后处理过程
那么,这个mask
要如何使用呢。我们接下来看一下其后处理过程
后处理过程中传入的参数为preds即预测的结果,即YOLOv8分割头输出的结果
img是输入的图像(归一化后的),orig_img是原始图像
在后处理过程的刚开始,便是利用非极大值抑制来筛选出部分数据:
得到的结果 p
即为(5,38)
判断preds[1]
是否是tuple
类型,是,则为preds[1][-1]
,即为(1,32,120,160)
随后进行下面的循环(预测只有一张图像,故只有一轮)
我们可以看到mask
的处理结果:
process_mask
方法是如何处理的呢?
我们先看一下其传入的参数,proto
为(1,32,120,160)
,pred
为(5,38)
,取从6
到38
,即只是mask
的32
维数据,即为(5,32)
,同时还有bbox
为 (5,4)
,img.shape[2:]
为宽高
接下来便是将保证mask在Bbox内。
得到的mask依旧为(5,120,160)
,随后对mask
进行上采样,使其与原本的图像大小一样的,这里就已经是蒙版了,通过插值的方式进行上采样,得到的mask为(5,480,640)
mask的理解
在检测头(分割头)中输出的32 维的向量可以看作是与每个检测框关联的分割
mask
的系数或权重。针对于分割头的输出
1x32x160x160
,一个关键的概念是prototype masks
。它是一个固定数量(32)的基础mask,每个mask
的尺寸为160×160
。这些基础mask
并不直接对应于任何特定的物体或类别,而是被设计为可以线性组合来表示任何可能的物体mask
。简单来说,模型不直接预测每个物体的完整
mask
,而是预测一组基本的masks
(称为prototype masks
)以及每个物体如何组合这些masks
(权重/系数)。这种方法的好处是,模型只需要预测一个较小的mask
张量,然后可以通过简单的矩阵乘法将这些小mask
组合成完整的物体masks
。大家可以把它类比于线性代数中基向量的概念,空间中的任何一个向量是不是都可以表示为一组基向量的线性组合,那么其中的
prototype masks
即32x160x160
的mask
张量可以把它理解为一组基向量,而之前在检测框中的32
维向量可以理解为组合这一组基向量的权重或者说系数。当我们从检测头得到一个
32
维的向量,分割头得到32
个基础masks
时,这个32
维的向量实际上表示了如何组合这些基础masks
来得到一个特定物体的 mask。具体来说,我们用这个32
维向量对32
个基础 masks进行线性组合,从而得到与检测框关联的最终 mask。简单来说,这就像你现在有 32 种不同的颜料,检测头给你一个配方(32 维向量),告诉你如何混合这些颜料来得到一个特定的颜色(最终的 mask)。这样做的优点是我们不需要为每个检测框都预测一个完整的
mask
,这个非常消耗内存和计算资源。相反,我们只需要预测一个相对较小的32
维向量和一个固定数量的基础masks
,然后在后处理中进行组合即可。
结果可视化
最后附上将结果可视化的代码