Multi-person head posture estimation method and device based on multi-task loss balance

Abstract

The invention discloses a multi-person head posture estimation method and device based on multi-task loss balance. The method comprises the following steps: constructing a multi-person head posture estimation model based on a Mamba-YOLO network; training a multi-person head posture estimation model, and controlling the training process by using a loss function based on a dynamic weight in the training process, the loss function based on the dynamic weight is obtained by weighting the loss function of the head prediction probability learning task, the loss function of the head position detection learning task and the loss function of the head posture estimation learning task by using the dynamic weight, and the dynamic weight of each branch is determined according to the loss value of the corresponding task; and obtaining a to-be-detected picture, and inputting the to-be-detected picture into the trained multi-person head posture estimation model to obtain a detection result. According to the method, end-to-end multi-person head posture estimation based on multi-task loss balance can be realized, and the detection real-time performance is kept while the estimation precision is improved.

Description

Technical Field

[0001] The present invention relates to the technical field of head posture estimation, and in particular to a method and device for multi-person head posture estimation based on multi-task loss balance. Background Art

[0002] Head posture is an important factor in judging human behavior and is also very helpful in understanding human behavior. Head posture estimation refers to the computer analyzing and predicting the input image or video sequence to determine the position and posture parameters of a person's head in three-dimensional space. Head posture estimation can be applied in various application fields such as human-computer interaction, biometrics, virtual reality, and fatigue detection.

[0003] At present, head posture estimation can be mainly divided into two directions: single-person head posture estimation and multi-person head posture estimation. Most of the research focuses on single-person head posture estimation. Although the single-person head posture estimation algorithm has high accuracy, its detection efficiency is low in the multi-person head posture estimation scenario, which is not conducive to engineering deployment and application. For multi-person head posture estimation, the existing technology usually adopts a one-stage multi-person head posture estimation algorithm based on deep learning models such as CNN (Convolutional Neural Network), which can improve the efficiency of the multi-person head posture estimation algorithm. However, the backbone feature extraction ability of the CNN model is weak and cannot fully extract the detailed features of the head area. In addition, since multiple learning tasks such as category prediction, positioning and posture estimation need to be performed during the detection process, each task algorithm usually adopts a fixed weight during the learning process, which will lead to an imbalance between the loss values ​​of each learning task, which will affect the accuracy and reliability of the model head posture estimation. Summary of the invention

[0004] The technical problem to be solved by the present invention is: in response to the technical problems existing in the prior art, the present invention provides a method and device for multi-person head posture estimation based on multi-task loss balance, which is simple to implement, low in cost, high in detection accuracy and reliability, and strong in real time. It can realize end-to-end multi-person head posture estimation based on multi-task loss balance, improve the estimation accuracy while maintaining the real-time detection.

[0005] In order to solve the above technical problems, the technical solution proposed by the present invention is: A method for multi-person head posture estimation based on multi-task loss balance, comprising the following steps: A multi-person head posture estimation model is constructed using a YOLO network based on the Mamba model. The multi-person head posture estimation model includes a backbone network, a PAFPN (Pyramid Attention Feature Pyramid Network) module, and a detection head. The backbone network is an ODMamba network structure for processing input images and extracting preliminary features. The PAFPN module is used to extract and fuse features of different scales. The detection head is used to achieve the tasks of category prediction, head position detection, and head posture estimation. Using an image training set containing multiple head objects to train the multiple head posture estimation model, during the training process, a loss function based on dynamic weights is used to control the training process, wherein the loss function based on dynamic weights is obtained by weighting the loss function of the head object prediction classification task, the loss function of the head position detection task, and the loss function of the head posture estimation task using dynamic weights, and the dynamic weight of each task is determined according to the loss value of the corresponding task to balance the loss of each task; Obtain a picture to be detected containing multiple heads, input it into the trained multi-person head posture estimation model, and obtain detection results of multiple head objects.

[0006] Furthermore, the ODMamba network structure includes an image segmentation module, an ODSSBlock module, a visual cue merging module and an SPPF module. The image segmentation module is used to perform preliminary processing on the input image through multiple convolutional layers, and the ODSSBlock module is used to further perform feature extraction; the visual cue merging module is used to further fuse and optimize the features output by the ODSSBlock module, and the SPPF module is used to aggregate features at multiple scales; the core module in the PAFPN module is the ODSSBlock module; ACConv convolution is used to perform convolution operations at the input end of the ODSSBlock module, and deep separable DW-ACConv convolution is used to perform deep separation convolution operations in the LS Block module, SS2D module and RG Block module in the ODSSBlock module, the SS2D module is used to scan and merge image features to extract multi-directional global information, the LS Block module is used to extract local spatial information, and the RG Block module is used to capture local dependencies by combining the gating mechanism and residual connection.

[0007] Furthermore, the calculation expression of the loss function based on dynamic weight is:

[0008] in, represents the loss function based on dynamic weights, Indicates the batch size of the multi-person head pose estimation model during training. , , Respectively , and The dynamic weights of each dynamic weight change with the change rate of the loss value in the number of training rounds. represents the loss function of the head position detection task, represents the loss function of the category prediction task, Represents the loss function for the head pose estimation task.

[0009] Furthermore, the dynamic weight of each branch is calculated according to the following expression:

[0010]

[0011] in, represents the dynamic weight of the pth task, p=1, 2, 3, corresponding to the category prediction task, head position detection task, and head posture estimation task, respectively, t represents the number of training times, K represents the total number of tasks, represents the loss value of the pth task, T Indicates the preset temperature used to control the softness of task weighting.

[0012] Furthermore, the loss function of the head position detection task is constructed based on the focus regression box loss of gamma transform, and the calculation expression is:

