A Lightweight underwater detector enhanced by Attention mechanism, GSConv and WIoU on YOLOv8

YOLOv8 algorithm

The object detector can be classified into a single-stage detector and a two-stage detector. The two-stage detectors, such as Fast R-CNN¹⁸, Faster R-CNN¹⁹, Mask R-CNN²⁰, firstly generate candidate regions, and then classify them. So, its inference speed is reduced greatly by its large number of repeated calculation operations. A single-stage object detector directly maps input images to bounding boxes and class probabilities. YOLOv8²¹, a single-stage object detector, is developed by Ultralytics on basis on the YOLOv7 architecture²² to improve accuracy and speed. The YOLOv8 model consists of input, backbone, and head.

In backbone, firstly, by the convolution, containing batch normalization and SiLU activation functions, the down-sampling is done for feature extractions; and then, by the more connections among branches with cross layers of the C2f module, referring to the E-ELAN structure²², the gradient flow of model is enriched for further feature extraction to improve detection results; Lastly, by the joint of three max pooling layers in SPPF block, which is at the end of the backbone network, the various scale of features are processed for enhancing the feature abstraction ability of network. 

In neck network, by FPN²³ and PAN²⁴ structures, the information of various size feature maps are fused for the head.

In the head, a decoupled detection is used to calculate the loss of regression and category by the convolutions of two parallel branches, each of which focuses on its own tasks to improve the detection accuracy of model greatly.

There are five models, YOLOv8n, YOLOv8s, YOLOv8m, YOLOv8l, and YOLOv8x in YOLOv8 series. The computing complexity of YOLOv8n is the lowest of them. So, the improvement of this paper is done on the basis of YOLOv8n.

Data-set

The data-set of Underwater Robotics Competition (URPC) 2020 is used here (shown in Fig. 1). In Fig. 1, a and b show small targets in the data set, c and d show complex terrain, and e and f show dim waters. There are 5543 training images of echinus, holothurian, scallops and starfish in the data-set. The data-set are divided into training, validation, and testing sets in the 7:2:1 ratio. There are some complex situations in the data-set, such as the masking of underwater biological aggregation, the uneven lighting, the blurring of motion lens, the uneven distribution of samples, and their different resolutions of categories. So, all these above make the model training significantly challenging.

Data augmentation

The poor detection performance of models is resulted from the insufficient training samples and the imbalanced data-set categories. To alleviate these unfavorable conditions, some data augmentations are applied to provide better training data for the models, and to improve both performance and generalization ability²⁵:

In underwater target detection, the movement, shape, and direction of underwater organisms are uncertain. So, the training effect of the model can be improved by flipping (shown in Fig. 2).

Mosaic concatenates four images by random scaling, cropping, and layout²⁶. So, the background of detecting objects, especially the background of the small targets is enriched by Mosaic (shown in Fig. 3).

Haar wavelet transformation²⁷ is used for style transfer and image restoration. Haar wavelet transform produces four sub-bands, namely LLT, LHT, HLT and HHT low-pass and high-pass filters. In this paper, the low-frequency sub-band (LL) is used to capture the smooth surface and texture information of the image, and the high-frequency sub-band (LH, HL, HH) is used to extract the edge information. By high-fidelity image restoration, the overall contrast is enhanced, the impact of noise on the image is reduced, the distorted images are restored and enhanced without affecting their original details (shown in Figs. 4,5). In Fig. 5, a is the original image, b is the image after wavelet change, and c is the image after LCAHE-WT transformation. The top image shows images of deep water, and the bottom image shows images of shallow water

When the sunlight passes through water, the energy of red light is absorbed by the water molecules, and the energy of the red band rapidly weakens. The blue and green bands of light have relatively higher energy, and propagate easier in deep water. So, the blue and green band are more pronounced in the images of deep water areas. However, the normal wavelet transform is only used to analyze local features, edges, textures, and other information of an image, and cannot deal with a specific color band. Therefore, the influence of green waves or other color bands cannot be directly reduced.

LCAHE (Limited Contrast Adaptive Histogram Equalization) is a local contrast enhancement method. By LCAHE, the image is divided into small regions, and then histogram equalization is performed in each region to avoid excessive enhancement and noise limitation. The LCAHE and Wavelet Transform are combined (called LCAHE-WT) in this paper to improve fidelity and detail of images, and to adapt to underwater environments. For the images of deep water areas, it works as following (shown In Fig. 6):

1. The images are firstly normalized to the range of 0-1 for wavelet transform;

2. The LL, LH, HL, and HH coefficients are extracted. For low-frequency coefficients, the LL coefficient is enhanced by linearly combining the original LL and its mean. For high-frequency coefficients, the green channels of the HL coefficient are suppressed to reduce the influence of green waves;

