Mine underground object detection algorithm based on TTFNet and anchor-free

2.1 CenterNet algorithm

In contrast, CenterNet [19] is more concise and can achieve higher accuracy and speed. It is easy to extend to other tasks [20] such as target tracking, pose detection, instance segmentation, and 3D detection. It can be extended to a variety of scenes in the mine. As shown in Figure 1, CenterNet models the target as a point to complete the identification and localization of the object in an image. First, the image is input into the convolutional neural network to generate a Heatmap, and the peak in the Heatmap is the center of the corresponding category object, and then, the object detection is realized by directly regressing the bounding box from the center.

Figure 1

Centernet algorithm schematic diagram.

The structure of the CenterNet algorithm is shown in Figure 2. The network inputs an image with a size of 512 × 512 . The backbone network will first extract a feature map with a size of 16 × 16 , and then upsample it to a size of 128 × 128 . The three branches of the input prediction module, respectively, predict the center, offset, and border width and height of the target.

Figure 2

CenterNet algorithm structure.

The loss function of CenterNet mainly consists of three parts: type, center point offset, and frame size loss. For any input image I ∈ R W × H × 3 , the goal of the CenterNet algorithm is to obtain heatmap through convolutional neural networks:

where W and H represent the width and height of the image, r = 4 represents the output stride, and C represents the number of classes. When Y ˆ x y c = 1 , it means that ( x , y ) is the center point of C category on the heatmap, and when Y ˆ x y c = 0 , it is the background, and only the center point is regarded as the positive sample.

On the downsampling feature map, we calculate the radius of the Gaussian circle based on the size of the object box and then apply a Gaussian kernel:

splat all ground truth center point onto heatmap Y ∈ [ 0 , 1 ] W r × H r × C . In formula (2), σ p is a standard deviation related to the size of the target. If there is an overlap in the Gaussians distribution of multiple targets of the same category, the target corresponding to the maximum value of elements in each position of the overlap area is taken. P ˜ x and P ˜ y represent the values of the center point coordinates on the downsampled picture and rounded down, respectively. In addition, for a category in the label graph, the algorithm calculates the center coordinates of the actual border:

where x 1 x 2 y 1 y 2 represents the coordinate point of the object’s box. Therefore, given the prediction heatmap value Y ˆ x y c and calculated heatmap value Y x y c , the location loss is designed as

where α and β represent the hyperparameters of focal loss, which are, respectively, taken as 2 and 4, and N is the number of key points in the image. This method can adjust the imbalance of positive and negative samples and improve the mining ability of difficult samples.

Since the prediction of the center point of the target is performed on the downsampled feature map, a precision error will occur when it is remapped to the original image size, so the algorithm adds a center point offset loss, and the formula is

where O p is the predicted offset of the network.

On the size regression branch, let any target box belonging to category C k be represented as ( x 1 ( k ) , y 1 ( k ) , y 2 ( k ) , x 2 ( k ) , y 2 ( k ) ) , CenterNet selects s k = ( x 2 ( k ) − x 1 ( k ) , y 2 ( k ) − y 1 ( k ) ) as the regression target for the target size prediction and represents the predicted value with s ˆ ∈ R W r × H r × 2 , so the specific dimensional loss can be expressed as

To sum up, the total loss function of the CenterNet algorithm can be expressed as

where λ size = 0.1 , λ off = 1 is the weight coefficient of the loss function.

2.2 TTFNet detection algorithm

Although CenterNet can achieve good comprehensive performance, its training convergence rate is slow compared with other algorithms. TTFNet believes that the reason for this problem is that only the target midpoint is regarded as a positive sample for size regression during CenterNet training. In order to shorten the training time, TTFNet proposed to use all pixel points within a certain range of the target center as the sample of border regression for training and adopted Gaussian probability as the weight of each regression sample, emphasizing those samples close to the target center. This method allows the network to use more positive samples for training, thus speeding up the algorithm convergence in a similar way to increasing the learning rate. Combined with CenterNet’s original design, TTFNet achieves a better balance between training time, reasoning speed, and accuracy.

As shown in Figure 3, TTFNet splits object detection into two parts: object localization and size regression. In terms of center positioning, considering the actual aspect ratio of the object, TTFNet further uses two-dimensional Gaussian kernels with different widths and heights to generate heatmap. Given mth annotated box belongs to the Cm class, it is first mapped linearly onto the feature map scale. And then 2D Gaussian kernel with different widths and heights

