This project provides a script to train a convolutional neural network (CNN) on a dataset of images. The script is written in Python and uses TensorFlow and OpenCV. The model can classify images into different classes based on the dataset provided.
- Tensorflow
- Tensorflow-gpu
- OpenCV-python
- Matplotlib
- Clone or download the repository.
- Make sure that you have the required libraries installed.
- Place your dataset in a directory called 'training_data'
- Open the Jupyter notebook by running the command "jupyter notebook" in your terminal in the directory.
- Locate and open the training_model.ipynb notebook
- Make sure that the kernel for the notebook is set to Python 3.
- Run the cells in the notebook to train the model on your dataset.
- The script will create a unique folder with the current date and time in the format dd.mm.yyyy hhmmss-train_results for storing the model files and training results.
- To use this script, you will need to place your dataset of images in a directory called 'training_data'. The script expects the images to be organized into different subfolders within the 'training_data' directory, one for each class. For example, if you have two classes, "class1" and "class2", you should create two subfolders within the 'training_data' directory named "class1" and "class2" and place the corresponding images in each folder. Make sure that the images are in jpeg, jpg, bmp, or png format.
training_data/ class1/ image1.jpg image2.jpg ... class2/ image1.jpg image2.jpg ...
The model architecture is a convolutional neural network (CNN) with the following layers:
- Conv2D with 16 filters of size (3,3), stride of 1, and ReLU activation
- MaxPooling2D
- Conv2D with 32 filters of size (3,3), stride of 1, and ReLU activation
- MaxPooling2D
- Conv2D with 16 filters of size (3,3), stride of 1, and ReLU activation
- MaxPooling2D
- Flatten layer
- Dense layer with 128 units and ReLU activation
- Dense layer with the number of units equal to the number of classes in the dataset and softmax activation
You can change the architecture of the model, training parameters, and data set path accordingly.
The script will save the following files in the unique folder created after training:
- The trained model file in the format model.h5
- A plot figure of a batch of the training images and their labels in the format batch_result.jpg
- An Accuracy and a Loss chart as accuracy_chart.jpg and loss_chart.jpg
- A file with the training and validation accuracy and loss values in the format histroy.json
- Final summary and result as .txt
The following images are examples of the results generated by the script after training.
- This script is set to work with images of size 256x256, you can change the image size accordingly.
- The script automatically loads the images and splits them into train, validation, and test sets.
If you want to contribute to this project, please feel free to submit a pull request.
This project is licensed under the MIT License - see the LICENSE.md file for details.
The code for this project is greatly inspired by Nicholas Renotte.
You can find an example of a pre-trained model in the model folder, which was trained to classify images as notes and non-notes from WhatsApp images. You can use this pre-trained model as a starting point for fine-tuning on your own dataset or as a reference for comparison to your own model's performance. The model is trained using two classes (notes and non-notes) and is available in this repository NotesSeparator.