Optimized YOLOv8n-TBN identification method for water surface floating garbage detection

Abstract

The invention discloses an optimized YOLOv8n-TBN identification method for water surface floating garbage detection, and relates to the technical field of computer vision and environment monitoring. Comprising the steps of obtaining and preprocessing a water surface floating garbage image data set; constructing an improved YOLOv8n target detection model, wherein the improvement comprises the following steps: embedding a lightweight Triplet attention mechanism in a C2f module of a backbone network to enhance cross-dimensional feature interaction; a bidirectional feature pyramid network BiFPN is adopted to replace an original neck network so as to realize efficient multi-scale feature fusion; the NWD-CIoU loss function is used for replacing the original CIoU loss function to improve the positioning precision of the small target; training and optimizing the model by using the training set; and inputting a to-be-detected image into the trained model, and outputting a bounding box and confidence of the garbage target. According to the method, the detection precision and robustness of small target garbage in a complex water surface environment are effectively improved while the model parameter quantity is remarkably reduced, and the method is suitable for real-time monitoring scenes with limited resources.

Description

Technical Field

[0001] This invention relates to the fields of computer vision and environmental protection technology, specifically to a YOLOv8n-TBN recognition method optimized for detecting floating debris on water surfaces. Background Technology

[0002] With the increasing severity of global water pollution, efficient and real-time monitoring of floating debris on the surface of rivers, lakes, and oceans has become a crucial aspect of environmental protection and governance. Traditional manual inspection methods suffer from inefficiency, high costs, limited coverage, and certain safety risks.

[0003] In recent years, deep learning-based object detection algorithms, especially the YOLO series models, have been attempted to be applied to surface debris detection due to their speed and high accuracy. However, surface debris detection faces several unique challenges: First, debris targets are usually small in size and occupy a small proportion of the image; second, the surface environment is complex and variable, with strong interference from waves, reflections, and shadows; third, debris comes in various shapes and colors and may be partially submerged or clustered. These factors lead to problems when directly applying general object detection models, such as high false negative rates for small targets, weak resistance to interference from complex backgrounds, and high computational costs that hinder deployment.

[0004] In existing technologies, improvements for small object detection often focus on designing more complex feature pyramids or introducing attention mechanisms. However, these improvements are often accompanied by a significant increase in model parameters and computational load, making it difficult to achieve real-time detection on embedded devices or mobile devices with limited computing resources.

[0005] Therefore, how to significantly improve the detection accuracy and robustness of small-target debris on the water surface while maintaining or even reducing model complexity is a technical problem that urgently needs to be solved in this field. Summary of the Invention

[0006] The purpose of this invention is to provide a YOLOv8n-TBN recognition method optimized for detecting floating debris on the water surface. By performing triple collaborative improvements on the YOLOv8n model, the detection accuracy and anti-interference ability of small target debris on the water surface are significantly improved while achieving model lightweighting.

[0007] Therefore, this invention provides a YOLOv8n-TBN identification method optimized for detecting floating debris on the water surface, comprising the following steps: S1. Data preparation: Obtain a dataset of images of floating garbage on the water surface, clean and label the images, unify the category as "garbage", and divide the dataset into training set, validation set and test set; S2. Model Building: Construct an improved YOLOv8n object detection model; S3. Model Training: The improved YOLOv8n-TBN model is trained using the training set, and the model is validated and hyperparameters are tuned using the validation set. S4. Model Application: Input the water surface image to be detected into the trained YOLOv8n-TBN model, and output the detection results of floating garbage on the water surface. The detection results include the target bounding box position and class confidence.

[0008] Preferably, in step S1, the data in the dataset of floating garbage images on the water surface is obtained from the PaddlePaddle AiStudio platform and the public dataset FloW-Img; The ratio of training set, validation set and test set is 7:2:1.

[0009] Preferably, in step S2, the specific steps for constructing the improved YOLOv8n target detection model are as follows: S2.1 Introduce a lightweight Triplet attention mechanism into the C2f module of the model backbone network to construct the C2f-Triplet module, which is used to enhance the model's ability to extract cross-dimensional interaction information. S2.2 Replace the original neck network with the Bi-directional Feature Pyramid Network (BiFPN) module to achieve efficient cross-scale feature fusion and reduce computational redundancy; S2.3 Replace the original CIoU loss function with the NWD-CIoU loss function, and optimize the localization accuracy of small targets by introducing the normalized Wasserstein distance.