[0013]

[0014]

[0015]

[0016] Where N represents the number of head detections in the image to be detected. , The i-th head prediction detection box and the true label are respectively, represents the Siou-based regression loss, Indicates the preset threshold value with a value range of (0,1). Indicates the loss value of the center coordinates and frame length of the head detection frame. represents the focus regression box loss based on gamma transform, It represents the intersection-over-union ratio between the head region detection frame and the real head frame. It represents the intersection-over-union ratio between the reconstructed head region detection frame and the real head frame. represents the gamma shift adjustment factor, represents the gamma shift adjustment factor.

[0017] Furthermore, the loss function of the head object prediction classification task is calculated according to the following formula:

[0018] in, represents the probability of the head predicted by the i-th head detection box, represents the real label of the head of the i-th head detection frame, Indicates the overlap between the i-th head detection frame and the true head frame.

[0019] Furthermore, the loss function of the head posture estimation task is constructed based on the tangent function, and the calculation expression is:

[0020]

[0021]

[0022] Among them, i represents the i-th dimension of the head posture matrix, represents the loss weight on the i-th dimension of the head posture matrix, represents the head pose estimation matrix, represents the true head pose matrix, , express , The elements in .

[0023] Furthermore, it also includes obtaining the head posture estimation matrix The steps include: Use the head pose estimation 6D decoupled detection head set in the multi-person head pose estimation model to detect the head pose 6D rotation representation parameters ; The two three-dimensional vectors corresponding to the 6D rotation parameters of each head object's head posture are , Convert to get the head pose estimation matrix The elements in: ; ; ; in, They are u exist x ,y , z The coordinate values ​​of the axis, They are v exist x , y , z The coordinate values ​​of the axis, , , Respectively represent the 6D rotation parameter matrix Transformed into three three-dimensional unit vectors in a 3×3 rotation matrix, , , , ~ Respectively , , Each element in Representation and A vector that is perpendicular in three-dimensional space.

[0024] An electronic device comprises a processor and a memory, wherein the memory is used to store a computer program, and the processor is used to execute the computer program to perform the above method.

[0025] A computer-readable storage medium storing a computer program, wherein the computer program implements the above method when executed by a processor.

[0026] Compared with the prior art, the advantages of the present invention are: 1. The present invention adopts the YOLO architecture based on the Mamba model as the basic model for head posture estimation to construct a multi-person head posture estimation model, which can expand the contextual semantic connection of each pixel and improve the feature extraction capability of the head posture estimation model for the head area. At the same time, a dynamic weighted strategy is used to determine the loss value to balance the loss values ​​between tasks, thereby reducing the impact of the imbalance between the loss values ​​of each task in the head posture estimation model on the model accuracy, and realizing end-to-end multi-person head posture estimation based on multi-task loss balance. The accuracy of the end-to-end multi-person head posture estimation algorithm can be improved while maintaining its detection real-time performance, which is convenient for practical engineering deployment and application.

[0027] 2. The loss function of the head position detection task of the present invention further adopts a focus regression frame loss based on gamma transform, so that it can focus on regression samples of different difficulties, enhance the model's attention to the difficult sample regression frame, and improve the regression accuracy of the head detection frame, thereby effectively improving the performance of the detector in different detection tasks.

[0028] 3. The loss function of the head posture estimation task of the present invention is further constructed based on the tangent function, which can effectively update and optimize the head posture estimation parameters, so that the model can pay more attention to the learning of difficult-dimensional vectors and further improve the accuracy of head posture estimation. BRIEF DESCRIPTION OF THE DRAWINGS

[0029] Figure 1 It is a schematic diagram of the implementation flow of the multi-person head posture estimation method based on multi-task loss balance in this embodiment.

[0030] Figure 2 Schematic diagram of the structure principle of a multi-person head posture estimation model constructed using a YOLO network based on the Mamba model in this embodiment.

[0031] Figure 3 It is a schematic diagram of the structural principle of the ODSSBlock module in this embodiment.

[0032] Figure 4 Schematic diagram of the structure principle of the LS Block module in this embodiment.

[0033] Figure 5 It is a schematic diagram of the structural principle of the SS2D module in this embodiment.

[0034] Figure 6 FIG. 4 is a schematic diagram of the structural principle of the RG Block module in this embodiment.

[0035] Figure 7 It is a schematic diagram of the structural principle of ACConv training in this embodiment.

[0036] Figure 8 It is a schematic diagram of the structural principle of ACConv reasoning in this embodiment.

[0037] Fig. 9 It is a schematic diagram of the ACConv inference convolution re-parameterization principle in this embodiment.

[0038] Fig.10 Schematic diagram of the structure principle of the Scan module in this embodiment.

[0039] Fig.11 It is a schematic diagram of the tangent function curve.

[0040] Fig.12 It is a schematic diagram of the implementation flow of the post-processing using the non-maximum suppression algorithm that is independent of confidence and Iou and is used in this embodiment. DETAILED DESCRIPTION

[0041] The present invention is further described below in conjunction with the accompanying drawings and specific preferred embodiments, but the protection scope of the present invention is not limited thereby.

[0042] For ease of understanding, the relevant technical background of the present invention is first introduced by way of example.

