International Journal of Science, Engineering and TechnologyAuthors: Gurdarshan Singh, Shivam, Yash Singh
Abstract: Sorting waste by hand is slow, error-prone, and still one of the biggest bottlenecks in recycling streams. This paper presents WasteVision AI, an end-to-end system for fine-grained waste classification built around a Vision Transformer (ViT-Base/16), fine-tuned on the 36-class TriCascade WasteImage dataset (35,264 images) through a two-stage transfer-learning schedule. The backbone is first warmed up as a whole with a frozen classifier head under a class-weighted, label-smoothed cross-entropy loss, after which the last four transformer blocks are unfrozen and fine-tuning continues. On a held-out test set of 3,527 images, the resulting model reaches 95.66% top-1 accuracy, 98.27% top-3 accuracy, and a macro-averaged F1-score of 0.93. To explain its predictions, we adapt Grad-CAM to work on the transformer’s patch-token hidden states, producing pixel-wise heatmaps that localise the discriminative regions of each waste item; predictions below 50% confidence are flagged “Unknown Waste” rather than forced into a class. The whole pipeline is exposed through a lightweight Flask REST API that returns the predicted class, a calibrated confidence score, and a base64-encoded Grad-CAM overlay in a single call. Against a recent three-stage DP-CNN/Ensemble-ELM benchmark evaluated on the same dataset, our single-stage ViT pipeline gains over 10 percentage points of fine-grained (36-class) accuracy, though at a higher parameter cost — a trade-off worth weighing when choosing between transformer- and CNN-based deployments for waste sorting.
