U-Garbage

A garbage sorting model using VGG16, Residual Blocks, and Convolutional Block Attention Module, implemented in PyTorch.

How to get started

Clone the repository:

git clone [email protected]:ycshao21/U-Garbage.git
cd U-Garbage

Set up the conda environment:

conda create -n garbage python=3.12
conda activate garbage

conda install numpy pandas matplotlib
conda install pillow
conda install ruamel.yaml
pip install seaborn=0.13.1

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

conda install tensorboardX

Download the dataset Garbage Classification (12 classes) by Mostafa Mohamed from Kaggle, and unzip it under data/.

Preprocess the dataset:

python utils/preprocess.py

To train the model, modify configs/train.yaml and run:

python train_vgg.py --name <name>

To test the model, modify configs/train.yaml and run:

python test.py