[0043] Taking the multi-person head posture estimation algorithm based on Yolov5 in the prior art as an example, this type of method adds an Euler angle detection head branch on the basis of YOLOv5, and uses MSE (mean square error) as the loss function to optimize the parameters of the head posture prediction branch, so as to achieve end-to-end multi-person head posture recognition model training and recognition. It can simplify the head posture recognition algorithm model training and recognition process, improve the real-time performance on various platforms, and also prove that embedding head position detection and head posture estimation into a feature representation layer is conducive to the network learning the intrinsic connection between the two. However, the algorithm model will have the following problems: 1. The backbone feature extraction capability of the model is weak In the downstream tasks of target detection, CNN and Transformer structures are mainly used. CNN and its improvements can provide fast execution speed while ensuring accuracy, but the convolution operation of CNN can only perceive the characteristics of local areas at each layer, and it is difficult to capture long-distance dependencies. For multi-task models such as multi-person head pose estimation, this will make it difficult for the target pixel to establish a contextual connection with the surrounding larger range of pixels, which is not conducive to the model to extract more effective features from the head area, thereby affecting the end-to-end multi-person head pose estimation model performance. Transformer performs well in global modeling, can effectively capture long-distance dependencies, and significantly expand the receptive field of the model. For example, the DETR series relies on the powerful global modeling ability of the self-attention mechanism to solve the problem of the small receptive field of CNN and improve the model's feature extraction ability, but the square complexity of its self-attention mechanism will increase the computational burden of the model, which is not conducive to the real-time processing and engineering deployment of the multi-person head pose estimation model.

[0044] 2. Imbalance between loss values ​​of multiple learning tasks End-to-end multi-person head pose estimation algorithms usually involve multiple learning tasks such as classification, positioning, and pose estimation. Traditional end-to-end multi-person head pose estimation algorithms usually directly sum the loss values ​​of each learning task with fixed weights as the final loss value, ignoring the imbalance problem between the loss values ​​of each learning task, which in turn affects the model accuracy.

[0045] In view of the above problems of the one-stage multi-person head pose estimation method in the prior art, the present invention adopts the YOLO architecture based on the Mamba model as the basic model for head pose estimation to construct a multi-person head pose estimation model, which can expand the contextual semantic connection of each pixel, improve the feature extraction capability of the head pose estimation model for the head area, and solve the limitations of models such as traditional CNN (local receptive field limitation in capturing long-distance information) or Transformer (excellent performance in global modeling, but high computational complexity). At the same time, a dynamic weighted strategy is adopted to determine the loss value to balance the loss values ​​between tasks, reduce the impact of the imbalance between the loss values ​​of each task in the head pose estimation model on the model accuracy, and realize end-to-end multi-person head pose estimation based on multi-task loss balance. It can improve the accuracy of the end-to-end multi-person head pose estimation algorithm while maintaining its detection real-time performance, which is convenient for practical engineering deployment and application.

[0046] like Figure 1 As shown, the steps of the multi-person head posture estimation method based on multi-task loss balance in this embodiment include: Step S01. Use the YOLO network based on the Mamba model to build a multi-person head posture estimation model. The multi-person head posture estimation model includes a backbone network, a PAFPN module and a detection head. The backbone network is an ODMamba network structure for processing the input image and extracting preliminary features. The PAFPN module is used to extract and fuse features of different scales. The detection head is used to realize category prediction, head position detection and head posture parameter estimation.

[0047] The Mamba model is a method based on the Selective State Space Model (SSM). By dynamically adjusting the model behavior based on the context of the input content, it can effectively capture long-distance dependencies and effectively alleviate the problems caused by quadratic complexity, so that the multi-person head posture estimation model has both local feature extraction capabilities and global modeling capabilities while maintaining a certain degree of real-time performance. This embodiment adopts the YOLO network based on the Mamba model to form a new backbone network Mamba-YOLO that combines the advantages of SSM and CNN, in which the feature extraction backbone of the network adopts the Mamba structure as the basic network of the end-to-end multi-person head posture estimation model, so that it has a strong feature extraction capability for the head area while maintaining a certain degree of reasoning real-time performance, thereby facilitating engineering deployment.

[0048] like Figure 2As shown, in this embodiment, the multi-person head posture estimation model constructed using the YOLO network based on the Mamba model includes a backbone network, a PAFPN module, and a detection head. The backbone network (feature extraction network part) of Mamba-YOLO is an ODMamba network structure (ODMamba Backbone), which is used to process the input image and extract preliminary features. It is the basic part of feature extraction and is used to generate sufficient information for further processing by subsequent modules. The PAFPN module is a feature pyramid network used to fuse and refine features of different scales. It combines upsampling and downsampling with additional convolutional layers to enhance the model's detection capabilities for objects of various scales, especially to improve the detection capabilities in multi-scale target detection tasks. The model sets up three different detection layers according to the different sizes of feature maps, which are responsible for small, medium and large-sized head object target detection respectively (each grid includes head prediction probability Cls, position detection box Box and head posture parameter Head Pos, a total of 11 channels) and generates dense head object output. Each detection head uses a decoupled detection head to use sliding convolution in three ways to achieve three tasks: category prediction, head position detection and head posture parameter estimation. The use of decoupled detection head method can speed up the model convergence speed and improve the model detection accuracy.

[0049] like Figure 2As shown in the figure, the ODMamba network structure consists of four modules: image segmentation, ODSSBlock, visual clue merging and SPPF. The image segmentation module is used to perform preliminary processing on the input image through multiple convolutional layers. Specifically, the image segmentation module contains multiple convolutional layers, which are used to perform preliminary processing on the input image such as scaling and basic feature extraction to reduce the dimension of the input data, thereby reducing the burden for subsequent deeper feature extraction. ODSSBlock is the core module in the ODMamba network structure, which is used to process the input image and extract preliminary features. The SPPF module is the spatial pyramid pooling fast layer for feature extraction. The visual clue merging module is used to further fuse and optimize the features output by the ODSSBlock module to ensure that the features extracted from different levels can be effectively combined to enhance the detection ability of small and large targets. The SPPF module is used to aggregate features at multiple scales. By aggregating features at multiple scales, the model can better understand and process objects of different scales. The core module in the PAFPN module is ODSSBlock, which is used to abstract image features so that the module can maintain a good balance between accurate feature expression and efficient extraction quality inspection. The PAFPN module also includes the use of dilated convolution to replace traditional convolution to further abstract image features and establish semantic connections between feature pixels and surrounding pixels in a larger range, thereby improving the contextual semantic connections of the output features and further improving the feature expression capabilities of the output features.

