Garbage classification refers to the separation of several types of different categories in accordance with the environmental impact of the use of the value of the composition of garbage components and the requirements of existing treatment methods.
The garbage classification dataset is from Kaggle. There are totally 2467 pictures in this dataset. And this model is an image classification model for this dataset. There are 6 classes for this dataset, which are cardboard (393), glass (491), metal (400), paper(584), plastic (472), and trash(127).
The model is based on the
ViT model, which is short for the Vision Transformer. It was introduced in the paper
An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, which was introduced in June 2021 by a team of researchers at Google Brain. And first released in
this repository. I trained this model with PyTorch. I think the most different thing between using the transformer to train on an image and on a text is in the tokenizing step.
I trained the model with 10 epochs, and I use Adam as the optimizer. The accuracy on the test set is 95%.
Huggingface model card is
here.
-
Next step: build a CNN model on this dataset and compare the accuracy and training time for these two models.
-
Didn’t use the Dataset package to store the image data. Want to find out how to use the Dataset package to handle image data.
