Pruning-Based YOLOv4 Algorithm for Underwater Gabage Detection

摘要

To tackle the problem of aquatic environment pollution, a vision-based autonomous underwater garbage cleaning robot has been developed in our laboratory. This paper proposes a garbage detection method based on a modified YOLOv4, allowing high-speed and high-precision object detection. More specifically, the YOLOv4 algorithm is chosen as a basic neural network framework to perform object detection. With the purpose of further improvement on the detection speed, the channel pruning and layer pruning are implemented on the trained YOLOv4 model, while the fine-tuned mechanism assists the pruned model to restore accuracy. In virtue of the improved detection methods, the robot has the ability to collect garbage autonomously. The experimental results indicate that the pruned YOLOv4 detection method can still maintain the high performance even though the parameter amount is 7.062% of the original model.