that are used to produce H m ∈ R 1 × H r × W r is generated with σ x = η w ∕ 6 , σ y = η h ∕ 6 , and η = 0.54 , w and h represent the width and height of the box, respectively, and the peak of the heatmap is the center of the corresponding box C m . On the other hand, TTFNet will use a similar method to CenterNet to predict the H ˆ m ∈ R 1 × H r × W r . During training, only the peak value of Gaussian distribution is regarded as a positive sample. Similar to CenterNet, given the prediction heatmap value H ˆ x y c and calculated heatmap value H x y c the positioning loss is represented by the corrected focal loss, as shown in the following:

where M is the number of annotated boxes.

Figure 3

TTFNet algorithm structure.

In terms of regression, as shown in Figure 4, TTFNet can use the same Gaussian kernel as the positioning branch to construct a sub-region around the target center and treat all pixels in the region as positive samples of size regression. The target of the algorithm regression is the distance of the sample point to the four boundaries. For one pixel ( i , j ) in the subarea and the downsampling rate r , the target of regression is the distance ( w l , h t , w r , h b ) i j m from ( i r , j r ) to the four boundaries of the corresponding box, so the predicted box at ( i , j ) can be expressed as

where s = 16 is a fixed scalar used to increase the prediction for optimization purposes; therefore, the algorithm does not need to set additional center offset loss. ( x ˆ 1 , y ˆ 1 , x ˆ 2 , y ˆ 2 ) and ( w ˆ l , h ˆ t , w ˆ r , h ˆ b ) are the values predicted by the network. TTFNet uses the Generalized Intersection over Union (GIOU) loss between the ground truth and the bounding box as the loss function of the regression branch, which is expressed as

where N reg is the number of regression samples, B ˆ i j is the bounding box, and B m is the ground truth. Furthermore, considering that the scale difference of the target will cause the difference in the number of samples, to improve the ability of the model to detect small targets, W i j is used to balance the loss values contributed by targets of different sizes, and the algorithm designs different sample weights W i j as follows:

where a m is the area of the m th annotated box, G m ( i , j ) is the Gaussian probability value at ( i , j ) , and A m is the Gaussian region.

Figure 4

CenterNet and TTFNet define the training sample strategy.

In summary, the total loss of the TTFNet algorithm can be expressed as follows:

where w loc = 1.0 and w reg = 5.0.

2.3 Algorithm improvement

The performance of the TTFNet algorithm is very excellent, which can shorten the training time by several times or even more than ten times while maintaining a very fast reasoning speed and better accuracy than CenterNet. Therefore, we made the following improvements based on the TTFNet algorithm.

2.3.1 Backbone network

Wang et al. [21] attributed the reason that the previous neural network architecture requires a large amount of reasoning calculation to the repetition of gradient information in network optimization, and they designed the CSPNet network structure as shown in Figure 5 using the idea of splitting and merging. In this structure, the feature map is first divided into two parts according to the proportion, which is propagated forward, respectively, and then merged into a complete feature map. This structure can not only increase the gradient propagation path but also reduce the network memory flow, balance the calculation, and improve the network learning efficiency.

Figure 5

CSPNet structure.

The underground object detection task is not many types of targets, and the computing equipment is limited, so it does not need too deep and too wide backbone network to extract features. But considering the complex underground environment, in order to enhance the robustness of the algorithm and further improve the computing speed, this article introduces pooling in the basic module of CSPNet. When the downsampling rate is 2, the CSPNet structure is improved as shown in Figure 6. The right structure of the module is kept unchanged, and 2 × 2 average pooling or maximum pooling is added before the left convolution to realize downsampling. To facilitate the description, these two structures are called AvgCSPNet and MaxCSPNet, respectively.

Figure 6

The CSPNet structure is used in this paper.

2.3.2 Multi-layer feature fusion

As shown in Figure 7(a), the multi-layer feature fusion method adopted by the TTFNet algorithm is to directly add the low-level features with the same dimension to the upsampled features. However, this method may weaken some learned features when adding. In order to maximize the retention of these features, this study changes it to the splicing of the low-level features and the upsampled features along the channel dimension. Then, 1 × 1 convolution is used to reduce the dimension to the required dimension to achieve multi-layer feature fusion, and the improved structure is shown in Figure 7(b).

Figure 7

