The goal of this repository is to implement a number detection using Tensorflow with a custom Convolution Neural Net (CNN) architecture specifically for fast inference. The CNN will be trained using a custom created dataset that contains numbers from 1-9 and a category for 'not numbers (-1)'.
The CNN is fast enough to run in real-time on a Raspberry Pi.
- Image size (input to the model): 320 x 240
- Camera fps: 38
- Processing time (with isolator): 25ms - 45ms (on Raspberry Pi)
- DataExtractor (to extract data for the dataset)
- Isolator (to find the numbers inside an image)
- Trainer (to train the model)
- Tester (to test the model)
You need to have the following packages installed (check requirements.txt
):
- Python 3.6
- Tensorflow 1.4.1+
- OpenCV 4.0
- Etc.
Clone the repo and install 3rd-party libraries
$ git clone https://github.com/FabianGroeger96/cnn-number-detection
$ cd cnn-number-detection
$ pip3 install -r requirements.txt
- Create a folder named
images_to_extract
in the data extractor directory (The folder can be named differently, but don't forget to change theINPUT_DATA_DIR
variable in theconstants.py
file). This directory will be used to extract the regions of interest to train your CNN. - Copy the data to extract the regions of interest into the
images_to_extract
folder - Specify which categories you want to extract in the
constants.py
file, by changing theCATEGORIES
variable - Run the
extract_data.py
file and call the methodextract_data()
from theExtractor
- After the method is finished your extracted regions of interest are located in the
data_extracted
folder. In there you will also find folders for each of your categories. These folders are used to label the regions of interest for then training your CNN.
- First of all you will have to extract the regions of interest with the DataExtractor (follow the step Extract data with DataExtractor)
- Classify the images, by dragging them in the corresponding category folder
- First of all you will have to extract the regions of interest with the DataExtractor (follow the step Extract data with DataExtractor)
- Specify in the
constants.py
file where your model will be located, by modifying theMODEL_DIR
constant - Place your existing model in the directory that you specified before
- Run the
extract_data.py
file and call the methodcategorize_with_trained_model()
, this will categorize your regions of interest - Verify that the data was labeled correctly, by checking the
data_extracted
folder
- If you are finished labeling the images, run the
extract_data.py
file and call the methodrename_images_in_categories()
from theExtractor
, this will rename the images in each category - Run the
extract_data.py
file and call the methodcreate_training_data()
, this will create your pickle files (X.pickle
andy.pickle
) which contain the data and the labels of your data
- Check if the pickle files (
X.pickle
andy.pickle
) were created in the root directory of the project - Run the
train_model.py
file within the trainer, this will train your model and save it to the directory specified in theconstants.py
(MODEL_DIR
)
usage: test_model.py [-h] [--model_type MODEL_TYPE] [--model_path MODEL_PATH] [--test_image TEST_IMAGE]
[--test_folder TEST_FOLDER] [--test_on_random]
cnn-number-detection
optional arguments:
-h, --help show this help message and exit
--model_type MODEL_TYPE type of model to use
--model_path MODEL_PATH path to the saved model
--test_image TEST_IMAGE path to the image for testing the model
--test_folder TEST_FOLDER folder with images for inference
--test_on_random if a randomly generated image should be used for inference
- Check if the model is in the directory specified in
constants.py
(MODEL_DIR
) - Upload a testing image or a folder full of test images to the
tester
directory - Specify the correct model type (
--model_type
) and model path (--model_path
) - Run inference
- Run the
test_model.py
file within theTester
by specifying the image with--test_image
- Run the
test_model.py
file within theTester
by specifying the folder with--test_folder
- Run the
test_model.py
file within theTester
on a random image with--test_on_random
- Run the
Command to create a video out of a bunch of images.
convert *.jpg recognized.mpeg