[0050] The structure of the ODSSBlock module in this embodiment is as follows Figure 3As shown in the figure, it includes LS (LocalSpatial) Block module, SS2D (Scanning and Merging) module and RG (ResGated) Block module. SS2D module is used to scan and merge image features to extract global information in multiple directions. LS Block module is used to extract local spatial information to enhance the model's ability to capture local features. RG Block module is used to combine gating mechanism and residual connection to capture local dependencies to enhance the robustness of the model. At the same time, ACConv (Asymmetric ConvolutionBlock) is used at the input end of the ODSSBlock structure to replace the traditional ordinary two-dimensional convolution for convolution operation, and deep separable DW-ACConv convolution is used in each module (LS Block module, SS2D module and RG Block module) for convolution operation. Deep separable DW-ACConv convolution is formed by replacing the ordinary convolution unit in the deep separable convolution DW-Conv with ACConv convolution unit. ACConv convolution can reduce the amount of calculation and the number of parameters while maintaining or improving the performance of the model by decomposing the traditional symmetric convolution kernel into asymmetric convolution kernels in two directions. ODSSBlock uses the state-space model to capture and process image features, so that the LS Block module, the SS2D module, and the RG Block module can deeply extract complex features in the image, including but not limited to shape, texture, and dynamic changes, which can help the model better perform classification and positioning in subsequent steps.

[0051] Specifically, the data processing flow of ODSSBlock includes: S1.11. Use deep separable convolution DW-Conv to perform preliminary feature abstraction on the input features, then perform batch normalization on the output features, and use activation functions to better maintain the distribution of feature map information.

[0052] S1.12, the output features are input to the LS Block module to enhance the model's ability to capture local features, and the feature map is normalized. The normalized feature map is input to the SS2D module to extract and fuse the input features globally, so as to enhance the contextual semantic information of the feature layer and improve the feature expression ability of the feature layer. Finally, the residual cascade is used to combine with the input feature Z l-2 The sum is the output Z l-1 .

[0053] S1.13 Z l-1 Normalization is performed, and the normalized post-processed output is input to the RG Block module to capture more global features and is cascaded with the input feature Z through the residuall-1 The final output Z is obtained by summing l .

[0054] The Mamba architecture can capture long-range ground dependencies, but is not suitable for extracting local features when processing tasks involving complex scale changes. This embodiment uses the LS Block module in the ODSSBlock module to enhance the capture of local features. Specifically, Figure 4 As shown in the figure, for a given input feature, the LS Block module first performs a depth-separable DW-ACConv convolution to replace the traditional depth-separable convolution DW-conv. By adopting this convolution, each input channel can be operated separately without mixing channel information. Compared with the traditional depth-separable convolution DW-conv, the local spatial information of the input feature map can be extracted more effectively by adopting DW-ACConv, and the amount of calculation will not be increased. Compared with the traditional convolution, the computational cost and the number of parameters can be reduced at the same time by adopting DW-ACConv. After the depth-separable DW-ACConv convolution, batch normalization is performed to provide a certain degree of regularization effect while reducing overfitting. The intermediate state mixes channel information through 1×1 convolution, and better maintains the distribution of information through activation functions, so that the model can learn more complex feature representations. These feature representations can extract rich multi-scale contextual information from the input feature map, and then further use 1×1 convolution to change the number of channels of the feature without changing the spatial dimension, thereby enhancing the feature representation. Finally, the original input is fused with the processed features through residual concatenation, which enables the model to understand and integrate features of different dimensions in the image, thereby improving the robustness to scale changes.

[0055] like Figure 5 As shown, in this embodiment, the SS2D module first performs a linear transformation on the input features, then uses the deeply separable DW-ACConv to abstract the features, and uses the activation function to better maintain the distribution of information, so that the model can learn more complex feature representations; and then uses the Scan module to realize global feature extraction and fusion of the input features to enhance the contextual semantic information of the feature layer and improve the feature expression ability of the feature layer.

[0056] The RG Block module can capture more global features while only slightly increasing the computational cost. Figure 6As shown, in this embodiment, the RG Block module creates two branches from the input and implements a fully connected layer in the form of a 1×1 convolution on each branch; a deep separation DW-ACConv convolution is used on the right branch as a position encoding module, and the gradient is more effectively reflected through the residual cascade during the training process, so that it has a lower computational cost, and the performance is significantly improved by retaining and using the spatial structure information of the image, which can effectively reduce the computational cost and improve the performance of the model. Specifically, the RG block can use nonlinear GeLU as an activation function to control the information flow at each level, and then merge with the left branch through element multiplication, and then refine it with the global feature through 1×1 convolution to mix channel information, and finally sum it with the features in the original input through residual cascade. The Scan module implements global feature extraction and fusion of the input features to enhance the contextual semantic information of the feature layer and improve the feature expression ability of the feature layer. Finally, the output of the Scan module is batch normalized to provide a certain degree of regularization while reducing overfitting.