Multi-layer feature fusion. (a) TTFNet up-sample, and (b) Ours up-sample.

2.3.3 Residual shrinkage structure

Due to the harsh underground working environment, dust is easily produced; coupled with uneven lighting and equipment movement, the photographs taken by the camera on the scraper usually contain a lot of noise. In addition, due to the irregular shape of the target during the image annotation process, it is inevitable to contain some background information or overlap the target, which can also be regarded as noise during the target detection process. Therefore, it is necessary to strengthen the network’s ability of the network to deal with noise.

Deep residual shrinkage network [22] is a kind of network structure used for fault diagnosis. Its core idea is to let the neural network learn a group of small positive thresholds and then set the weights whose absolute values are less than the thresholds to zero. In theory, because the noise contributes little to the detection target and the weight value is lower, it will be set to zero in the training process, so as to achieve the purpose of noise reduction and accelerate the convergence speed of the algorithm. Therefore, this method is introduced into the network to reduce the impact of underground noise on detection. The residual shrinkage unit (RSBU) can be regarded as a special filter, and its basic structure is shown in Figure 8(a). The structure first takes the absolute value of the input network features, obtains the mean value of each channel, and then learns the scale factors of each layer through the two fully connected layers and the sigmoid function. Finally, the threshold value of each channel is obtained by multiplying it with the mean value of the corresponding layer, so as to achieve soft thresholding. Let x c and y c represent the input and output values of the C channel, respectively, and the soft threshold method can be described as

(14) y c = x c − τ c x c > τ c , 0 − τ c ≤ x c ≤ τ c , x c + τ c x c < − τ c ,

(15) τ c = α c ⋅ average h , w ∣ x h , w , c ∣ ,

(16) α c = 1 1 + e − z c ,

where τ c is the threshold of the C channel, α c is the scale factor, and Z c is the output of the fully connected layer. The residual shrinkage unit uses two fully connected layers to achieve the purpose of learning the soft threshold, but the process involves dimensionality reduction operations, and the fully connected layer has low efficiency and slow propagation speed, which consumes computing resources and time.

Figure 8

Structure of residual contraction unit before and after improvement. (a) Residual shrinkage unit and (b) efficient residual shrinkage unit.

Therefore, the original structure is improved as shown in Figure 8(b), and the fully connected layer in the original structure is replaced by 1-D convolution, so as to avoid dimension reduction in the operation process, improve the efficiency of the neural network, and reduce the consumption of computing resources [23]. In addition, the threshold learned through the fully connected layer is only the processing learning of the single-layer feature layer, and the interaction between multiple different feature layers can be used to further enhance the learning ability of the network after replacing it with 1D convolution.

In addition, we use the SiLU function as the activation function to retain some negative features and enhance the nonlinear effect while achieving a certain regularization effect. The improved structure is called the efficiency residual shrinkage unit (ECRSBU). The high-efficiency residual collection unit only adds a very small amount of parameters and calculations, and its role can be regarded as a special attention mechanism, and it can also achieve the effect of plug-and-play.

2.3.4 Overall structure of the algorithm

A feature extraction network as shown in Figure 9 is constructed on the basis of the CSPNet network structure. After preprocessing and data enhancement, the original image is resized to 512 × 512 . Input it into the backbone network. After five times downsampling, the 16 × 16 size dimension feature is obtained. Then, the 128 × 128 size feature is obtained by three times upsampling using deformable convolution and bilinear interpolation. At the same time, the shallow features are fused with the current features during each upsampling, and finally, the network output is obtained by the decoding of the input detection head after ECRSBU processing.

Figure 9

Overall structure.

3.1 Data preparation

The research of this study is based on the largest underground non-ferrous metal mine in China, Shangri-La Pulang Copper Mine in Yunnan Province, whose underground LHD model is Sandvik LH514, which is operated by remote video remote control. All the data in the experiment are collected from different angles of the LHD camera during operation, with a size of 1,280 × 720 pixels. Some frames were selected as training data, and a total of 2,149 data sets were obtained. The establishment of the experimental data set marked three categories of ore in the bucket, ore heap, and bucket to study the automatic shoveling and working state judgment of the scraper. All images were annotated by Labeling software after preprocessing. The labeling is shown in Figures 10 and 11 shows the data set details.

Figure 10

Downhole data labeling.

Figure 11

Dataset details.

3.2 Experiments and results

