Embedded Files
Architectures
Architectures
We chose to approach our problem statement with three different image classification models; VGG-16, DenseNet201, ResNet50. For all three models, we used the TrashNet Dataset, otherwise known as the Garbage Classification Dataset compiled by Gary Thung. This dataset is popular among garbage classification tasks as it has 6 different classes with plenty of images. Gary Thung is a computer science student at Stanford who made this dataset for a computer science class and it has gone on to be used by thousands of different developers.
Dataset can be found here: TrashNet
There are about 2527 images in total with the following image counts per class:
501 images of Glass
594 images of Paper
403 images of Cardboard
482 images of Plastic
410 images of Metal
137 images of Trash
The pictures were all taken with a white board as a background with similar lighting conditions. The images were also all of the same resolution: 512x384. The dataset is meant to be compact and small with robust data. We decided to use this dataset as it was quite popular and we originally wanted to create a top down structure for the camera in our solution which is the angle the photos are taken in this dataset. The dataset is very easy to use with every class being a labelled folder with the images within it making it ideal for creating image data generators to feed models the information.
We also looked at other datasets like TACO, for images of garbage in the wild, and also a waste classification dataset with images scraped from google search. There is a comprehensive list of datasets along with reviews for different forms of waste classification at AgaMiko's GitHub repository: https://github.com/AgaMiko/waste-datasets-review
In order to utilize this dataset we split the overall directory into three sets, a training set, validation set and test set.
VGG-16
VGG-16
VGG-16 is a CNN comprised of 21 total layers and 16 learnable layers. We chose this as one of our potential models as it uses small convolutional filters in order to increase layer depth such that memory and training time increases are neutralized to an extent. Most previous architectures traditionally use larger convolutions.
[2] VGG-16 Architecture
Loss
Loss
In contrast to the above graphs we can see how the training curve is quite smooth and has a steady decrease and improvement. On the other hand the validation loss is fairly high and is also not very smooth as it has an almost zig zag shape. This means that while the training loss got better the validation loss was suffering.
Accuracy
Accuracy
Similarly to the loss graph, the training accuracy rose to almost 100%, however the validation loss started to stagnate at around 80% and quite early as well. This probably means that while it was still improving on the training data, its actual performance was not entirely improving.
DenseNet201
DenseNet201
The Dense Convolutional Network, also known as DenseNet is a Feed-Forward CNN model. This model is another good candidate because it uses multiple Dense Blocks that contain many Dense layers which strengthen the feature propagation and enhance the accuracy. While this model has more overall node connections, it has less layers and so it decreases the distance between the input and output.
[3] DenseNet201 Architecture
Loss
Loss
From this graph we can clearly see that the training loss has a very smooth curve and decreases at a good rate throughout the models training. On the other hand we can see how the validation loss did not really decrease past 0.4 despite the training loss going much lower down.
Accuracy
Accuracy
The training accuracy also increases quite steadily throughout the training and ends at an accuracy above 95% which is very good. We also see the validation accuracy grow stagnant after around 90% which means that it continued to train but its performance was not really increasing.
RESNET50
RESNET50
ResNet50 is a CNN comprised of 50 layers with 48 convolutional layers, one average pool layer, and one MaxPool layer. It uses the bottleneck design for the building block. [5] Residual Networks have many variants that run on the same concept but have different numbers of layers. Identity shortcut connection can skip one or more layers which helps retain information in deep layers.
[4] ResNet50 Architecture
Loss
Loss
From this graph we can clearly see that the model did not perform too well. Despite RESNET being the most popular of our 3 models it performed quite poorly. We believe that there may have either been an error in our implementation that caused such results but more investigation must be done to be certain. The training loss and validation loss decrease together however they are quite high at around 1.5 at the end of training which is not good.
Accuracy
Accuracy
The training accuracy and validation accuracy also increase together and similarly to the loss its not very smooth and extremely sporadic. The ending accuracy was around 38% at the end of training which is poor for RESNET.
Evaluation
Evaluation
RESNET50
RESNET50
VGG-16
VGG-16
DenseNet201
DenseNet201
These heatmaps plot the confusion matrices when evaluating all three models. A confusion matrix is a special kind of contingency table, with two dimensions ("actual" and "predicted"), and identical sets of "classes" in both dimensions (each combination of dimension and class is a variable in the contingency table). As we can see from the data, both VGG-16 and DenseNet have fairly accurate predictions with a few outliers whereas RESNET50 on the other hand has very varied results with a lot of mistakes in its predictions over all the classes. After looking at these graphs we were certain that we would use either VGG or DenseNet for our final model. The heatmaps definitely helped us get an idea of what classes were being predicted the most and were tested most by the test dataset which was 'Paper' in this case. We also found that a lot of other non-paper objects were being classified as paper inaccurately which was a problem for us and we believe that the paper class may have not been too accurate despite a white cell for it on the heatmap.
While these were the final testing results for our model it looks like our model did not perform with these kind of metrics when testing in a real life setting. We believe this is due to a multitude of reasons:
Our model was trained on a dataset that comprised of images that were taken from a specific angle with a plain white background. This angle and background were not replicated completely and so as a result we believe that our hardware testing had a lower accuracy overall. We believe this can be solved by using a more comprehensive database, for example, the TACO dataset has a lot of images of trash in the wild and continues to grow due to its crowd-sourcing nature. This will help us as we want this solution to be readily available without the need for them to have a plain white background for the device.
The dataset we trained our model on had 2527 images which is about a couple hundred images per class. This is not much data for such a complicated prediction as classifying trash can often be difficult due to the fact that the object can be disfigured or contaminated. A more robust dataset with more diverse images may help with the accuracy of the model and the performance of it in real-time predictions on the hardware.
Our models can be found here:
Codebase
Codebase
Page updated
Report abuse