[0057] In a specific application embodiment, during the ACConv training phase, such as Figure 7 As shown in the figure, the model training is performed by using 3×3-BN, 1×3-BN and 3×1-BN multi-branch structures to replace the traditional 3×3-BN structure. -BN means that a batch normalization layer is set after the convolution layer, and the output feature layers of each branch (feature layer 1~feature layer 3) are summed to obtain the final feature layer 4, and then the GELU activation function is used for nonlinear transformation. After the training is completed, as shown in Figure 8 As shown in Figure 2, during ACConv inference, each 1×3-BN and 3×1-BN asymmetric kernel and the corresponding BN layer parameters are added to the skeleton of the 3×3-BN branch and the number of BN layers, that is, the cross part of the square kernel, to achieve ACConv inference convolution reparameterization, as shown in Figure 2. Fig. 9 As shown; then the three parallel convolution models during training are converted back to the same structure as the original ordinary convolution. Through the above-mentioned re-parameterized convolution process, the ordinary convolution can enhance the feature recognition ability of the convolution backbone part without increasing the calculation and changing the original structure, thereby further improving the overall model's ability to represent the features of the scene image.

[0058] In a specific application embodiment, the structural principle of the Scan module is as follows: Fig.10 As shown in the figure, the execution process of the Scan module specifically includes: S1.21. Scan expansion: The input image is expanded into a series of sub-images through a scan expansion operation, each sub-image represents a specific direction, and the scan expansion operation is processed along four symmetric directions when observed from a diagonal viewpoint.

[0059] Specifically, the above four directions can be upper left, lower right, lower left, and upper right. The above scanning and expansion layout method can not only fully cover all areas of the input image, but also provide a rich multi-dimensional information library for subsequent feature extraction through the direction transformation of the system, thereby improving the efficiency and comprehensiveness of multi-dimensional capture of image features.

[0060] S1.22. Send the four one-dimensional vectors obtained in step S1.21 to the S6 block module for S6 operation.

[0061] Specifically, the S6 operation is an efficient selective information retention and parallel scanning algorithm that can independently process the flattened one-dimensional vector to ensure that information in all directions is thoroughly scanned, thereby capturing features in different directions and forming a global receptive field without increasing the linear computational complexity, effectively extracting the global features of the image.

[0062] S1.23, scan merging: The four one-dimensional vectors obtained in step S1.22 are fused into a two-dimensional feature output through the scan merging module.

[0063] Step S02. Use an image training set containing multiple head objects to train a multi-person head posture estimation model. During the training process, a loss function based on dynamic weights is used to control the training process. The loss function based on dynamic weights is obtained by weighting the loss function of the head object prediction classification task, the loss function of the head position detection task, and the loss function of the head posture estimation task using dynamic weights. The dynamic weight of each task is determined according to the loss value of the corresponding task to balance the loss of each task.

[0064] The end-to-end head posture estimation algorithm involves three learning tasks: classification, positioning, and posture estimation. The traditional detection method directly uses fixed weights to sum the loss values ​​of each learning task as the final loss value, which will lead to an imbalance between the loss values ​​of each learning task and affect the accuracy of the model. In this embodiment, the detection head includes three tasks: category prediction (detecting the head prediction probability), head position detection (detecting the head detection frame), and head posture estimation. The loss function in the model training process is determined by adopting a dynamic weighting combined with a cubic square root strategy, that is, the loss values ​​of the three tasks are processed by cubic square roots, and dynamic weights are used to weight them to form a total loss function, so as to reduce the impact of the imbalance of the magnitude of the loss values ​​of different tasks and the loss rates of different tasks on the model, solve the problem of imbalanced losses of the three learning tasks in the head posture estimation model, and further improve the accuracy of the head posture estimation model.

[0065] As an optional implementation, the loss function may be specifically expressed as follows: (1) in, represents the loss function based on dynamic weights, Indicates the batch size of the multi-person head pose estimation model during training. , , Respectively represent the tth training time , and The dynamic weights of the three will change with the change rate of the loss value in the number of training rounds. represents the loss function of the head detection branch, represents the loss function of the head prediction probability branch, Represents the loss function of the head pose estimation branch.

[0066] The dynamic weight of each branch is calculated according to the following expression: (2) (3) in, represents the relative decreasing rate of the loss value of the pth task, represents the dynamic weight of the pth task, p represents the order of the tasks, p=1, 2, 3, corresponding to the category prediction task, head position detection task, and head posture estimation task, respectively. represents the number of training times, represents the total number of learning tasks, Represents the loss value of the p-th task; T Indicates the preset temperature, which is a fixed value used to control the softness of task weighting. T This will make the distribution between different tasks more even. T If it is large enough, then ≈1, that is, the weights of tasks are equal.

[0067] According to the above formulas (2) and (3), the dynamic weights corresponding to the three tasks can be determined respectively: , , , through these three real-time dynamic weights , , , which can effectively reduce the impact of learning rate imbalance between different learning tasks on the model. Specifically, by calculating the relative decreasing rate of task loss value through formula (3), the overall loss value of the model will show a downward trend during the training process. If the loss value of a task decreases significantly during the training process, its loss value will be greatly reduced, and its relative decreasing rate of loss value will be larger. Combined with the dynamic weight of the task obtained through formula (2), it will become larger. Therefore, by adjusting the dynamic weights of different tasks, the loss values ​​between different tasks can be dynamically balanced, reducing the impact of loss rate imbalance on the model.

[0068] The loss function of the head detection branch is constructed based on the focus regression box loss of gamma transform, and the calculation expression is: (4) (5) (6) (7) Where N represents the number of head detections in the image to be detected. , The i-th head prediction detection box and the true label are respectively, represents the Siou-based regression loss, Indicates the preset threshold value with a value range of (0,1). Indicates the loss value of the center coordinates and frame length of the head detection frame. represents the focus regression box loss based on gamma transform, It represents the intersection-over-union ratio between the head region detection frame and the real head frame. It represents the intersection-over-union ratio between the reconstructed head region detection frame and the real head frame. represents the gamma shift adjustment factor, Represents the gamma shift adjustment factor.