[0010] Preferably, in step S2.1, the C2f-Triplet module is implemented by embedding a TripletAttention network into the original C2f module; The Triplet Attention network comprises three branches: the first and second branches rotate the input features and then generate cross-dimensional attention weights through Z-Pool pooling, convolutional layers, and the Sigmoid activation function; the third branch is used to calculate spatial attention weights; finally, the output weights of the three branches are averaged and aggregated, and the input features are weighted accordingly.

[0011] Preferably, in step S2.2, a multi-input edge node and adaptive weighting mechanism are employed, and its workflow is as follows: Remove nodes with only one input edge and add skip connections between input and output nodes at the same level to fuse more feature information; A fast normalization fusion strategy is introduced, which uses learnable weights to perform weighted fusion of features at different scales; The formula for calculating weighted feature fusion is: In the formula For output features, For input features, and For the corresponding weights, To prevent the value from becoming an unstable minimum.

[0012] Preferably, the NWD-CIoU loss function in S2.3 is constructed by modeling the bounding box as a two-dimensional Gaussian distribution, calculating the normalized Wasserstein distance between the Gaussian distributions corresponding to the predicted box and the ground truth box, and combining it with the CIoU loss. Its loss function expression is as follows: In the formula and These are the weighting coefficients. It is a normalized exponential form based on the Wasserstein distance. For intersection, union, and comparison, This is a loss for CIoU.

[0013] Preferably, the input image size for model training is set to 640×640, the SGD optimizer is used, the initial learning rate is set to 0.01, and the training epochs are 100.

[0014] The present invention proposes a YOLOv8n-TBN identification method optimized for detecting floating debris on the water surface, the advantages of which are as follows: This invention introduces the C2f-Triplet module, enabling the model to establish a fine interaction between channels and spatial dimensions, enhancing its ability to extract discriminative features of small debris from complex water backgrounds and effectively suppressing background interference. Combined with the adaptive multi-scale feature fusion provided by the BiFPN structure, the model's ability to represent targets of different sizes is strengthened. The introduction of the NWD-CIoU loss function optimizes bounding box regression in a way that is more sensitive to small geometric deviations, directly improving the localization accuracy of small targets. The Triplet attention mechanism employed is a lightweight design that adds almost no extra parameters. The BiFPN structure enhances feature flow while simplifying the network path and achieving more efficient computation by removing redundant nodes and introducing fast normalization fusion. This makes the model of this invention very suitable for deployment on embedded devices, drones, or unmanned vessels with limited computing resources, enabling real-time, online monitoring of surface debris. All three improvements are designed to address the specific challenges of water surface monitoring. Triplet attention enhances the model's ability to focus on targets in dynamic water surface contexts; BiFPN improves the model's adaptability to debris at different distances and scales; and NWD-CIoU loss improves the model's localization stability when targets are partially occluded or slightly deformed, making it more stable and having more complete detection capabilities than the original model.

[0015] This invention effectively solves the core contradictions of existing methods for detecting floating debris on water surfaces—namely, the difficulty in detecting small targets, significant environmental interference, and high requirements for model deployment—through a collaborative and innovative architectural improvement. It provides an efficient, accurate, and practical technical solution for intelligent water surface cleaning and environmental monitoring. Attached Figure Description

[0016] To more clearly illustrate the technical solutions in the embodiments of this application or related technologies, the accompanying drawings used in the description of the embodiments or related technologies will be briefly introduced below. Obviously, for those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0017] Figure 1 This is a diagram of the improved YOLOv8n network model structure of the present invention; Figure 2 This is a diagram of the TripletAttention network structure of the present invention; Figure 3 This is a diagram of the TripletAttention branch network structure of the present invention; Figure 4 This is a structural diagram of the Bidirectional Feature Pyramid Network (BiFPN) of this invention; Figure 5 This invention provides a visual comparison chart of the detection results of the improved algorithm and the original algorithm. Detailed Implementation

[0018] To make the objectives, technical solutions, and beneficial effects of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings and specific embodiments. It should be understood that the specific embodiments described in this specification are merely for explaining the invention and are not intended to limit the invention.

