One of the major challenges in sorting waste at the source is most people are still not properly informed about waste segregation. Artificial intelligence has advanced to a level where I think it can assist humans in quickly sorting junk.
- Data
In order to recognize waste in wild, I train models on the dataset which similar to the real world scenario. Trashnet is the dataset to be used to train models. It consists of 2527 images categorize into 6 classes:
- 501 glass
- 594 paper
- 403 cardboard
- 482 plastic
- 410 metal
- 137 trash
For more details, please check out the repository.
In order to mimic the real world environment, I applied some augmentation technique like RandomBrightnessContrast, ISONoise, Blur, RandomFog with the help of albumentation.
- Models
I experiment with various Classification models:
- ConvNeXt
- ResNet 50
- ResNet DINO
- ViT
- ViT-MAE
The accuracy on test set is impressive, these models get F1 score at 9x%. But when I test it on the real images, most of the cases it failed. I'm still working on finding a better solution.
- Install dependecies
pip install -r requirements.txt
- Check the notebook for more details
Testing on the real world images, I realize that ConNeXT although the F1 score on test set isn't as good as ViT but it recognize well on these images. So I create the API with ConvNeXt as the classifier instead of ViT. The API load the ONNX model stored on Google Drive, run the code and it will automatically download the model.
For your specific needs, checkout helper.py. Change PRETRAINED_MODEL to one of the Hugging Face models, the image preprocessing is based on this. DRIVE_PATH as the link to your ONNX model, train your own and put it on Google Drive.
- Build the docker image
bash build.sh
- Run the container
bash run.sh
- Inside the container, execute the following command to run the api
bash run_api.sh