[Paper Notes] DiffusionTrack: Diffusion Model For Multi-Object Tracking

Original link: https://arxiv.org/abs/2308.09905

1 Introduction

  Multi-target tracking is usually divided into two-stage tracking after detection (TBD) and one-stage joint detection and tracking (JDT). After TBD detects objects in a single frame, it uses a tracker to associate the same object across frames. Trackers used include motion-based tracking using Kalman filters, object association using re-identification techniques, graph-based trackers (modeling the association process as a minimum cost flow problem).
  The JDT method handles tracking and detection in a unified manner and can be divided into three categories: query-based trackers (using implicit unique queries to force queries to track the same object), offset-based trackers (using motion features to estimate motion offsets) ) and trajectory-based trackers (which use spatiotemporal information to deal with severe object occlusion problems).
  However, most TBD and JDT methods have the following problems:

1. Global or local inconsistency: TBD's separate training process for detection and tracking leads to global inconsistency, while JDT also places detection and tracking in different branches or modules, which does not completely resolve the inconsistency.
2. The balance between robustness and model complexity is suboptimal: TBD is simple, but detection fluctuations affect performance; JDT is robust, but it damages detection accuracy.
3. Inflexibility of different scenes within the same video: Videos are processed under uniform settings without adaptively applying strategies to different scenes.

  This article proposes DiffusionTrack, which establishes a new paradigm from noise to tracking. This method forms object associations directly from a set of random bounding box pairs within adjacent frames, as shown in the figure. The coordinates of bounding box pairs are refined so that they cover the same object in both frames, allowing implicit simultaneous detection and tracking under a unified model.

3. Method

3.1 Preparation knowledge

  The goal of multi-target tracking is to obtain a set of input target pairs $(X_t,B_t,C_t)$ , where$X_t$Yestt _Input image at time$t$$B_t$and $C_t$Yestt _The set of bounding boxes and class labels at time$t\; .$$B_t$的元素$B_t^i=(c_x^i,c_y^i,w_i,h_i)$ , part$i$ is the identification number of the object. When$X_t$middle iiWhen object$i$$B_t^i=\varnothing$ .
  Diffusion modelstypically utilize two Markov chains, a forward chain that adds noise to the image and a backward chain that extracts the image from the noise. Given data distribution$x_0\sim q(x_0)$ and define the variance sequence$b_1,b_2,\cdots ,b_T$, the forward process is defined as $$p(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-b_t}{x}_{t-1},b_{t}I)$$ Given$x_0$, by sampling the Gaussian vector $ϵ\sim \mathcal{N}(0,I)$ , xt x_tcan be obtained according to the following formula$x_t$Sample of: $$x_t=\sqrt{{\bar{a}}_t}{x}_0+(1-{\bar{a}}_t){ϵ}_t$$Where ${\bar{a}}_t={\prod }_{s=0}^t(1-b_s)$ . During training, the grid will be based on xt x_tat different times$x_t$Predict $x_0$. When extrapolating, from random noise $x_T$Start and iterate through the reverse process to obtain $x_0$。

3.2 DiffusionTrack

  The network structure is shown in the figure, including the feature extraction backbone and the data association denoising head (diffusion head). The former is run once to extract two adjacent frames in depth $(X_{t-1},X_t)$ , which uses features as conditions to gradually extract predicted associated bounding box pairs from noisy bounding box pairs. The data sample is a set of bounding box pairs$z_0=(B_{t-1},B_t)\in {\mathbb{R}}^{N\times 8}$。神经网络 f θ ( z s , s , X t − 1 , X t ) , s ∈ { 0 , 1 , ⋯   , T } f_\theta(z_s,s,X_{t-1},X_t),s\in\{0,1,\cdots,T\} fi(zs​,s,Xt−1​,Xt​),s∈{ 0,1,⋯,T } is trained as$z_s$Predict $z_0$, and based on two adjacent frames of images. Corresponding category label $(C_{t-1},C_t)$ and the associated confidence score$S$ is also predicted at the same time. If$X_{t-1}=X_t$, then the multi-target tracking task degenerates into a target detection task.
  The backbone uses YOLOX+FPN.
  The diffusion head uses proposal boxes to crop RoI features from the feature map and feeds them into different blocks to obtain bounding box regression values, classification results, and associated confidence scores. Additionally, a spatio-temporal fusion module (STF) and associated fractional head are added for each block of the diffusion head.
  The spatiotemporal fusion module enables pairs of bounding boxes to exchange temporal information to ensure the integrity of data association. Given the RoI feature $f_{roi}^{t-1},f_{roi}^t\in {\mathbb{R}}^{N\times R\times d}$ and self-attention output query$q_{pro}^{t-1},q_{pro}^t\in R^{N\times d}$ , use linear projection and batch matrix multiplication to obtain the object query$q^{t-1},q^t\in {\mathbb{R}}^{N\times d}$：
