A single-stage multi-target tracking method based on attention mechanism

Abstract

A single-stage multi-object tracking method based on an attention mechanism is proposed. The method involves selecting a multi-object tracking dataset and preprocessing the images; building a network model, adding a designed discrimination enhancement module after the backbone network, followed by a prediction head for the detection branch and a designed identity preservation module; and finally, adding a prediction head for the re-ID branch after the identity preservation module. During training, a frame from the dataset video is randomly selected as input, and the coordinates of each pedestrian in the dataset annotations are used as samples for the detection branch, while the ID of each pedestrian is used as training samples for identity features. During testing, the current frame of the video is used as input to obtain the detection result and identity features of the current frame; and a matching process is then implemented. This invention improves the adaptability of object matching with scale variations.

Description

Technical Field

[0001] This invention belongs to the field of computer vision technology, specifically relating to a single-stage multi-target tracking method based on an attention mechanism. Background Technology

[0002] Multi-object tracking is one of the most fundamental and challenging tasks in computer vision. The goal of multi-object tracking is to track multiple objects of interest in a video by estimating their trajectories from temporal and spatial information. Successful solutions to this problem can immediately benefit many applications, such as autonomous driving, human-computer interaction, and intelligent surveillance video.

[0003] Existing methods primarily employ two main paradigms: detection-based tracking and joint detection tracking. Past research has largely followed detection-based tracking, dividing the multi-object tracking task into two models: a detector model and an identity embedding (re-ID) model. This method first uses detectors to obtain the bounding boxes of objects in each frame, then associates the obtained bounding boxes with existing trajectories by matching the predicted locations with the extracted identity features of each bounding box. While achieving good results, these methods face significant computational costs and still adhere to the detection-per-detection tracking paradigm, limiting real-time tracking. For example, SORT uses a Kalman filter to track bounding boxes and associates each bounding box with a Hungarian algorithm; ByteTrack tracks objects by associating each detection box with a single detection box, rather than just associating high-resolution detection boxes. In recent years, tracking methods that learn detection and association within a single network have demonstrated their ability to achieve a good balance between inference speed and tracking accuracy, gradually becoming the mainstream approach.

[0004] To achieve high levels of detection and tracking accuracy, MeMOT proposes a long-range spatiotemporal memory algorithm to improve the ability to connect objects over long time spans; FairMOT proposes a method that treats detection and association tasks equally, constructing two parallel branches for object detection and appearance feature extraction, respectively. However, since joint detection and tracking methods share the same backbone features between detection and matching tasks, they may lead to conflicting information requirements from feature representations. Summary of the Invention

[0005] To overcome the shortcomings of existing multi-target tracking methods that involve conflicting requirements for shared features between detection and association tasks, this invention provides a single-stage multi-target tracking method based on an attention mechanism. First, a discrimination enhancement module is designed, which uses a self-attention mechanism to acquire contextual information to obtain better feature representation. Then, an identity preservation module is designed, which maintains the semantic alignment of identity features through channel recalibration design, thereby improving the adaptability of object matching with scale variations.

[0006] The technical solution adopted by this invention to solve its technical problem is:

[0007] A single-stage multi-target tracking method based on an attention mechanism includes the following steps:

[0008] (1) Select a multi-target tracking dataset and perform image preprocessing;

[0009] (2) Build a network model, add the designed discrimination enhancement module after the backbone network, add the prediction head of the detection branch and the designed identity preservation module after the discrimination enhancement module, and finally add the prediction head of the re-ID branch after the identity preservation module.

[0010] (3) During training, a frame of the video in the dataset is randomly selected as the input, the coordinates of each pedestrian in the dataset are used as the sample for the detection branch, and the ID of each pedestrian is used as the sample for identity feature training.

[0011] (4) During testing, the current frame of the video is used as input. After feature extraction, the enhanced features are first obtained through the discrimination enhancement module. Then, the enhanced features are input into the prediction head of the gold detection and the identity preservation module respectively. Finally, the detection result and identity features of the current frame are obtained.

[0012] (5) The matching process consists of three steps. First, the current frame is connected to the previous frame to match the appearance and distance of the newly detected target. Second, the current frame is connected to the previous frame again to calculate the IOU distance that was not matched in the previous step. Third, the state is updated for matching in the next frame.

