Research on floating object classification algorithm based on convolutional neural network - Jikai Yang; Zihan Li; Ziyan Gu; Wei Li (Scientific Reports)
Abstract With the advancement of artificial intelligence technology, unmanned boats utilizing deep learning models have shown significant potential in water surface garbage classification. This study employs Convolutional Neural Network (CNN) to extract features of water surface floating objects and constructs the VGG16-15 model based on the VGG-16 architecture, capable of identifying 15 common types of water surface floatables. A garbage classification dataset was curated to obtain 5707 images belonging to 15 categories, which were then split into training and validation sets in a 4:1 ratio. Customized improvements were made on the base VGG-16 model, including adjusting the neural network structure to suit 15 floating object categories, applying learning rate decay and early stopping strategies for model optimization, and using data augmentation to enhance model generalization. By tweaking certain parameters, the study analyzed the impact of the number of epochs and batch sizes on the model’s classification effectiveness. The results show that the model achieves the best performance with 20 epochs and a batch size of 64, reaching a recognition accuracy of 93.86%. This is a 10.09% improvement over the traditional VGG-16 model and a 4.91% increase compared to the model without data augmentation, demonstrating the effectiveness of model improvements and data augmentation in enhancing image recognition capabilities. Additionally, the few-shot test demonstrates the fine-tuned model’s improved generalization capability. This research illustrates the applicability of transfer learning in the task of water surface garbage classification and provides technical support for the application of unmanned boats in environmental protection.