3. The enhanced frequency domain information is reconstructed to restore the enhanced image from the range of [0,1] to [0,255].

4. LCAHE-WT is used to improve the local contrast of the image, to make the details of the image clearer in different regions.

The green waves are efficiently suppressed in the high-frequency coefficients of wavelet transform. The reduction degree of green wave can be controlled more finely by preserving other features of the image better. For the images of shallow water areas, it works similarly. But, the red wave coefficients of all frequency band is enhanced to compensate for color shift, and to make the underwater image closer to the real image.

Convolutional block attention module

The attention mechanism is used to shrink and lock the focus area of the network, and to ultimately form an attention focus, which is crucial for object detection. By attention mechanisms, the perceptual information is refined, and the contextual connections are enhanced. There are channel, spatial, and combined channel-spatial attention mechanism, which dynamically measure the significance of input features to improve detection by suppressing noise and capturing relevant information. It is proved that the integration of attention mechanisms into deep neural networks is effective for boosting performance in computer vision tasks. In last years, lots of works have been done to integrate various attention mechanisms into deep convolutional neural network architectures to improve the performance of detection tasks effectively, and they have been proven effective²⁸.

The CBAM (Convolutional Block Attention Module) (shown in Fig. 6)²⁹ enhances the representation capability of network by introducing attention mechanisms. It consists of the channel attention module and the spatial attention module. In CBAM, by emphasizing the meaningful features of both channel and spatial dimensions, essential information is focused on, the unnecessary details are suppressed, and the target recognition of network is more accurate. Moreover, CBAM is a lightweight module, and its parameters and computational overhead is almost no in most cases. So, it is widely applied to enhance the representation capability of convolutional neural networks.

CBAM is divided into two parts, spatial attention and channel attention. In the channel attention, firstly, the global average of each channel is pooled and normalized in each feature graph; Secondly, a channel weight vector is generated by a series of linear transformations (usually fully connected layers); Thirdly, the vector is converted into probability distribution by softmax function, which indicates the importance of each channel; Finally, the important features are enhanced, and the unimportant features are suppressed through weighting channel by channel. In spatial attention, the maximum value and the average value of each position are captured by maximum pooling and average pooling; And then, the local spatial attention is calculated by the same linear transformation, where the softmax function is applied to determine the position importance of each pixel. Normally, These two modules are connected in series, with spatial attention following channel attention, to obtain more comprehensive feature context information.

The backbone plays the role of feature extraction in the whole YOLOv8 network. So, the CBAM module is added here to enforce the importance of the correct information, and to remove some useless information for improving the feature extraction ability and detection accuracy of the network.

GSConv module

The storage and computing resources of underwater equipment are so limited that it is necessary for underwater detection to acquire high accuracy with low computing overhead³⁰. Now, the model compression and the lightweight network architecture design are the two main ways of lightweight.

Knowledge distillation³¹ and network pruning³² are the two main methods for model compression. In knowledge distillation, model size and computing requirements are reduced without affecting performance by a lighter and smaller student model, which are trained from a teacher model. In pruning, the model redundancy is reduced, and the network is made sparsity by removing unimportant channels or convolutional kernels.

The design of lightweight network architecture is to make convolutional neural networks more lightweight by different convolutional methods and structures. Such as SqueezeNet³³, MobileNet³⁴, ShuffleNet³⁵, etc. In SqueezeNet, a fire module structure, which is the pioneer in model compression, is introduced. In MobileNet, the standard convolution is divided into depth-wise convolution and point-wise convolution to reduce computing overhead. In ShuffleNets, the computing complexity of convolution is reduced by grouping convolution, the speed and accuracy are balanced better by the channel shuffling for channel information interaction. In GSConv (Grouped Spatial Convolution) module, there are standard Conv module, DWConv module (consist of Depth-wise Separable Convolution and point-wise convolution module), Concat module, and shuffle module (shown in Fig. 7). GSConv has the advantages of both MobileNet and ShuffleNet, and is used in this paper for both accuracy and lightweight. In Fig. 7, C1 is the number of channels for the input feature map; C2 is the number of channels for outputting feature maps (shown in formula (1)).

1

DSConv (Depth-wise Separable Convolution), which consists of two different convolutional modules (shown in Fig. 8), is performed after a SC (Standard Convolution)³⁶. The output channel number of both convolutions are half of C2. So, a new feature map can be formed by the concatenation of the two convolution results along the channel dimension. Finally, shuffle shuffling is performed, and an output tensor, reordered along the channel dimension, is returned.