In this study, all the experiments are built and trained in the win10 system, AMD Ryzen7 4800H CPU, NVIDIA GeForce GTX 1650Ti GPU, pytorch1.8.1+CUDA11.1 deep learning framework, and python 3. 9.2 programming language environment.

In order to evaluate the quality of the model, this study uses the commonly used indicators in object detection, such as precision (P), recall (R), mean average precision (mAP), and mAP @.5:.95 to evaluate the model comprehensively.

In the experiment, the improved model adopts mosaic, flip, and color transformation, and other data enhancement methods, and SiLU is used as the activation function. Because the TTFNet algorithm converges quickly, only 24 or 60 epochs are trained. The network input image size is 512 × 512 , the minimum training batch is set to 4, Adam optimizer is selected, momentum is set to 0.9, cosine annealing strategy is used for learning rate setting, and the initial setting is 0.0002.

3.3 Backbone network comparison experiment

Figure 12 shows the convergence of different backbone networks when trained on the downhole data set. It can be seen that all the networks tend to converge after training 20 epochs, and the network structure proposed in this study can still converge to a similar level with a large number of reduced parameters.

Figure 12

Loss function.

Table 1 shows a comparison of the specific data from the experiment. It can be seen from the table that the recall rate and mAP of the improved network model are slightly reduced, but the accuracy is improved, and the overall level is similar. At the same time, MaxCSPNet with maximum pooling on network branches had a 1.5% higher recall rate, 4.3% higher mAP.5:.9, and 1.7% higher mAP than AvgCSPNet using average pooling on branches, indicating that MaxCSPNet feature extraction effect was significantly better than the latter. In addition, the MaxCSPNet backbone network can achieve close accuracy and reduce computational overhead while significantly reducing the number of parameters compared with other backbone networks.

3.4 Ablation experiment

Table 2 shows the experimental comparison results of adding different structures to the model. The performance of the algorithm is improved by improving the network feature fusion method. Adding the RSBU structure to the backbone network improves the precision of the model by 0.5%, the recall by 0.9%, and the mAP@.5:.95 by 0.5% on the underground data set, which illustrates the effectiveness of RSBU in the object detection algorithm. When replacing RSBU with ECRSBU, the precision decreases by 0.5%, while the recall improves by 0.2% and mAP@.5:.95 improves by 0.7%. Moreover, ECRSBU introduces fewer parameters to the network and is computationally more efficient. However, when the CBAM attention structure is further added to the model, the overall performance of the algorithm decreases rather than improves as expected. When it is replaced by efficient channel attention (ECA) attention structure, the performance of the algorithm is further improved. The possible reason is that the residual shrinkage structure itself is equivalent to a spatial attention mechanism, while the addition of mixed attention CBAM will lead to its excessive effect and eventually reduce the performance of the algorithm. Therefore, algorithm performance can be further improved when combined with ECA that contains only channel attention.

In order to explore the influence of the residual shrinkage structure and the relative position of the attention mechanism in the network on the algorithm, experiments are conducted as shown in Table 3 and Figure 13 shows the corresponding loss function changes. The results show that the algorithm can achieve faster speed and higher performance by placing the residual shrinkage structure in front of the attention mechanism structure.

Figure 13

Different structural sequence loss.

3.5 Comparative experiment

To further verify the performance level of the algorithm, we compared the algorithm in this article with the current mainstream object detection algorithm on the underground data set, and the results are shown in Table 4. It can be seen that the Faster RCNN algorithm has a high recall rate on the underground data set, but it has a lot of misrecognition of the background as the target; the indicators of YOLO series algorithms have reached a high level; the precision of CenterNet algorithm is significantly higher than other algorithms, but its recall rate is low. Compared with the CenterNet algorithm, the precision is increased by 1.6%, the recall is increased by 3.8%, and mAP@.5:.95(%) is increased by 4%, which is close to the performance indicators of the current mainstream object detection algorithms based on anchor and Non-Maximum Suppression (NMS). Compared with other anchor-free algorithms, the performance is slightly lower than YOLOx, but our algorithm is simpler. In addition, through experiments, we find that the algorithm based on CenterNet method always shows a high level of accuracy on the underground data sets with poor image quality, which is very important for the underground equipment to accurately identify the target and make accurate actions.

3.6 Presentation of results

The heatmap and results of the algorithm are shown in Figure 14.

Figure 14

Algorithm-detected heatmap (top) and results (bottom).