A Novel Ensemble Learning Model Using CNNs And Vision Transformers for Waste Classification
Mohamed Lamine Cisse, David Adeyinka Aderinwale, Melike Şah · 2024
Automatic waste classification is challenging due to the variety and complexity of waste types. Automatic waste classification with deep learning can enable faster and more accurate sorting for waste management. In this work, we focus on this problem and propose a novel ensemble approach by applying transfer learning on various pre-trained model combinations of Convolutional Neural Networks (CNNs) and Vision Transformers (ViT) for precise classification of waste images. We used the six-class Trashnet dataset for evaluations. Among the tested pre-trained CNN models, DenseNet121 and InceptionV3 achieved the best accuracies of $\mathbf{9 5. 2 5 \%}$ and $\mathbf{9 3. 4 6 \%}$ respectively on the Trashnet dataset. For the vision transformer, ViT_b16, and ViT_b32 achieved accuracies of $\mathbf{9 5. 4 5 \%}$, and $\mathbf{9 3. 4 7 \%}$ respectively. We call these models base models. Results analysis on base models showed that specific models learned certain features of the dataset better for different waste class types. Then, we applied two types of ensemble approaches to combine the predictions of base learners: Weighted averaging and model averaging. Weighted averaging assigns different importance to each model’s predictions based on their accuracy, while model averaging combines all models’ predictions equally. Furthermore, after checking the overall performance of those models on the test set, it was noticed that the model averaging ensemble approach had $97,92 \%$ accuracy, and the weighted averaging ensemble approach had $\mathbf{9 8}, \mathbf{2 2} \%$ accuracy. Base models performed worse than the weighted average ensemble approach and the model averaging ensemble approach showing that the ensemble approach improved the accuracy by up to 4%. Our contribution is developing a novel ensemble method based on both ViTs and CNNs for the first time on the TrashNet that demonstrates state-of-the-art performance.