This project demonstrates simple object recognition using a Convolutional Neural Network (CNN) trained on the CIFAR-10 dataset with PyTorch. It also supports real-time object classification using your webcam.
A Convolutional Neural Network (CNN) is a type of deep learning model especially effective for image recognition tasks. CNNs use layers of filters (convolutions) to automatically learn features from images, such as edges, shapes, and objects, making them well-suited for visual data.
- Trains a CNN on CIFAR-10 dataset (10 classes)
- Evaluates model accuracy
- Real-time webcam object recognition
- Demo class labels include: plane, car, bird, cat, deer, dog, frog, horse, ship, truck, mobile phone, pen
- Python 3.8+
- torch
- torchvision
- opencv-python
- mlflow
- Install dependencies:
pip install torch torchvision opencv-python mlflow
- Run the script:
python simple_object_recognition.py
- After training, a webcam window will open. Show an object to the camera to see the predicted class. Press 'q' to quit.
This project uses MLflow to track experiments, log hyperparameters, metrics, and save trained models.
- After running the script, experiment data will be logged automatically.
- To view and compare experiment runs, launch the MLflow UI:
mlflow ui
- Open http://localhost:5000 in your browser to explore runs, metrics, and models.
You can find the logged models and metrics in the mlruns directory created in your project folder.
- The model is only trained on CIFAR-10 classes. The accuracy is quite low, 68% and some classes like 'mobile phone' and 'pen' are not in the training set.
MIT