Many objects in TACO dataset have Unknown label. This waste is mostly invisible or destroyed. To address this challenge at the early stage of the project we trained classifier to this type of waste to know their true category.

Additionally during our project we realized that existing datasets do not provide a large number of object classes with sufficient annotated training data. In addition, as we managed to find out, differentiating waste instances under a single class label is also challenging. In this regard, we decided to formulate our problem as a one-class object detection, and classification in next step.

A PyTorch script for litter classification:

1. ResNet - based on implementation from Tutorial on training ResNet,
2. EfficientNet - EfficientNet implementation with pseudo-labeling technique. (implemented using PyTorch Lightning)

Additionally to address class imbalance we used WeightedRandomSampler.

pip install -r requirements.txt

To track logs (for example training loss) we used neptune.ai. If you are interested in logging your experiments there, you should create account on the platform and create new project. Then:

• Find and set Neptune API token on your system as environment variable (your NEPTUNE_API_TOKEN should be added to ~./bashrc)
• Add your project_qualified_name name in the train_<net_name>.py

    neptune.init(project_qualified_name = 'YOUR_PROJECT_NAME/detect-waste')

  Currently it is set to private detect-waste neptune space.
• install neptun-client library

    pip install neptune-client

To run experiments with neptune simply add --neptune flag during launch train.py.

For more check LINK.

• we used TACO dataset with additional annotated data from detect-waste,
• we used few waste detection/segmentation dataset mentioned in main README.md,
• we used TrashNet and waste_pictures, and some pictures collected using Google Images Download.

We expect the images directory structure to be the following:

path/to/images/          # all images
  images_square/         
    pseudolabel/         # unlabeled data used in pseudo-labeling task
    test/                # images divided into categories - test subset
      background/
      bio/
      glass/
      metals_and_plastic/
      non_recyclable/
      other/
      paper/
      unknown/
    train/                # images divided into categories - train subset
      background/
      bio/
      glass/
      metals_and_plastic/
      non_recyclable/
      other/
      paper/
      unknown/

• Modified ResNet50

  Backbone of classificator is ResNet50 taken from torchvision storage.

  To run test

  python train_resnet.py --data_img path/to/images/images_square/test/ \
                         --out path/to/checkpoints/ \
                         --mode test \
                         --name test.jpg \
                         --device cpu \

  To run training (on GPU id=1)

  python train_resnet.py --data_img path/to/images/train/ \
                         --out path/to/checkpoints/ \
                         --mode train \
                         --device cuda:1 \

• EfficientNet

  This implementation uses lukemelas/EfficientNet-PyTorch implementation. EfficientNet is implemented with PyTorch Lightning.

  To run training use train_effnet.py

  python train_effnet.py --data_img path/to/images/train/ \
                         --save path/to/checkpoint.ckpt \
                         --model efficientnet-b2 \
                         --gpu 1 \
                         --pseudolabel_mode per-batch \
                         --neptune \

  • efficientnet - any efficientnet form b0 to b7 can be used
  • pseudolabeling - allow to use unannotated data. Can be use in per-epoch and per-batch modes, which refers to how often pseudolabels will be upadted
  • data augmentation - any data augmentation form albumentations can be applied

• 8 classes - 8th class for additional background category
• we provided 2 methods to update pseudo-labels: per batch and per epoch