$$\begin{array}{rl}[P_1^i;P_2^i]& =Linear1(q_{pro}^i),P_1^i,P_2^i\in {\mathbb{R}}^{N\times Cd}\\ \mathbf{feat}& =[f_{roi}^i,f_{roi}^{2t-1-i}],feat\in R^{N\times 2R\times d}\\ feat& =\mathtt{Bmm}(\mathbf{feat},P_1^i.\mathtt{view}(N,d,C))\\ feat& =Bmm(feat,P_2^i.view(N,C,d))\\ q_i& =Linear2(\mathbf{feat}.\mathtt{flatten}(1)),q^i\in R^{N\times d}\end{array}$$Among them $[\cdot ,\cdot ]$ means splicing,$[\cdot ;\cdot ]$ represents division,i ∈ { t − 1 , t } i\in\{t-1,t\}i∈{ t−1,t } .
  The association score headerinputs the fused features of the bounding box pairs into the linear layer, which is used to obtain the confidence score of the data association. This score is used in the subsequent NMS post-processing link to determine whether the output bounding box pair belongs to the same object.

3.3 Model training and inference

  When training, randomly sample two frames from the video (the frame interval is 5), and fill the bounding box so that its number is $N_{train}$. Then, $a_t$Use monotonic decreasing cosine scheduling and add Gaussian noise to the ground-truth bounding boxes. Finally, a denoising process is used to obtain associated bounding box pairs from the noisy bounding boxes. In addition, a baseline scheme is designed that only adds noise to the current frame and denoises based on the bounding boxes of past frames to demonstrate the necessity of destroying both frames simultaneously.
  Loss function : Extended 3D GIoU using GIoU. For the matching set MM obtained using the Hungarian algorithmMatching pair (T d, T gt) (T_d,T_{gt})in$M$$(T_d,T_{gt})$ , record its classification score, bounding box and associated confidence score as$(C_d^{t-1},C_d^t),(B_d^{t-1},B_d^t)$ and$S_d$。损失函数如下：$$\begin{gathered} {\mathcal{L}}_{cls}(T_d,T_{gt})=\sum _{i=t-1}^t{\mathcal{L}}_{cls}(\sqrt{C_d^i\times S_d},C_{gt}^i) \\ {L}_{reg}(T_d,T_{gt})=\sum _{i=t-1}^t{L}_{reg}(B_d^i,B_{gt}^i) \\ {L}_{de}=\frac{1}{N_{pos}}\sum _{(T_d,T_{gt})\in M}{l}_1{L}_{cls}(T_d,T_{gt})+l_2{L}_{reg}(T_d,T_{gt})+l_3(1-{GIoU}_{3D}(T_d,T_{gt})) \end{gathered}$$ where$T_d$和$T_{gt}$Contains respectively the estimation results and real bounding boxes of the same target in two adjacent frames. $N_{pos}$is the number of prospect targets. $L_{cls}$is the focal loss, ${\mathcal{L}}_{reg}$for L1 loss.
  The figure below shows the inference process. Different from detection, since the estimation results of past frames are known, N test N_{test} can be generated from the bounding boxes of past frames.$N_{test}$An initial noisy bounding box. In the baseline scheme, existing bounding boxes are repeated instead of adding random bounding boxes, and noise is only added to the current frame. After obtaining the correlation results, IoU is used to measure the similarity and connect the object trajectories. To handle potential occlusions, a Kalman filter is used to re-associate missing objects.

4. Experiment

4.1 Settings

  Implementation details : using YOLOX detector + DiffusionTrack. The training is divided into two phases: first training with detection tasks, and then training with tracking tasks.

4.2 Features

  DiffusionTrack has dynamic bounding boxes, progressive refinement features, and robustness to detecting perturbations.
  Dynamic bounding boxes and progressive refinement : Once the model is trained, the number of bounding boxes and the number of sampling steps in DiffusionTrack can be modified during inference. Thus a balance of performance and accuracy can be achieved without retraining.
  Robustness to detecting disturbances : Almost all past methods are sensitive to detecting disturbances, which poses safety issues in autonomous driving. Experiments show that DiffusionTrack is almost unaffected by detection disturbances.

4.3 Ablation studies

  Based on Figure 3, this paper studies the impact of each part.
  Proportion of prior information : Since the tracking task has prior information about the object in the previous frame, the proportion of prior information can be controlled by copying the number of existing bounding boxes. Experiments show that the right ratio can improve performance.
  Bounding box filling strategy : Filling random bounding boxes following a Gaussian distribution outperforms random bounding boxes following a rest distribution, image-sized bounding boxes, and copying existing bounding boxes.
  Disturbance scheduling : In order to deal with complex situations, $a_t$. For example, when the object moves nonlinearly, $a_t$. Disturbance scheduling is available $t$ modeling, expressed as$t=\max (999,\min (0,1000f(x)))$ , where$x$ is the average percentage of objects spanning two frames,$f$ is the disturbance scheduling function. Experiments show that the optimal scheduling is logarithmic scheduling, that is,$f(x)=\frac{\log (x+1)}{\log 2}$.
  Efficiency comparison : More refinement steps improve performance but also reduce speed.

4.4 Comparison with SotA

  The method in this article can reach the SotA level of the one-stage method on the MOT17, MOT20 and Dancetrack data sets.
  The performance of the baseline scheme is similar on MOT17, but drops severely on the remaining datasets. This is because it only learns the given $t-In\; the\; case\; of\; 1$ time characteristics,$B_{t-1}$and $B_t$coordinate regression, but cannot handle nonlinear object motion. This article speculates that the diffusion process is a special data enhancement method that enables DiffusionTrack to distinguish different targets.