Track to Detect and Segment: An Online Multi-Object Tracker

Track to Detect and Segment: An Online Multi-Object Tracker

论文和代码：https://jialianwu.com/projects/TraDeS.html
在交流群里看到作者公布了前段时间收录于CVPR2021的多目标跟踪论文，花了半天的时间看了一下，做个笔记。正片文章最主要的内容在于构建了一个端到端的模型，将跟踪信息用于辅助检测或者分割，在多目标跟踪（MOT）和像素级多目标跟踪（MOTS）上均达到了sota。

---

Introduction

现如今主流的MOT方法大致分为两种：track-by-detection（TBD）的二阶段方法和joint-detection-and-tracking的单阶段方法。TBD范式的方法将MOT分为两个独立的任务，先用一个现成的检测器检测出图像中的目标，再用其他的网络将其与现有轨迹相连，用的比较多的是Re-ID的网络。这种范式的问题便在检测和跟踪是两个独立的任务，分开训练，并不是一个端到端的模型，无法达到全局的最优解，而且通常来说速度较慢。为了解决这种二阶段带来的弊端，JDE提出了新的joint-detection-and-tracking范式，在一个前向过程中，同时进行检测和外貌特征提取，将MOT做到了端到端实现。
JDE范式依旧存在着问题：（1）虽然在这种范式中，骨干网络是两个任务共享的，但检测依然是独立在外的，没有用到跟踪的信息。作者认为，对于稳定连续的轨迹来说，检测是一个基础，反过来，跟踪线索也将有助于检测，尤其是在目标发生部分遮挡或者运动模糊的情况下；（2）在联合训练检测和跟踪任务时，普通的Re-ID损失和检测损失不太兼容，检测任务的目的是区分目标，使得类间距变大，类内距离缩小，然而Re-ID任务是关注同一类内部的不同之处。
出于以上思考，作者提出了一个新的在线跟踪模型TraDeS（Track to Detect and Segment），通过将跟踪的信息反哺给检测以及设计了一个与检测损失较兼容的Re-ID损失，来提升多目标跟踪的效果。整个模型最主要的在于两个模块：CVA（cost volume based association）模块和MFW（motion-guided feature warper）模块。CVA模块构建了一个cost volumn（类似于匹配代价矩阵），存储着两帧的embedding之间的相似度匹配的信息。接着，从cost volumn中推理出一个偏移量，即目标中心在两帧之间的位移。采用这个位移以及embedding，模型进行一个两轮的数据关联。然后，MFW模块通过偏移量，将前一帧的信息融合到当前帧中，以完成检测或者分割。
在CVA模块中，cost volumn用于监督Re-ID embedding的训练，这样目标和背景信息均在考虑的范围内，也就是说Re-ID模块也会考虑类别之间的不同。不仅提高Re-ID embedding的学习，而且在一起训练时不会伤害到检测模块的性能。此外，这里的偏移量是通过两个外貌embedding的之间的相似度学习到的，所以即使当目标有很大程度的移动或者视频帧率很低的时候，偏移量也能匹配到目标。因此，这个偏移量也可以作为稳定的跟踪信息传入到MFW模块中，来提高检测的效果。比如说，在当前帧，某个目标被部分遮挡了，但它在之前肯定是有不被遮挡的时候，将之前帧跟踪的信息融合到当前帧，可以使模型很好地应对这种部分遮挡的情况。

---

Preliminaries

cost volumn方法在深度估计和光流法中的应用很成熟，一般用于对两帧之间的像素进行关联。作者由此受启发，通过cost volumn来学习外貌embedding和偏移量，实验证明cost volumn是有效的。
TraDeS采用的检测框架为CenterNet，一开始的输入为图像$\mathbf{I}\in {\mathbb{R}}^{H\times W\times 3}$，由骨干网络提取出的特征为$\mathbf{f}\in {\mathbb{R}}^{H_F\times W_F\times 64}$，$H_F=\frac{H}{4}$，$W_F=\frac{W}{4}$。随后，将$\mathbf{f}$输入进几个head，得到分类的中心点热度图$\mathbf{P}\in {\mathbb{R}}^{H_F\times W_F\times N_{cls}}$以及其他不同任务需要的特征图（比如用于预测2D或3D的目标大小特征图等等），其中$N_{cls}$为目标类别的数量。CenterNet检测器是通过目标的中心点（$\mathbf{P}$中的峰值）以及其他分支相应位置的预测来检测目标。和CenterTrack一样，TraDeS也是额外添加了一个预测偏移量的分支来完成跟踪${\mathbf{O}}^B\in {\mathbb{R}}^{H_F\times W_F\times 2}$。