[0013] Furthermore, in step (1), the multi-target tracking datasets used are MOT17 and MOT20, and the preprocessing includes HSV enhancement, rotation, scaling, translation, and shearing.

[0014] Furthermore, in step (2), the framework of the network model can be considered to consist of three parts: a feature extraction backbone network, used to extract features of the current frame image; a discrimination enhancement module, used to enhance the discrimination features of the current frame object, which consists of a spatial attention enhancement and a channel attention enhancement; and an identity preservation module, which consists of channel attention, used to recalibrate channels to achieve semantic alignment of information for specific objects and preserve identity features.

[0015] Furthermore, in step (3), during training, the detection branch consists of three parts: a heat map of the object's center point, the distance of the object's center from the top, bottom, left, and right sides, and the offset of a center point. The matching branch only trains identity features, and both branches are trained simultaneously.

[0016] In step (4), during testing, the current image is first processed by a backbone network for feature extraction. Then, the detection result is obtained through a discrimination enhancement module, and the identity features are obtained through an identity preservation module to distinguish different targets and complete the classification task. Finally, post-processing is used to obtain the bounding boxes of all objects in the current frame and the identity features used for matching.

[0017] In step (5), preliminary results are first obtained using Kalman filtering and Re-ID features, and then the Hungarian algorithm is used for matching. Then, for unmatched trajectories and detections, a second matching is performed by setting an IOU threshold. Finally, the unmatched detections are initialized as new trajectories, and the unmatched trajectories are saved for 30 frames for future recurrence.

[0018] The beneficial effects of this invention are mainly reflected in the following aspects: based on the general joint detection and tracking framework, attention mechanism is used to improve the feature representation of the entire tracking process, and instance-specific information is perceived for re-identification.

[0019] Since effective feature representation with sufficient spatial information is key to multi-target detection with only one frame as input, a discriminative enhancement module is designed to promote the feature representation of the normal backbone by utilizing self-attention mechanism and channel attention, and to enhance discriminability by utilizing spatial and contextual information, thereby balancing accuracy and efficiency.

[0020] The contextual information enhanced by the designed discrimination enhancement module can mostly facilitate detection. However, research shows that low-level features contain more object-specific information for re-ID learning. Based on this, an identity preservation module is proposed to recalibrate channels. This module is developed using a channel attention mechanism to achieve information specific to the object and preserve the semantic alignment of identity features.

[0021] Experiments on the MOT-17 and MOT-20 datasets show that this method has certain advantages over other single-stage multi-target tracking algorithms. Attached Figure Description

[0022] Figure 1 This is a flowchart of the overall framework.

[0023] Figure 2 To identify the enhancement module diagram.

[0024] Figure 3 Module diagram for identity preservation. Detailed Implementation

[0025] The present invention will now be further described with reference to the accompanying drawings.

[0026] Reference Figures 1-3 A single-stage multi-target tracking method based on an attention mechanism includes the following steps:

[0027] (1) Select a general multi-target tracking dataset and perform image preprocessing.

[0028] Input the current frame pedestrian image to be predicted, and apply different transformations to the entire image, including HSV enhancement, rotation, scaling, translation and cropping for preprocessing;

[0029] (2) Build a network model, add the designed discrimination enhancement module after the backbone network, add the prediction head of the detection branch and the designed identity preservation module after the discrimination enhancement module, and finally add the prediction head of the re-ID branch after the identity preservation module.

[0030] ResNet-34 is used as the backbone to achieve a good balance between accuracy and speed. An enhanced version of Deep Convergence (DLA) is applied to the backbone to fuse multi-layer features. The input image size is 608×1088×3, and the extracted feature size is 152×272×3. First, the features are input into the discrimination enhancement module, which consists of spatial attention enhancement and channel attention enhancement. To enable the network to focus on foreground information and suppress background noise, spatial attention enhancement is achieved through a self-attention mechanism. This mechanism encodes remote spatial information with local features and background information. The self-attention mechanism process can be represented by the following formula:

[0031]

[0032] First, spatial attention is enhanced, F∈R C×H×W F represents the features extracted by the backbone network, and Q, K, V∈R are obtained through three learned weight matrices. C×H×W Then by Spatial attention map is calculated Where N = H × W, it represents the characteristic influence of the i-th position on the j-th position. The correlation between these two positions is... The values ​​are positively correlated; then The transpose of F is used to calculate a new feature map, which is then adjusted to R. C×H×W The final output F is obtained. SA ∈R N ×N It is obtained by multiplying a learnable parameter α by the feature F:

[0033]

[0034] in In F SA For each position, the final feature is calculated as a weighted sum of all positions with the initially extracted features; therefore, the enhanced features can adaptively obtain a global context view and aggregation from the spatial attention graph, thereby improving consistency within the object class.

[0035] Next, in the channel attention enhancement section, since each channel of a feature can be viewed as a response map of a specific object, channel attention can be used to enhance the feature representation of specific object information. For example... Figure 2 As shown, the features F∈R extracted from the backbone network are first... C×H×W Remodeling into R C×N Where N = H × W. Similarly, Q, K, V ∈ R are learned through three learnable weight matrices. C×H×W Then, the transpose of K and Q are computed and passed through a softmax layer to obtain the channel attention map. It represents the influence of the i-th channel on the j-th channel in terms of features; next, for Multiply the transpose of F with F, and adjust the dimension of the result to R. C×H×W Finally, multiply by a learnable parameter β and sum element-wise with the feature map F to obtain the final output F. CA ∈R C×H×W :

[0036]

[0037] in For F CA Each channel is a weighted sum of features across all channels and the original features, which can improve the discriminability of features based on remote context information;

[0038] Finally, the two enhanced features are fused as the final output of the discrimination enhancement module to predict detection results. The input features are then fed into the identity preservation module. The contextual information enhanced by the proposed discrimination enhancement module largely facilitates detection; however, research shows that low-level features contain more object-specific information for re-ID learning. Therefore, a channel attention mechanism is used to develop this module, such as... Figure 3 As shown, it achieves this through three steps: First, the F output by the feature discrimination module is... DEM ∈R C×H×W Feature division involves multi-scale features with different kernel sizes on each channel. This operation yields richer positional information of the input tensor and enables parallel processing across multiple scales. Next, the SEWeight module is used to further extract channel attention from the input feature maps at multiple scales, generating channel-level attention weight vectors, represented as follows:

[0039]

[0040] Z i ∈R C′×1×1 These are attention weights, where each feature map has a different scale. With a common channel dimension C' = C / S, this operation allows for greater focus on shallow characteristics suitable for re-ID tasks, by adjusting the weights with the corresponding scale. By multiplying features in the local and global channels, enhanced features are obtained through better information exchange between local and global channels.

[0041]

[0042] Finally, these features are combined to generate a refined feature map for predicting identity features:

[0043] F IPM =Cat([Y0,Y1,...,Y)) (S-1) ].)

[0044] The identity preservation module preserves the semantic alignment of identity features. It complements the discrimination enhancement module and improves the resilience of objects at different scales during the association process, thereby improving the adaptability of matching objects with different scale variability.

[0045] (3) During training, a frame of an image from the video in the dataset is randomly selected as input. The coordinates of each pedestrian in the dataset annotation are used as samples for the detection branch, and the ID of each pedestrian is used as samples for identity feature training.

[0046] We use uncertainty loss to automatically balance detection and re-ID tasks:

[0047]

[0048] w1 and w2 are learnable parameters used to balance the two tasks. The detection branch is built on top of CenterNet and learns re-ID features through a classification task. All object instances with the same identifier in the training set are considered to be in the same class. It is trained using a cross-entropy loss.

[0049] The experimental parameters were set as follows: a variant of DLA-34 was used as the default backbone network; model parameters pre-trained on the COCO dataset were used to initialize our model; the model was trained for 20 epochs using the Adam optimizer; the initial learning rate was 1r-4; the learning rate was reduced to 1r-5 when training reached the 10th epoch; the batch size was set to 12; and an RTX3090 was used.

[0050] (4) During testing, the current frame of the video is used as input. After feature extraction, the enhanced features are first obtained through the discrimination enhancement module. Then, the enhanced features are input into the detection prediction head and the identity preservation module respectively. Finally, the detection result and identity features of the current frame are obtained.

[0051] The network takes 1088×608 frames as input, the same as the previous work JDE; on the predicted heatmap, non-maximum suppression (NMS) is performed based on the heatmap score to extract peak keypoints. The NMS is implemented by a simple 3×3 max pooling operation, which preserves the positions of keypoints whose heatmap scores are greater than a threshold; then the corresponding bounding boxes are calculated based on the estimated offset and box size, and the identity embedding at the estimated object center is extracted.