The parameter number and the computational cost of DSC are much lower than that of SC. However, the information of different layers on the same spatial position is not effectively utilized. So, the feature maps, generated in depth-wise separable convolution, are weightly combined to generate a new feature maps with expanded interaction feature in the following point-wise convolution module (shown in Fig. 9), whose convolution kernel size is 1.

By the DSC, the speed of the detector can be greatly improved; By the SC, the interaction information among channels can be effectively preserved to ensure their accuracy. So, the SC, DSC, and the mixing operations are all utilized simultaneously in GSConv to light-weight without losing its accuracy. When GSConv is used in the backbone network, the layers of the backbone network are deepen, the resistance of data flow is exacerbated, the inference time is significantly increased, and some semantic information loses. So, the GSConv module is used in the neck in this paper.

SE-C2f

The channel-wise attention of SE (Squeeze-and-Excitation) attention mechanism is used to enhance the performance of convolutional neural networks. Squeeze and Excitation are two main steps of SE module (shown in Fig. 10).

In the Squeeze step, by global average pooling, each channel of the feature map is compressed into a scalar, which represents the overall global information, and reflects the comprehensive importance of each channel.

In the Excitation step, a weight vector is generated by two fully connected layers. Firstly, the dimension of scalar, obtained from the Squeeze step, is reduced; And then, the dimension is expanded by ReLU activation; Finally, the weights are limited among 0 and 1 by the Sigmoid activation function. So, the relative importance of each channel can be indicated by the weight vectors.

In SE attention mechanism, by multiplying the original feature map with this weight vector, an enhanced feature map can be obtained to increase attention, the inter-channel information can be better utilized to improve the performance, and the representation power of convolutional neural networks can be effectively enhanced. The parameter number of the SE module is so relatively small that the overall parameter number of the network module is hardly increased when it is integrated into the C2f module of the neck(the structure is shown in Figs. 11,12).

Loss function

In the field of object detection, the accuracy are directly affected by the effectiveness of bounding box loss functions, such as CloU³⁷, EloU³⁸, and SloU³⁹. The positional error between the predicted bounding box and the ground truth is optimized by the loss function for gradually approaching a minimum. In most research data-sets, the objects are annotated with high quality for bounding boxes to highly effectively improve fitting accuracy.

CloU is an earlier bounding box regression loss function, and is used in the original YOLOv8 model. However, CLoU has the following drawbacks: (1) it does not consider the balance between high-quality and low-quality samples; (2) the aspect ratio is used as a penalty factor for the loss function, and cannot effectively reflect the real aspect ratios because the predicted aspect ratios significantly differs from the real aspect ratios. The CloU formula is as follows:

where, IoU represents the Intersection Over Union between the predicted box and the true box. The remain parameters are shown in Fig. 13

Subsequently, EIoU improved CIoU by using length and width as penalty terms, which can more accurately reflect the length and width difference between the real box and the predicted box. The formula for EIoU is shown in formula (3):

Based on EloU, the angle between the predicted box and the actual box is used as a penalty factor of SIoU. So, the predicted box can quickly move towards the nearest axis, regress to the actual box, and accelerate the convergence speed.

The above loss functions all adopt a static focusing mechanism, whose effect on data-sets with high-quality images is good. However, for data-sets with large number of low-quality images, the detection performance improvement of model with the above loss function is limited. In WIoU⁴⁰, considering the slope direction, overlapping area, and centroid distance, a dynamic non-monotonic focusing mechanism, the outliers for evaluating anchor boxes, and a reasonable gradient gain strategy are used for focusing on low-quality anchor boxes to improve overall accuracy. Now, WIoU has three versions. In WIoUv1, the distance attention is constructed on distance measurement; In WIoUv2, a monotonic focusing mechanism is designed for cross entropy to focus on low-quality samples better; In WIoUv3, a small gradient gain is added furthermore to reduce harmful gradients generated by low-quality samples. The formula of WIoUv1 is defined as follows:

The formula of WIoUv2 is defined as follows:

The formula of WIoUv3 is defined as follows:

where, β is defined as the quality of the anchor box, and is assigned a small gradient gain; δ and α are adjustable hyper-parameters, which can adapt to different models.

WIoU has the advantages of both EIoU and SIoU. Therefore, in response to the adverse conditions and poor datasets quality of underwater target detection tasks, WIoU is used as the bounding box loss function to obtain more accurate detection results.

Parameter settings

The experimental environment configuration of this paper are shown in Table 1. The Pytorch version 1.12.0 + cu113 deep learning framework was used, with an Intel(R) Core (TM) (TM) i7-13620H CPU, 8G memory, NVIDIA GeForce GTX 4060 graphics cards, and Windows 11 operating system; the software programming environment was Python 3.9.