---

TraDeS Tracker

总的来说，作者的目标是想用跟踪的信息来反馈检测的子网络，提出了Cost Volume based Association（CVA）模块来更好的学习Re-ID任务的embedding以及一些目标的运动信息和Motion-guided Feature Warper（MFW）模块来对融合之前帧的跟踪信息，以便于更好的进行检测。

Cost Volume based Association

此模块用来学习外貌embedding和偏移量，用来进行数据关联，学到的embedding不仅关注类内的不同，也关注类间的差距，与检测损失兼容性更好，CVA模块还会将跟踪的信息添加到MFW模块以加强检测的性能。
Cost Volume： 假设来自两帧图像的特征分别为${\mathbf{f}}^t$和${\mathbf{f}}^{t-\tau }$，通过一个简单的Re-ID网络来学习其外貌embedding，${\mathbf{e}}^t=\sigma \left({\mathbf{f}}^t\right)\in {\mathbb{R}}^{H_F\times W_F\times 128}$。通过两个embedding，构造一个cost volumn矩阵来存储两帧之间点与点之间的匹配相似度信息。为了更好的计算cost volumn矩阵，现将得到的embedding下采样两倍，得到${\mathbf{e}}^{\prime }\in {\mathbb{R}}^{H_C\times W_C\times 128}$，$H_C$和$W_C$很明显就是$H_F$和$W_F$的一半。通过以下公式进行计算：$$C_{i,j,k,l}={\mathbf{e}}_{i,j}^{\prime t}{\mathbf{e}}_{k,l}^{t-\tau \top }$$得到最终的cost volumn矩阵$\mathbf{C}\in {\mathbb{R}}^{H_C\times W_C\times H_C\times W_C}$，$C_{i,j,k,l}$代表着$t$时刻的$(i,j)$点与$t-\tau$时刻的$(k,l)$点之间的相似度。这里的一个点，也就是上面公式里的${\mathbf{e}}_{i,j}^{\prime t}$，代表原embedding里$(i,j)$位置的所有通道的值，是个一维向量。
Tracking Offset： 在上面$\mathbf{C}$的基础上，计算所有$t$时刻点到$t-\tau$时刻相应位置的偏移量$\mathbf{O}\in {\mathbb{R}}^{H_C\times W_C\times 2}$。
假设现在要计算目标$x$的偏移量，其$t$时刻的中心点的位置为$(i,j)$，那么便可以找到其在$\mathbf{C}$中对应的一个两维的cost volumn图${\mathbf{C}}_{i,j}\in {\mathbb{R}}^{H_C\times W_C}$，${\mathbf{C}}_{i,j}$里是目标$x$的中心点与$t-\tau$时刻所有点的匹配相似度。在拿到${\mathbf{C}}_{i,j}$之后，（1）将${\mathbf{C}}_{i,j}$进行两个方向的最大值池化，
池化核大小分别为$H_C\times 1$和$1\times W_C$，然后再接softmax函数，得到两个向量${\mathbf{C}}_{i,j}^W\in [0,1{]}^{1\times W_C}$和${\mathbf{C}}_{i,j}^H\in [0,1{]}^{H_C\times 1}$。这两个向量分别代表着目标$x$的中心点在$t-\tau$时刻中的水平和垂直位置的可能性。（2）定义了水平和垂直方向的两个偏移量模板，${\mathbf{M}}_{i,j}\in {\mathbb{R}}^{1\times W_C}$和${\mathbf{V}}_{i,j}\in {\mathbb{R}}^{H_C\times 1}$，计算方法如下：$$\{\begin{array}{ll}{M}_{i,j,l}=(l-j)\times s& 1\le l\le W_C\\ V_{i,j,k}=(k-i)\times s& 1\le k\le H_C\end{array}$$其中$s$是$e^{\prime }$的步长（推测应该就是下采样倍数，即当前特征图一个像素在原来的图像中代表多远）。$M_{i,j,l}$即为目标$x$出现在$t-\tau$时刻的$(\ast ,l)$位置的水平位移。最终，偏移量的计算为：$${\mathbf{O}}_{i,j}={\left[{\mathbf{C}}_{i,j}^{H\top }{\mathbf{V}}_{i,j},{\mathbf{C}}_{i,j}^W{\mathbf{M}}_{i,j}^{\top }\right]}^{\top }$$类似于从1到$W_C$的对${\mathbf{C}}_{i,j}^H$的积分。得到$\mathbf{O}$后，再将其上采样2倍，得到${\mathbf{O}}^C\in {\mathbb{R}}^{H_F\times W_F\times 2}$。这就是最终的偏移量，也会输入到MFW模块用于之后帧的检测。
Training： 由于之前的普通Re-ID损失注重类内不同，与检测损失不是很兼容，所以作者提出用cost volumn来监督网络的学习。所以作者设定，当$t$时刻的点$(i,j)$出现在$t-\tau$时刻的$(k,l)$位置时，令$Y_{i,j,k,l}=1$，否则$Y_{i,j,k,l}=0$。CVA模块的损失定义为：$$L_{CVA}=\frac{-1}{{\sum }_{ijkl}{Y}_{ijkl}}\sum _{ijkl}\{\begin{array}{clc}{\alpha }_1\log \left(C_{i,j,l}^W\right)+{\alpha }_2\log \left(C_{i,j,k}^H\right)& \text{ if }{Y}_{ijkl}=1& \\ 0& \text{ otherwise }& \end{array}$$其中，${\alpha }_1={\left(1-C_{i,j,l}^W\right)}^{\beta }$，${\alpha }_2={\left(1-C_{i,j,k}^H\right)}^{\beta }$，不难看出这是两个方向上的focal loss形式的逻辑回归损失，$\beta$即为focal loss中的超参数。由于$C_{i,j,l}^W$和$C_{i,j,k}^H$是由softmax函数计算得到，所以其中并不只是包含了位置$(i,j)$和位置$(k,l)$之间的相似度信息，还包含了$(i,j)$和前一帧中其他点的相似度信息，所以当$C_{i,j,l}^W$和$C_{i,j,k}^H$向着1的方向进行优化时，可以使得目标不仅更加靠近前一帧中的自己，还能使得目标和其他目标以及背景之间的区别更大。

Motion-guided Feature Warper

MFW模块旨在采用CVA模块得到的偏移量$O^C$作为跟踪信息，用过可变卷积将前几帧的信息传播到当前帧，使得在当前帧进行的检测更加准确。
Temporal Propagation： 将$O^C$送入一个3*3的卷积后，得到${\mathbf{O}}^D\in {\mathbb{R}}^{H_F\times W_F\times 2K^2}$作为DCN要使用的偏移量，$K$为DCN卷积核的宽或高。为了获得更多的跟踪上的信息，将${\mathbf{f}}^t-{\mathbf{f}}^{t-\tau }$也作为输入，与$O^C$一起进行卷积。由于TraDeS的检测或者分割是基于目标中心点的，所以直接对${\mathbf{f}}^{t-\tau }$进行卷积不是很适合。作者想要对${\mathbf{f}}^{t-\tau }$进行一个类似中心点的注意力操作，先从${\mathbf{P}}^{t-\tau }$中提取到${\mathbf{P}}_{agn}^{t-\tau }\in {\mathbb{R}}^{H_F\times W_F\times 1}$，压缩${\mathbf{P}}^{t-\tau }$中的通道维度，因为并不需要其中的类别信息。提取到${\mathbf{P}}_{agn}^{t-\tau }$后，通过以下方法来得到改善后的${\overline{\mathbf{f}}}_q^{t-\tau }$：
$${\overline{\mathbf{f}}}_q^{t-\tau }={\mathbf{f}}_q^{t-\tau }\circ {\mathbf{P}}_{agn}^{t-\tau },q=1,2,\dots ,64$$其中的$\circ$为Hadamard积，即按位相乘。最终用于增强检测任务的特征便为：$${\hat{\mathbf{f}}}^{t-\tau }=DCN\left({\mathbf{O}}^D,{\overline{\mathbf{f}}}^{t-\tau }\right)\in {\mathbb{R}}^{H_F\times W_F\times 64}$$
Feature Enhancement： 有了之前帧的信息，在当前帧被遮挡的目标，便更容易被检测出来了。为了提高鲁棒性，作者融合了前$T$帧的信息：$${\tilde{\mathbf{f}}}_q^t={\mathbf{w}}^t\circ {\mathbf{f}}_q^t+\sum _{\tau =1}^T{\mathbf{w}}^{t-\tau }\circ {\hat{\mathbf{f}}}_q^{t-\tau },q=1,2,\dots ,64$$其中，${\mathbf{w}}^t\in {\mathbb{R}}^{H_F\times W_F\times 1}$为$t$时刻的权重，${\sum }_{\tau =0}^T{\mathbf{w}}_{i,j}^{t-\tau }=1$，与Learning Spatial Fusion for Single-Shot Object Detection中类似，$\mathbf{w}$是自适应的，通过两层卷积和softmax函数预测得到。作者实验证明，这种自适应的权重比平均求和的效果好一点。融合之后的特征${\tilde{\mathbf{f}}}_q^t$输入到网络的head中，用来预测当前帧的检测框或者分割掩膜，有效地恢复消失的目标，降低FN，得到更高的MOTA和IDF1，使得轨迹更完整。

Tracklet Generation

上图是TraDeS的整体结构（一眼看上去还是挺复杂的……但图画得好啊）。MFW模块输出融合后的特征${\tilde{\mathbf{f}}}_q^t$后，将其输入进网络的各个head，预测出当前帧的2Dbbox或者3Dbbox等等，再与之前的轨迹进行关联。
Head Networks： 网络的每个head均由几层轻量级的卷积层组成。对于目标检测，head采用与CenterNet一致的设计，对于实例分割，采用与CondInst相同的设计。
Data Association： 得到检测或者分割的结果后，作者设计了一个两轮的数据关联，将其与之前的轨迹相连。第一步，将新的结果与上一帧中以$(i,j)+{\mathbf{O}}_{i,j}^C$为圆心，$r$为半径的圆中未匹配的最相近的检测进行关联，作者设定的$r$为检测bbox的宽与高的均值。第二步，如果在第一步中，新结果并没有得到任何的匹配，便计算其与所有未匹配到的以及历史中的轨迹的embedding之间的余弦相似度，如果最高的相似度高于阈值（文中设为0.3），则将其与最高得分的轨迹相连。两轮匹配后，若新的检测或分割依旧没有匹配到轨迹，则将其设为新开始的轨迹。
TraDeS Loss： 整个结构的总损失函数为：$$L=L_{CVA}+L_{det}+L_{mask}$$$L_{CVA}$为CVA模块的损失，上文已经给出，$L_{det}$和$L_{mask}$推测应该就是CenterNet与CondInst中的损失相同，还没看过源码。之前有人在群里问作者，这里的检测和分割损失为何会加在一起，作者的回答是这里的检测和分割任务其实是并行的，两者之间并没有关系，都用到了前面跟踪的信息，只是为了实验一下文章的idea是否也适用于分割级别的多目标跟踪。如果单独做检测的话，分割分支和损失应该可以直接去除。

---

实验

这是评估各个模块效果的实验结果。两个模块效果均挺好，尤其是加了CVA模块后，IDS下降了很多。

CVA模块与数据关联的第二步均能有效的降低IDS。

是否采用偏移量的对比实验，可以看出，添加了偏移量之后的MOTA和IDF1指标均有所上升。

上表是TraDeS与近几年来的sota跟踪器之间的对比，在MOTA、IDF1指标上均比用来对比的其他检测器高，在FP、FN指标上也是，也有很低的IDS，最终生成的轨迹的完整度也超过其他的检测器。证明了采用跟踪的信息，是可以很大程度上提高检测的效果的，最明显的一点就是对当前帧被遮挡，而在之前帧未被遮挡的目标的检测更加容易更加准确。