[0052] (5) Matching is divided into three steps. First, the current frame is concatenated with the previous frame to match the appearance and distance of the newly detected target. Second, the current frame is concatenated with the previous frame again to calculate the IOU distance of the unmatched target in the previous step. Third, the state is updated for matching in the next frame.

[0053] Following the MOTDT approach, a hierarchical online data association method is used. First, some trajectories are initialized based on the bounding boxes detected in the first frame. Then, in subsequent frames, a two-stage matching strategy is used to connect the detected bounding boxes to the existing trajectories. In the first stage, Kalman filtering and re-ID features are used to obtain initial tracking results. Specifically, a Kalman filter is used to predict the trajectory position in the next frame, and then the Mahalanobis distance D between the detected and predicted bounding boxes is calculated. m By fusing the Mahalanobis distance with the cosine distance calculated based on re-ID features, a similarity matrix is ​​obtained:

[0054] D=λD r +(1-λ)D m

[0055] Where λ is a weighting parameter, set to 0.98; referring to JDE, if the Mahalanobis distance is greater than a threshold, it is set to infinity to avoid obtaining trajectories with large motions. The Hungarian algorithm is used for matching, and its matching threshold is τ1 = 0.4 to complete the first stage of matching.

[0056] Then, for unmatched detections and trajectories, they are matched based on the degree of overlap between their detection boxes, with a matching threshold τ2 = 0.5. The appearance features of the trajectories are updated at each time step to handle appearance changes. Finally, the unmatched detections are initialized as new trajectories, and the unmatched trajectories are saved for 30 frames in case they reappear in the future.

[0057] Table 1 compares this application with other methods and trackers on the MOT17 and MOT20 datasets. MOTA reflects the average tracking performance of the tracker, IDF1 reflects the correlation performance, and IDs represent the number of identity changes.

[0058]

[0059] Table 1.

[0060] The scheme in this embodiment is designed with a discrimination enhancement module and an identity preservation module. The discrimination enhancement module is designed through spatial attention and channel attention, and uses a self-attention mechanism to capture contextual information to obtain better feature representation. The identity preservation module is designed through channel attention, which maintains the semantic alignment of identity features and improves the adaptability of target matching under scale changes.

[0061] The embodiments described in this specification are merely examples of implementations of the inventive concept and are for illustrative purposes only. The scope of protection of this invention should not be considered limited to the specific forms described in these embodiments; rather, it extends to equivalent technical means conceived by those skilled in the art based on the inventive concept.

Claims

