This project uses a deep learning model to identify if an image contains Frank the dog or not. The model is based on the VGG16 architecture and uses transfer learning for binary classification.
- Python 3.7+
- PyTorch
- torchvision
- OpenCV
- Pillow
git clone https://github.com/jbfly/frank_or_not.git
cd frank_or_notYou can install the required dependencies using pip:
pip install torch torchvision opencv-python pillowTo run the model on your local machine using your webcam, make sure you have the model file best_frank_or_not_model_run_1.pth in the same directory as frank_or_not.py. Then execute:
python frank_or_not.pyThis will open a window showing the webcam feed with the prediction overlaid on the video.
If you need to retrain the model, you can use the train_and_evaluate.py script. The training script is set up to handle multiple training runs and implements early stopping.
Ensure your dataset is organized as follows:
dataset/
train/
frank/
not_frank/
valid/
frank/
not_frank/
Then, execute:
python train_and_evaluate.pyThis script will perform multiple training runs, save the best model based on validation accuracy, and store it as best_frank_or_not_model_run_X.pth.
The frank_or_not.py script performs real-time inference using your webcam. Make sure the model file is in the same directory as the script and run:
python frank_or_not.pyThe script will display a window with the webcam feed and overlay the prediction ("Frank" or "Not Frank") on each frame.
- Fork the repository
- Create your feature branch (
git checkout -b feature/fooBar) - Commit your changes (
git commit -am 'Add some fooBar') - Push to the branch (
git push origin feature/fooBar) - Create a new Pull Request
This project is licensed under the MIT License - see the LICENSE file for details.