[0019] Example: Please refer to Figure 1-5 This invention provides a YOLOv8n-TBN identification method optimized for detecting floating debris on water surfaces, comprising the following steps: Images containing clearly visible floating debris, such as plastic bottles, plastic bags, foam, and twigs, were selected from the PaddlePaddle AI Studio open platform and the publicly available FloW-Img dataset for surface debris. Images were manually inspected and removed if they were blurry, contained extremely small and difficult-to-identify debris, or primarily consisted of underwater debris to ensure data quality. Using annotation tools such as LabelImg, all visible floating debris in the images was uniformly labeled as "litter" in YOLO format. The processed 3550 images were then randomly divided into a training set of 2485 images, a validation set of 710 images, and a test set of 355 images in a 7:2:1 ratio.

[0020] Implement the C2f-Triplet module: In the module definition file of the code, create a new TripletAttention class. Its forward function strictly follows... Figure 2 and Figure 3 The structure is as follows: For the input tensor x, branch one performs a permute(0, 2, 1, 3) transformation to (H, C, W) dimensions, then concatenates the results of max pooling and average pooling along the Z-Pool channel dimension, applies a 7x7 convolutional layer while preserving the dimensions, activates a sigmoid function, and then transforms it back to the original dimension. Branch two performs a permute(0, 3, 2, 1) transformation to (W, H, C) dimensions and performs the same operation. Branch three directly performs max pooling and average pooling on x along the channel dimension, concatenates the results, and then generates a spatial attention map through a 7x7 convolution and a sigmoid function. The final output is the product of the element-wise average of the three attention maps and the input x.

[0021] Create a class C2f_Triplet, which inherits from the original C2f class. In its forward method, after processing the features through the original bottleneck structure, the output features are fed into the TripletAttention module mentioned above for further operations. Modify the model structure configuration file to replace all calls to the C2f module in the backbone network with calls to C2f_Triplet.

[0022] Next, remove all layer definitions related to the path aggregation network from the Neck section of the original configuration file, according to the appendix. Figure 4 Rewrite the Neck section based on the existing structure. This section needs to explicitly receive multi-scale feature maps from the backbone network. Design bidirectional connection paths: for each layer of features, concatenate with the upper layer features through upsampling and with the lower layer features through downsampling. Implement a FastNormalizedFusion custom layer whose function is to... Learnable weighted fusion of multiple input feature maps is performed. In the configuration file, the Concat operator and the FastNormalizedFusion layer are used to construct the complete BiFPN data stream.

[0023] Locate the loss calculation file and add a method `wasserstein_distance` to the class to calculate the Wasserstein distance between the Gaussian distributions corresponding to two bounding boxes. Modify the forward function of the loss calculation. While calculating the traditional CIoU loss, calculate the Wasserstein distance between the predicted box and the ground truth box, and use this distance to calculate the normalized Wasserstein distance (NWD).

[0024] Model training and hyperparameter tuning were performed on a server running Ubuntu 20.04, using a single NVIDIA RTX 3090 GPU (24GB VRAM). Python 3.8, PyTorch 1.12.1, CUDA 11.3, and the corresponding cuDNN library were configured. After writing training scripts or directly using the modified YOLOv8 training command to start training, the framework outputs the loss value for each training epoch and calculates metrics such as precision, recall, mean precision, and mean sum on the validation set. The best-performing model weights during training are automatically saved as the best.pt file.

[0025] After training, the saved best.pt weight file is used for final evaluation on a separate test set (355 images). The evaluation command calculates and outputs a comprehensive metric report, including but not limited to: precision (P), recall (R), [email protected], number of model parameters, and GFLOPs. For practical use, the best.pt model can be exported to ONNX or TensorRT format using the official tools to optimize inference speed. During deployment, the model is integrated into the software of a surface monitoring system (such as one deployed on an unmanned vessel or a fixed camera on the shore). The system captures video streams in real time, extracts images frame by frame and scales them to 640x640, inputs them into the model for forward inference, and parses the output bounding box coordinates, class, and confidence. Finally, the detection results (such as highlighting trash locations with rectangular boxes) are overlaid on the video in real time, or alarms and coordinate recordings are triggered for subsequent cleanup operations.

[0026] To verify the effectiveness of the method of this invention, an ablation experiment was designed. Using the original YOLOv8n as a baseline, the results showed that adding C2f-Triplet improved mAP@0.5 by 1.4%; subsequently, introducing BiFPN further improved mAP@0.5 by 2.9%, and the number of parameters decreased significantly to 1.75M; finally, adding NWD-CIoU resulted in P, R, and mAP@0.5 reaching 88.8%, 80.4%, and 85.9%, respectively. The experiment demonstrates that the synergistic effect of each module significantly improves detection accuracy and achieves model lightweighting.