1. A single-stage multi-object tracking method based on an attention mechanism, characterized in that, The method comprises the following steps: (1) selecting a multi-target tracking data set and performing preprocessing on images; (2) building a network model, adding a designed discrimination enhancement module after a backbone network, adding a prediction head of a detection branch and a designed identity reservation module after the discrimination enhancement module, and finally adding a prediction head of an re-ID branch after the identity reservation module; (3) during training, randomly selecting a picture of a frame in a video in a data set as input, taking the coordinates of each pedestrian in data set annotation as a sample of the detection branch, and taking the id identification of each pedestrian as a sample for identity feature training; (4) during testing, taking the current frame of a video as input, obtaining enhanced features through the discrimination enhancement module after feature extraction, inputting the enhanced features into the prediction head of the detection and the identity reservation module, and finally obtaining the detection result and the identity feature of the current frame; (5) the matching process is divided into three steps, first, connecting the current frame and the previous frame to match the appearance and distance of the new detection target; second, connecting the current frame and the previous frame again to calculate the IOU distance of the unmatching in the previous step; third, updating the state for the matching of the next frame; In the step (2), the framework of the network model can be considered to have three parts, which are a feature extraction backbone network for extracting features of the current frame picture, a discrimination enhancement module for enhancing the discrimination features of the current frame object, which has a spatial attention enhancement and a channel attention enhancement, and an identity reservation module composed of channel attention, which is used for recalibrating channels to achieve information for specific objects and to reserve semantic alignment of identity features; ResNet-34 is used as the backbone, and the enhanced deep aggregation DLA is applied to the backbone to fuse multi-layer features, the input image size is represented as 608x1088x3, and the extracted feature size is 152x272x3; first, the features are input into the discrimination enhancement module, which is composed of a spatial attention enhancement and a channel attention enhancement, and the spatial attention enhancement is realized through a self-attention mechanism to encode the remote spatial information and background information with local features; the self-attention mechanism process is shown in the following formula: First is the spatial attention enhancement, F∈R C×H×W The features extracted by the backbone network, F is obtained by three learning weight matrices Q, K, V∈R C×H×W , and then The spatial attention map is calculated Where N=H×W, which represents the feature influence of the i-th position on the j-th position, the correlation of the two positions is positively correlated with the value of ; Then The transpose of and F are calculated to obtain a new feature map, and then the dimension is adjusted to R C×H×W , to obtain the final output F SA ∈R N×N It is calculated by multiplying a learnable parameter α by the feature F: wherein In F SA At each location, the final feature is computed as a weighted sum of all locations of the initially extracted features; thus, the enhanced features can adaptively obtain a global contextual view and aggregation from the spatial attention map, thereby improving the intra-class consistency of the object. Next, the channel attention enhancement part, first extract the features F ∈ R C×H×W Reshaped to R C×N , where N = H × W, also through three learnable weight matrix to learn Q, K, V ∈ R C×H×W Then the transpose of K and Q are calculated through a softmax layer to obtain the channel attention map It represents the influence of the i-th channel on the j-th channel in the feature; next, the transpose matrix of and F are multiplied, and the dimension of the result is adjusted to R C×H×W ; finally, multiply a learnable parameter β, and sum it with the feature map F in an element-level manner to obtain the final output F CA ∈ R C ×H×W : wherein For F CA For each channel, it is the weighted sum of the features and the original features over all channels, which can improve the discriminability of the features with respect to remote context information; Finally, the two parts of enhanced features are fused as the final output of the discrimination enhancement module, which is used to predict the detection result, and the input features are sent into the identity preservation module; first, the F DEM ∈R C×H×W The features are divided into multi-scale features of different kernel sizes on each channel; next, the SEWeight module is used to further extract the channel attention of the input feature map of multiple scales to generate a channel-level attention weight vector, which is represented as: wherein Z i ∈ R C′×1×1 are attention weights, each feature map has different scale feature maps, has a common channel dimension C' = C / S. Through this operation, more attention can be paid to the shallow features suitable for the re-ID task, and enhanced features can be obtained through better information interaction between local and global channels by multiplying the weights with the corresponding scale features in Z Finally, they are fused to generate a refined feature map for predicting identity features: F IPM = Cat([Y0, Y1,..., Y (S-1) ]) The identity reservation module preserves the semantic alignment of identity features, which is a supplement to the discrimination enhancement module, improves the flexibility of objects of different scales in the association process, and thus improves the adaptability of object matching with different scale change capabilities.

2. The single-stage multi-object tracking method based on attention mechanism according to claim 1, wherein, In the step (1), the multi-target tracking data set used is MOT17 and MOT20, and the preprocessing includes HSV enhancement, rotation, scaling, translation and shearing.

3. The single-stage multi-object tracking method based on attention mechanism according to claim 1 or 2, characterized in that, In the step (3), during training, the detection branch is composed of three parts: a heat map of an object center point, a distance of the object center point from up, down, left and right, and an offset of a center point, and the matching branch only trains identity features, and the two branches are trained simultaneously.

4. The single-stage multi-object tracking method based on attention mechanism according to claim 1 or 2, characterized in that, In the step (4), during testing, the current picture is first subjected to feature extraction through the backbone network, then a detection result is obtained through the discrimination enhancement module, and then an identity feature is obtained through the identity reservation module, which is used to distinguish different targets, complete the classification task, and finally obtain the coordinate frame of all objects in the current frame and the identity feature for matching through post-processing.

5. The single-stage multi-object tracking method based on attention mechanism according to claim 1 or 2, characterized in that, In the step (5), first, the Kalman filter and the Re-ID feature are used to obtain a preliminary result, and the Hungarian algorithm is used for matching; then, for the trajectories and detections that cannot be matched, a second matching is performed by setting an IOU threshold; finally, the unmatched detections are initialized as new trajectories, and the unmatched trajectories are saved for 30 frames for future reoccurrence.

Images