Evaluation indicators

In this paper, representative evaluation indicators, the AP and mAP, of target detection models are used. AP (Average Precision) refers to the model average recognition accuracy of each class. The AP value is approximated by smoothing the area calculation of the P-R (Precision-Recall) curve, where recall is the horizontal axis and precision is the vertical axis in the coordinate system. mAP (mean Average Precision) refers to the average AP value of all categories. The mAP can be obtained by calculating recall and accuracy metrics.

The formula for calculating the recall rate is defined as following:

11

The formula for calculating accuracy is defined as following:

where, T_P represents positive samples that have been correctly identified as positive; F_N represents positive samples that have been incorrectly identified as negative; F_P represents negative samples that have been incorrectly identified as positive.

The average accuracy (mAP) of all categories is obtained by calculating the average value of each category. The formula for calculating mAP is definde as following:

In the above equation, N represents the category number of objects, detected in the data-set. In the our experiments, N = 4, corresponding to echinus, holothurian, starfish, and scallops.

Detection results on URPC dataset

In this subsection, our proposed model is experimented on the augmented URPC2020 data-set. The initial learning rate is 0.01, the momentum is 0.9, the weight decay of optimizer is 0.0005, the epoch is 200, batch size is 4, and the other training parameters is set to the default values of YOLOv8 network. The RP curve is shown in Fig. 15.

The recall and accuracy of the model are two key performance indicators of object detection tasks. The PR curves in the above Fig. 15 demonstrate that our detection efficiencies of various target categories are much high. Especially, the average accuracy (AP) of sea urchins is 90.1%, and the calculated average accuracy (mAP) is 82.9%, Our model performs excellently.

Ablation experiments

In this section, ablation experiments are conducted to verify the effectiveness of different sub-modules. The mAP represents the detection accuracy of the algorithm. The higher the value of mAP, the better the performance of detection and recognition. In addition, the parameter size of the model is another important indicator. The fewer parameters, the faster the speed. The different modules are sequentially used to improve the YOLOv8 algorithm. The results of this ablation experiment are shown in Table 2. In the first row of the Table 2, The detection results of the original YOLOv8 network are shown as benchmark.

In the second row of Table 2, The detection results of improved YOLOv8 only with GSConv module are shown. Compared with the original YOLOv8, the parameter quantity is reduced, the accuracy remains unchanged. That is, the efficiency is improve by the depth-wise separable convolution of GSConv module. Simultaneously, the information exchanges among channels are still retained by the common convolution and channel shuffling, and the accuracy of model can be ensured by the richer learned feature representations.

In the third row of Table 2, the detection results of YOLOv8 + CBAM are shown. The global weight distribution can be learned by CBAM. So, the task key parts of the input can be more focused, and the attention of irrelevant (or noisy) information can be reduced. Therefore, the accuracy and robustness of the model can be improved, and the accuracy increases 0.9%.

In the fourth row of Table 2, the detection results of YOLOv8 + SE-C2f are shown. SE attention mechanism focuses on channels, and improves the network representation ability by explicitly modeling the interdependence among convolution feature channels. By Integrating SE attention mechanism into C2f module of neck, the feature fusion effect of neck and the feature extraction ability of the model can be enhanced, and the accuracy can be increased by 0.7% with similar parameters.

In the fifth row of Table 2, The detection results of YOLOv8 + WIoU lose function are shown. The accuracy of YOLOv8n + WIoU is 0.8% higer than that of original YOLOv8. WIoU uses dynamic non-monotonic focusing mechanism, evaluates the quality of anchor frame by using “outlier” instead of IoU, and provides a wise gradient gain allocation strategy. This strategy not only reduces the competitiveness of high-quality anchor frames, but also reduces the harmful gradient caused by low-quality examples. This allows WIoU to focus on the anchor frame of ordinary quality and improve the overall performance of the detector. It is proven that WIoU is an effective method without side effects.

In the sixth row of Table 2, the detection results of YOLOv8n + Data enhancement are shown. The accuracy of the model reaches 80.7%. It is proved that image enhancement is a very effective method to improve the accuracy of underwater targets.

In the last row of Table 2, the detection results of AGW-YOLOv8 are shown. From the table, it can be seen that the parameter number is reduced from 3.01 to 2.95M at the little cost of accuracy decrease by AGW-YOLOv8. However, the accuracy is still 82.9%, and is 2.5% higher than the benchmark model.

It is proved that our improvements are practical and effective by the above experiments. So, it can be said that our improvements is very successful, and lays a better foundation for the future research.

