This project demonstrates how to classify images of trash into different categories using deep learning techniques. The model is trained to recognize six types of waste: cardboard, glass, metal, paper, plastic, and trash, based on images. The data is sourced from the garythung/trashnet dataset hosted on Hugging Face, and the model is built using a Convolutional Neural Network (CNN) architecture with TensorFlow/Keras.
Dataset
The dataset used in this project is the TrashNet dataset, which contains labeled images of various types of trash. The dataset is divided into six categories:
Cardboard
Glass
Metal
Paper
Plastic
Trash
Each image is labeled with one of these categories, making the task a multi-class classification problem.
The dataset is loaded directly from Hugging Face using the datasets library and is split into training and testing sets. The dataset is further augmented for better generalization.
Data Preprocessing
Saving Images Locally: Images are downloaded and saved locally using a custom function. This is essential for use with ImageDataGenerator, which requires file paths to work.
Data Augmentation: Data augmentation is applied to the training data to improve the model's generalization ability. The augmentations include random rotations, width/height shifts, flips, zooming, and brightness adjustments.
Data Splitting: The dataset is split into 80% training and 20% testing using train_test_split from datasets.
Installation
To run the project, you need to install the following dependencies:
Additionally, if you're working in a Jupyter environment, use the following to ensure correct visualization:
pip install jupyter
Project Structure
bash
1/
2├── model_architecture_CNN.png # Image showing the CNN model architecture3├── trash_classifier_model.h5 # Trained deep learning model (saved after training)4└── README.md # This file
Model Architecture
The model is built using a Convolutional Neural Network (CNN), which is a type of deep learning model commonly used for image classification tasks. The architecture is as follows:
Conv2D Layers: Convolutional layers extract features from the images by applying filters. These filters help detect various patterns such as edges, textures, and shapes in the image. The filters are learned during the training process, allowing the model to automatically recognize relevant features.
MaxPooling2D Layers: These layers downsample the spatial dimensions, reducing computational complexity while retaining essential features. Max pooling helps make the model more invariant to small translations and distortions in the image.
Flatten Layer: This layer flattens the 2D output of the convolutional layers into a 1D array for classification. It prepares the features for input to the fully connected layers.
Dense Layers: These fully connected layers interpret the extracted features and make predictions based on the learned features. The more units in the dense layer, the more complex the model is.
Output Layer: The output layer has 6 units (one for each class), with a softmax activation function, which outputs probabilities for each class. The softmax function ensures that the outputs are normalized into a probability distribution.
The model is compiled with the Adam optimizer and categorical crossentropy loss, as this is a multi-class classification problem.
The model is trained for 50 epochs using a batch size of 32. The training and validation accuracy, as well as the loss values, are plotted during training to monitor the model's performance.
The model's performance is further evaluated using a confusion matrix and classification report, which provide insights into the model's accuracy, precision, recall, and F1-score for each class.
After training, the model is saved in the H5 format for later use:
model.save("trash_classifier_model.h5")
Conclusion
This project provides a robust solution for classifying trash into different categories using deep learning. The use of CNNs allows for automatic feature extraction from the images, and data augmentation improves the model's ability to generalize. The evaluation metrics ensure that the model performs well across all six classes.