[0069] Bounding box regression plays a crucial role in the field of object detection. The positioning accuracy of object detection largely depends on the loss function of bounding box regression. In traditional end-to-end head pose estimation algorithms, the head detection box branch usually adopts a loss function based on IOU (Intersection over Union) and its improvements. For example, the geometric relationship between bounding boxes is used to improve the regression performance. This type of loss function ignores the influence of the distribution of easy and hard samples on bounding box regression, thus further affecting the accuracy of head pose estimation. Hard samples refer to samples containing rich information, which can be manifested as a relatively large loss value for this sample during model training, making it difficult for the model to learn. Due to the rich information content, such samples are more conducive to model learning. Easy samples refer to samples containing less information, which can be manifested as a relatively small loss value for this sample during model training, making it easy for the model to learn. Therefore, an excessive number of easy samples will affect the training accuracy of the model. The loss function for the head position detection task in this embodiment adopts the focal regression box loss based on gamma transformation (as shown in Equation (6)), enabling attention to regression samples of different difficulties, enhancing the model's attention to the regression boxes of hard samples, improving the regression accuracy of the head detection box, and thus effectively improving the performance of the detector in different detection tasks.

[0070] As shown in Equation (7), when the value is relatively large, it indicates that the model has learned the sample detection regression box well at this time, and the sample at this time tends to be an easy sample. If the , then after transformation, the result > . Substituting the result of Equation (7) into Equation (6), it can be seen that as increases, its L Gama-Siou is a decreasing function and the loss decrease rate will increase. At this time, for easy samples, their loss proportion will decrease; by analogy, when the value is relatively small, it indicates that the model has learned the sample detection regression box poorly at this time, and the sample at this time tends to be a hard sample. If the < d, then after transformation by Equation (7), the result < . Substituting it into Equation (6), it can be seen that as becomes smaller, its L Gama-Siou will be an increasing function and the loss increase rate will increase. At this time, for hard samples, their loss proportion will increase. Through the above mechanism, it can effectively make different difficult regression samples be concerned during object detection.

[0071] The head object prediction classification task is to determine whether the current object is a head region. As an alternative implementation, the loss function of the head object prediction classification task can be calculated according to the following formula: (8) Where, represents the probability of the head predicted by the i-th head detection box, represents the real label of the head of the i-th head detection frame, It represents the overlap between the ith head detection frame and the real head frame. According to the above formula (8), the head object classification prediction loss can be calculated, that is, whether the current object is a head region.

[0072] Considering that the rotation representation method with less than four dimensions is discontinuous and not suitable for neural network learning, this embodiment adopts the 6D rotation representation method and estimates the head posture parameters based on the 6D rotation representation matrix, which can realize continuous rotation matrix representation, which is beneficial to neural network learning and further improves the accuracy of end-to-end one-stage multi-person head posture estimation. Based on the Mamba-YOLO framework, this embodiment adds a 6D decoupled detection head for head posture estimation to predict the head posture rotation matrix. The 6D branch at each anchor point predicts 6 values. , these 6 values ​​form two three-dimensional vectors:

[0073]

[0074] Then transform the above two three-dimensional vectors according to the following formula; ; ; (9) in, They are u exist x , y , z The coordinate values ​​of the axis, They are v exist x , y , z The coordinate values ​​of the axis, , , Respectively represent the 6D rotation parameter matrix Transformed into three three-dimensional unit vectors in a 3×3 rotation matrix, , , , ~ Respectively , , Each element in Representation and A vector that is perpendicular in three-dimensional space.

[0075] The two three-dimensional vectors corresponding to the 6D rotation parameters of each head object's head posture are , According to the above conversion, we can get the head posture estimation matrix By deriving the elements in , we can get an orthogonal matrix for head posture estimation prediction:

[0076] The label rotation matrix obtained by Euler angle is also an orthogonal matrix:

[0077] in, , , are three mutually orthogonal unit vectors of the true head posture.

[0078] The multi-person head posture model needs a loss function to measure the distance between the two matrices, so that the model can continuously optimize parameters and predict the head posture rotation matrix, and then obtain the head posture Euler angle. In order not to destroy the SO(3) popular geometry structure of the head posture 3×3 rotation matrix, this embodiment uses a method based on cosine similarity to characterize the head posture estimation matrix. With the real head pose matrix The similarity between .

[0079] Specifically, this embodiment uses a dense detection network to directly predict a group of head objects ,in Represents the head prediction score and position parameters, Represents the head pose 6D rotation representation parameters); The specific parameters are , Represents the probability of the current prediction being the head, Represents the center position of the current predicted head detection box, Represents the width and height of the current predicted head detection box; The specific parameters are , represents the 6D rotation representation parameter, where u and v represent two non-parallel vectors. The 6D rotation representation parameter is converted into a 3×3 rotation matrix according to the above formula (9). As an attribute of the head object, compare it with Combine to form the head prediction object , it is possible to achieve end-to-end one-stage multi-person head pose high-precision estimation based on 6D rotation representation.

[0080] In order to calculate the similarity between the two 3D rotation matrices and form the head posture estimation branch loss function, it is considered that the easier the sample, the smaller the loss value, the harder the sample, the larger the loss value, and when x is closer to , the value of y is The characteristics are consistent with those of the tangent function. The loss function of the head posture estimation task in this embodiment is constructed based on the tangent function to update and optimize the head posture estimation parameters, so that the model pays more attention to the learning of difficult-dimensional vectors and further improves the accuracy of head posture estimation.