In addition, LCAHE-WT is compared with Wavelet Transform, and the experimental results show that the accuracy of our proposed method has been improved slightly, As shown in Table 3. This method is proved to be effective.

Comparison with other algorithms

In this subsection, our proposed algorithm is compared with YOLOv4⁴¹, YOLOv5-s⁴², YOLOX-s⁴³, YOLOv8s, YOLOv8l, improved YOLOv4¹³ and B-YOLOX-S¹⁴ on the augmented URPC2020 data-set. All experiments were performed with the same parameters in the same environment. The results are shown in Table 4:

It can be seen from the table that, compared with YOLOv4-tiny, YOLOv5-s, YOLOX-s, YOLOv8s, and the improved YOLOv4, our accuracy is 7.3%, 5.1%, 3.0%, 1.5%, and 3.4% higher respectively, and the number of parameters slightly decreases simultaneously. Compared with YOLOv8l and B-YOLOX-S, the similar precision values can be achieved at the only cost of 6.7% and 23.9% parameters. That is, the satisfactory results can be achieved by AGW-YOLOv8 at the low cost. Therefore, AGW-YOLOv8 is more suitable for underwater target detection.

Visualization analysis

The heat map is a key visualization tool of object detection based on deep learning⁴⁴. The position information of the detected object can be reflected by heat map. So, a part of an image can be selected by heat map to help neural network making the final classification decision.

The necessary information is provided by the high response areas in the heatmap to support the final discrimination of the target detection network ⁴⁵. Some images are selected from the URPC data-set, and the comparative experiments of heat map visualization is done (shown in Fig. 16). In Fig. 16, the detection capabilities of YOLOv8 and AGW_YOLOv8 for echinus, starfish, holothurian, and scallop are respectively demonstrated.

From the prominent hotspots, it can be seen that it is difficult for YOLOv8 to focus on objects of interest due to the large receptive domain of the feature information, which leads to imprecise detection and recognition, when detecting and recognizing sea urchins, starfish and scallop. Therefore, in AGW-YOLOv8, attention mechanism is used to improve the ability of feature extraction, WIoU is used increase the performance of detection and recognition by measuring the matching degree between the detection box and the actual box better.

However, there are some issues of both YOLOv8n and AGW-YOLOv8 during the detection of sea cucumbers, not only because of the limited parameters and computing resources, but also because of lots of sea cucumbers being obstructed by seaweed and coral. So, the mAP of sea cucumbers is only 69%, and is significantly lower than that of others. Therefore, researches should be done furthermore.

The performance analysis is done in three representative complex underwater scenes of data-set to prove the feasibility of AGW-YOLOv8 (shown in Figs. 17,18,19). In these figures, the detection results of YOLOv8 is in the first column; the detection results of AGW-YOLOv8 is in the second column.

In Fig. 17 the missed detection rate of the detector is relatively high when the target volume is small with weak feature expression ability. There are 10 and 24 targets respectively in the tested scenario. YOLOv8n can detect 9 and 20 targets respectively; AGW-YOLOv8 can detect 10 and 22 targets respectively. So, it is proved that, by attention mechanism, the global information can be captured more effectively, the data understanding and expression ability of model is higher, and targets of different scales and positions can be adapted better. That is, it is very effective for Attention Mechanism to detect small objects.

In Fig. 18, the sea cucumbers, aquatic plants and reefs are highly similar. The targets are missed for seaweed and reefs by detector with any carelessness. So, the recognition is much low, and the detection is more difficult. From Fig. 18, it can be seen that YOLOv8n has missed detection; AGW-YOLOv8 does not has any missed detection, and localizes accurate. So, the recognition of AGW-YOLOv8 is much better than that of YOLOv8n. That is, in complex terrain, the labeling quality of the target object may not be high. WIoU can perform well in this case, because it will not produce a large IoU loss when generating high-quality anchor frames for low-quality examples.

In Fig. 19, the seawater is turbid with limited lighting, the detection is more difficult. The two detectors loss some detection. However, there are false detection in YOLOv8n because of the overly fuzzy data. In AGW-YOLOv8, the distorted images are repaired, and the image details are restored by data augmentation. So, its detection is better. Therefore, it is important for detection accuracy improvement to make images clearer by effective data augmentation methods in the deep-sea.

Overall, by AGW-YOLOv8, the missed and false detection can be reduced, and the optimization effect is improved in complex underwater environments. Especially, the excellent accuracy and detection speed can be guaranteed under the circumstances of many small targets, complex terrain, and distorted images. That is, our work is effective.

Competing interests

The authors declare no competing interests.