[0027] See Figure 5 The model of this invention was compared with current mainstream models (such as SSD, Faster R-CNN, YOLOv5n, YOLOv7-Tiny, etc.). On the same water surface debris dataset, the proposed algorithm outperformed all the compared models with an mAP@0.5 (85.9%). Compared to the original YOLOv8n, the accuracy was improved by 4.7% while the computational cost (GFLOPs) decreased from 7.5G to 6.5G. The proposed algorithm showed a significantly lower false negative rate than the original model in complex environments and scenes with dense small targets, demonstrating extremely high practical deployment value.

[0028] The above description is only a preferred embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any equivalent substitutions or modifications made by those skilled in the art within the scope of the technology disclosed in the present invention, based on the technical solution and inventive concept of the present invention, should be covered within the scope of protection of the present invention.

Claims

1. A YOLOv8n-TBN recognition method optimized for detecting floating debris on water surfaces, characterized in that: Includes the following steps: S1. Data preparation: Obtain a dataset of images of floating garbage on the water surface, clean and label the images, unify the category as "garbage", and divide the dataset into training set, validation set and test set; S2. Model Building: Construct an improved YOLOv8n object detection model; S3. Model Training: The improved YOLOv8n-TBN model is trained using the training set, and the model is validated and hyperparameters are tuned using the validation set. S4. Model Application: Input the water surface image to be detected into the trained YOLOv8n-TBN model, and output the detection results of floating garbage on the water surface. The detection results include the target bounding box position and class confidence.

2. The YOLOv8n-TBN identification method optimized for detecting floating debris on water surface according to claim 1, characterized in that: In step S1, the data for obtaining the floating debris image dataset comes from the PaddlePaddle AI Studio platform and the public dataset FloW-Img. The ratio of training set, validation set and test set is 7:2:

1.

3. The YOLOv8n-TBN identification method optimized for detecting floating debris on water surface according to claim 1, characterized in that: In step S2, the specific steps for constructing the improved YOLOv8n object detection model are as follows: S2.1 Introduce a lightweight Triplet attention mechanism into the C2f module of the model backbone network to construct the C2f-Triplet module, which is used to enhance the model's ability to extract cross-dimensional interaction information. S2.2 Replace the original neck network with the Bi-directional Feature Pyramid Network (BiFPN) module to achieve efficient cross-scale feature fusion and reduce computational redundancy; S2.3 Replace the original CIoU loss function with the NWD-CIoU loss function, and optimize the localization accuracy of small targets by introducing the normalized Wasserstein distance.

4. The YOLOv8n-TBN identification method optimized for detecting floating debris on water surface according to claim 3, characterized in that: In step S2.1, the C2f-Triplet module is implemented by embedding a TripletAttention network into the original C2f module; The Triplet Attention network includes three branches: the first and second branches rotate and transform the input features, and then generate cross-dimensional attention weights through Z-Pool pooling, convolutional layers and sigmoid activation function. The third branch is used to calculate spatial attention weights; Finally, the output weights of the three branches are averaged and aggregated, and the input features are weighted.

5. The YOLOv8n-TBN identification method optimized for detecting floating debris on water surface according to claim 3, characterized in that: In step S2.2, a multi-input edge node and adaptive weighting mechanism are used, and its workflow is as follows: Remove nodes with only one input edge and add skip connections between input and output nodes at the same level to fuse more feature information; A fast normalization fusion strategy is introduced, which uses learnable weights to perform weighted fusion of features at different scales; The formula for calculating weighted feature fusion is: In the formula For output features, As input features, and For the corresponding weights, To prevent the value from becoming an unstable minimum.

6. The YOLOv8n-TBN identification method optimized for detecting floating debris on water surface according to claim 3, characterized in that: The NWD-CIoU loss function in S2.3 is constructed by modeling the bounding box as a two-dimensional Gaussian distribution, calculating the normalized Wasserstein distance between the Gaussian distributions corresponding to the predicted box and the ground truth box, and combining it with the CIoU loss. Its loss function expression is as follows: In the formula and These are the weighting coefficients. It is a normalized exponential form based on the Wasserstein distance. For intersection, union, and comparison, This is a loss for CIoU.

7. The YOLOv8n-TBN identification method optimized for detecting floating debris on water surface according to claim 1, characterized in that: The input image size for model training was set to 640×640, the SGD optimizer was used, the initial learning rate was set to 0.01, and the training epochs were 100.

Images