[0081] Optionally, the loss function of the head pose estimation task is calculated as: (10) Among them, i represents the i-th dimension of the head posture matrix, represents the loss weight on the i-th dimension of the head posture matrix, represents the head pose estimation matrix, represents the true head pose matrix, , express , The elements in .

[0082] like Fig.11 As shown, the horizontal coordinate x of the tangent function is in the interval When , the range of its value y , and as x increases, the growth rate of y increases. This function change trend meets the needs of difficult sample mining (the easier the sample, the smaller the loss value should be, the larger the loss value of the difficult sample should be, and when x is closer to , the value of y is ), so using the tangent function to construct the loss function for the head pose estimation task can make the model pay more attention to the learning of difficult dimensional vectors. As shown in formula (10), considering the limit of y, the maximum value of x is set to ( The value range is [0,1]), so that the y value is within a certain range, and at the same time, by setting The ratio between the three can be further adjusted dynamically. If the similarity of a dimension component in the two matrices is small, the loss proportion contributed by it will be further increased, so that the head posture model can further focus on difficult samples, thereby achieving the purpose of further improving the accuracy of head posture estimation.

[0083] Step S03: Obtain a picture to be detected containing multiple heads, input it into a trained multi-person head posture estimation model, and obtain detection results of multiple head objects.

[0084] After inference, the multi-person head posture algorithm model will generate a series of head detection objects, which need to be post-processed to obtain the final real detection object. Non-maximum suppression (NMS) is a common post-processing method used by many target detection models. However, since the multi-person head posture estimation scene belongs to dense target detection, there are different degrees of overlap between targets. The traditional non-maximum suppression algorithm (NMS) only relies on the highest confidence and Iou of the target classification, which leads to a certain degree of inaccurate positioning and missed detection in overlapping detection scenes. This embodiment adopts a non-maximum suppression algorithm that does not rely on confidence and Iou to improve its detection accuracy for overlapping scenes, thereby improving the head posture estimation accuracy. Fig.12 As shown, the specific algorithm steps are as follows: 1) Filter out detection boxes whose confidence is less than the preset threshold; 2) Calculate the proximity P between the two detection frames used: First, normalize the relevant coordinates, calculate the center coordinates, length and width of the detection frame, and then use the Manhattan distance based on logarithmic transformation to calculate the proximity. The calculation formula is shown in Equation 11. (11) In the above formula, x1 represents the x-coordinate of the center point of the first detection frame, and y1 represents the y-coordinate of the center point of the first detection frame. w1 represents the width of the first detection box, h1 represents the height of the first detection box, and the same is true for x2, y2, w2, and h2.

[0085] 3) Group the detection frames with proximity P <= d (d is the proximity threshold (which can be determined according to the specific scenario) into a cluster; 4) Calculate its confidence-weighted proximity; 5) Find the detection box with the smallest confidence-weighted proximity in a cluster, save it and filter it out from the total box; 6) Filter out other detection frames whose proximity is less than the preset threshold k (d is the proximity threshold, which must be determined according to the specific scenario); 7) Repeat steps 2-6 for the remaining boxes.

[0086] This embodiment further provides an electronic device, including a processor and a memory, wherein the memory is used to store a computer program, and the processor is used to execute the computer program to perform the above method.

[0087] It is understandable that the above method of this embodiment can be executed by a single device, such as a computer or server, etc., and can also be applied to a distributed scenario and completed by multiple devices in cooperation with each other. In the case of a distributed scenario, one of the multiple devices can only execute one or more steps in the above method of this embodiment, and multiple devices interact to complete the above method. The processor can be implemented in the form of a general-purpose CPU, a microprocessor, an application-specific integrated circuit, or one or more integrated circuits, etc., for executing related programs to implement the above method of this embodiment. The memory can be implemented in the form of a read-only memory ROM, a random access memory RAM, a static storage device, and a dynamic storage device. The memory can store an operating system and other applications. When the above method of this embodiment is implemented by software or firmware, the relevant program code is stored in the memory and called and executed by the processor.

[0088] This embodiment further provides a computer-readable storage medium storing a computer program, and when the computer program is executed by a processor, the above method is implemented.

