Using a modified DenseNet121 neural network and OpenCV, this command line application will predict the breed of a dog in an image and output the top five classes. A transfer learning approach to training was taken by replacing the fully connected layer of the neural network with one particular to the number of classes in this project. The model was trained with the convolutional layers frozen so as not to backpropagate the parameters through the network. It can handle cases for a dog, person, and not a dog. Using a haarcascade alogrithm the app will detect if there is a person in the image, then return the dog breed that the person looks most likely resembles. Enjoy!
The implementation of this app, as seen below, can be found on my personal website.
-
Install dependencies - I generally use Conda for my environment and package management.
pip install requirements.txt
-
Clone this repo -
https://github.com/justinbellucci/app_dog_breed_classifier
cd app_dog_breed_classifier
-
Main Files - located in
/backend
folderDog Breed Classifier
app.py
- runs the main application
model_functions.py
- loads checkpoint and class names
image_helper.py
- functions to help with image prep
classifier.py
- CNN Classifier class
detectors.py
- human and dog detector functions
predict.py
- prediction functions using VGG16 model and haarcascade algorithmTraining a Model
train.py
- main function to train the CNN using transfer learning
loaders.py
- pytorch dataloaders and transforms
To run the applcation from the command line simply navigate to the main repo folder and enter the following:
python app.py --img_path _my_path.jpg_
Where my_path.jpg is the path to the image you want to classify.
This application uses two CNN models for classification purposes, one of which implements a custom classifier. The CNN model was trained on a dataset of over 8000 images.
-
Dog Detector - This function uses a VGG16 pretrained model to detect if a dog is in the image. It simply outputs a class then references a class dictionary from Imagenet. VGG16 Imagenet class dictionary keys 151-268, inclusive, correspond to dog names.
-
Transfer Learning - Starting with a pretrained densenet121 CNN model a custom multi layer classifier was constructed with 120 output nodes representing the 120 class labels in the dataset. With frozen convolutional layer parameters, the model was trained on a GPU for for around 6 hours over 40 epochs. At the time of stopping the validation loss was still decreasing. The model achieved accuracy of approximatly 79% when evaluated with the test image set.
Included is the code to train you own model. You will just need a training, validation and test set of data. The data_loader()
function in the loaders.py
file uses the torchvision.datasets.ImageFolder()
class to load image data. Please take a look at https://pytorch.org/docs/stable/torchvision/datasets.html#imagefolder for specifics about naming conventions. According to the documentation the images should be aranged in this way:
root/dog/xxx.png
root/dog/xxy.png
root/dog/xxz.png
root/cat/123.png
root/cat/nsdf3.png
root/cat/asd932_.png
It is highly recommended that you use GPU. The code has support for cuda and will check accordingly. From the command line run:
python train.py --chkpt_dir /model_checkpoint.pt --epochs 40 --img_dir /dogimages
Make sure that the /dogimages
folder has the following subfolder naming:
dogimages/train
dogimages/valid
dogimages/test
Stay tuned for more additions to this projects!