[0089] Those skilled in the art should understand that the above-mentioned embodiments of the present application can be provided as methods, systems, or computer program products. Therefore, the present application can take the form of a complete hardware embodiment, a complete software embodiment, or an embodiment combining software and hardware. Moreover, the present application can take the form of a computer program product implemented on one or more computer-readable storage media (including but not limited to disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code. The present application is described with reference to the flowcharts and / or block diagrams of the methods, devices (systems), and computer program products according to the embodiments of the present application. It should be understood that each process and / or box in the flowchart and / or block diagram, as well as the combination of the processes and / or boxes in the flowchart and / or block diagram, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer, an embedded processor, or other programmable data processing device to produce a machine, so that the instructions executed by the processor of the computer or other programmable data processing device generate instructions for implementing the processes in the process Figure 1 A process or multiple processes and / or boxes Figure 1 These computer program instructions can also be stored in a computer-readable memory that can guide a computer or other programmable data processing device to work in a specific way, so that the instructions stored in the computer-readable memory produce a product including an instruction device, which implements the functions specified in the process. Figure 1 A process or multiple processes and / or boxes Figure 1These computer program instructions can also be loaded onto a computer or other programmable data processing device, so that a series of operation steps are executed on the computer or other programmable device to produce a computer-implemented process, so that the instructions executed on the computer or other programmable device provide for implementing the process in the process. Figure 1 A process or multiple processes and / or boxes Figure 1 A step that specifies a function in one or more boxes.

[0090] The above is only a preferred embodiment of the present invention, and does not limit the present invention in any form. Although the present invention has been disclosed as a preferred embodiment, it is not intended to limit the present invention. Therefore, any simple modification, equivalent change and modification made to the above embodiment according to the technical essence of the present invention without departing from the content of the technical solution of the present invention shall fall within the scope of protection of the technical solution of the present invention.

Claims

1. A method for multi-person head pose estimation based on multi-task loss balance, characterized in that the steps include: A multi-person head posture estimation model is constructed using a YOLO network based on the Mamba model. The multi-person head posture estimation model includes a backbone network, a PAFPN module, and a detection head. The backbone network is an ODMamba network structure for processing input images and extracting preliminary features. The PAFPN module is used to extract and fuse features of different scales. The detection head is used to achieve category prediction, head position detection, and head posture estimation tasks. Using an image training set containing multiple head objects to train the multiple head posture estimation model, during the training process, a loss function based on dynamic weights is used to control the training process, wherein the loss function based on dynamic weights is obtained by weighting the loss function of the head object prediction classification task, the loss function of the head position detection task, and the loss function of the head posture estimation task using dynamic weights, and the dynamic weight of each task is determined according to the loss value of the corresponding task to balance the loss of each task; Obtain a picture to be detected containing multiple heads, input it into the trained multi-person head posture estimation model, and obtain detection results of multiple head objects.

2. The method for multi-person head pose estimation based on multi-task loss balance according to claim 1, characterized in that: The ODMamba network structure includes an image segmentation module, an ODSSBlock module, a visual cue merging module and an SPPF module. The image segmentation module is used to perform preliminary processing on the input image through multiple convolutional layers, and the ODSSBlock module is used to perform feature extraction; the visual cue merging module is used to fuse and optimize the features output by the ODSSBlock module, and the SPPF module is used to aggregate features at multiple scales; the core module in the PAFPN module is the ODSSBlock module; ACConv convolution is used to perform convolution operations at the input end of the ODSSBlock module, and deep separable DW-ACConv convolution is used to perform deep separation convolution operations in the LS Block module, SS2D module and RG Block module in the ODSSBlock module, the SS2D module is used to scan and merge image features to extract multi-directional global information, the LS Block module is used to extract local spatial information, and the RG Block module is used to capture local dependencies by combining a gating mechanism and a residual connection.

3. The method for multi-person head pose estimation based on multi-task loss balance according to claim 1, characterized in that: The calculation expression of the loss function based on dynamic weight is: in, represents the loss function based on dynamic weights, Indicates the batchsize of the multi-person head pose estimation model during training. , , Respectively represent the tth training time , and The dynamic weights of each dynamic weight change with the change rate of the loss value in the number of training rounds. represents the loss function of the head position detection task, represents the loss function of the category prediction task, Represents the loss function for the head pose estimation task.

4. The method for multi-person head pose estimation based on multi-task loss balance according to claim 3, characterized in that: The dynamic weight of each branch is calculated according to the following expression: in, represents the dynamic weight of the pth task, p=1, 2, 3, corresponding to the category prediction task, head position detection task, and head posture estimation task, respectively. represents the number of training times, Indicates the total number of tasks, represents the loss value of the pth task, T Indicates the preset temperature used to control the softness of task weighting.

5. The method for multi-person head pose estimation based on multi-task loss balance according to claim 3, characterized in that: The loss function of the head position detection task is constructed based on the focus regression box loss of gamma transform, and the calculation expression is: Where N represents the number of head detections in the image to be detected. , The i-th head prediction detection box and the true label are respectively, represents the Siou-based regression loss, Indicates the preset threshold value with a value range of (0,1). Indicates the loss value of the center coordinates and frame length of the head detection frame. represents the focus regression box loss based on gamma transform, It represents the intersection-over-union ratio between the head region detection frame and the real head frame. It represents the intersection-over-union ratio between the reconstructed head region detection frame and the real head frame. represents the gamma shift adjustment factor, Represents the gamma shift adjustment factor.

6. The method for multi-person head pose estimation based on multi-task loss balance according to claim 3, 4 or 5, characterized in that: The loss function of the head object prediction classification task is calculated according to the following formula: in, represents the probability of the head predicted by the i-th head detection box, represents the real label of the head of the i-th head detection frame, Indicates the overlap between the i-th head detection frame and the true head frame.

7. The method for multi-person head pose estimation based on multi-task loss balance according to claim 3, 4 or 5, characterized in that: The loss function of the head posture estimation task is constructed based on the tangent function, and the calculation expression is: Among them, i represents the i-th dimension of the head posture matrix, represents the loss weight on the i-th dimension of the head posture matrix, represents the head pose estimation matrix, represents the true head pose matrix, , express , The elements in .

8. The method for multi-person head pose estimation based on multi-task loss balance according to claim 7, characterized in that: It also includes obtaining the head pose estimation matrix The steps include: Use the head pose estimation 6D decoupled detection head set in the multi-person head pose estimation model to detect the head pose 6D rotation representation parameters ; The two unrelated three-dimensional vectors corresponding to the 6D rotation parameters of the head pose of each head object are , Convert to get the head pose estimation matrix The elements in: ； ； ； in, They are u exist x , y , z The coordinate values ​​of the axis, They are v exist x , y , z The coordinate values ​​of the axis, , , Respectively represent the 6D rotation parameter matrix Transformed into three three-dimensional unit vectors in a 3×3 rotation matrix, , , , ~ Respectively , , Each element in Representation and A vector that is perpendicular in three-dimensional space.

9. An electronic device comprising a processor and a memory, wherein the memory is used to store a computer program, wherein: The processor is configured to execute the computer program to perform the method according to any one of claims 1 to 8.

10. A computer-readable storage medium storing a computer program, characterized in that: When the computer program is executed by a processor, the method according to any one of claims 1 to 8 is implemented.

Images