Neural Network Sorter using Tensorflow, OpenCV on a Raspberry Pi
As I started thinking how to describe the software I built, this famous quote comes to mind:
“If you are not embarrassed by the first version of your product, you’ve launched too late.” — Reid Hoffman, LinkedIn founder
I am certainly embarrassed to show my code without going back and properly structure and documented in a modern way. As mentioned before, I followed the leanest path to achieve the outcome and it will take me a bit more time to go back and review my code to properly put in a form that can be shared.
That said, in this blog I will cover the overall software architecture approach and what I used from the community and what I built myself.
Tensorflow based Lego Sorter software architecture
Above is the overall map of the software I built, and roughly splits into 6 areas, with the Main Module being the best way to describe how it all comes together.
The Main Module
At the center of the software stack is what I call the Main Module, which orchestrates the operation of the Lego Sorter, and per the above image, the sequence of activities is pretty straight forward and brings in the rest of the modules.
What it accomplishes:
a) Enables automated creation of Training data sets by orchestrating feeding, image capture and correctly labeling 100–200 pictures at a time, with the only setup required is feeding a fixed type of block.
b) Executes the Sorting function by coordinating the feeding, image analysis and sorting of the blocks as they roll across the conveyors.
The main activities it performs are:
1. Brick Supply — controls the feeder pushing cycle and the speed to the conveyor belts using a linear regression model that keeps a fluid stream of bricks to flow through the system. There is a timing function that monitors each event and the progress of bricks to provide a feedback loop and adjust the triggers.
2. Camera Trigger — when a brick is in range of the IR Beam Sensor, it triggers the Camera and Image detection modules, which not only bring back an image, but also inform the Main Module about the number of objects in range and it’s location on the file.
3. Image Validation — if the return image from the Image Detection modules is a single object, it triggers one of two activities: a) start the logging module if this is a training set generation or b) starts the image recognition process. If there is no objects or more than one object in the image, it discards the bricks into Bin #0 (tags it as a Mechanical Separation failure).
4. Classify Image — if the return image from the Image Detection modules is a single object, it uses the Image Classifier module to evaluate and provide the return probabilities for classification.
5. Sort —the current logic uses the highest probability result to select the adequate Bin. The bins are not hardcoded, so the logic will assign a new bin as it encounters new brick types (constrained to 11 bin types due to the physical limitation of the sorter).
Interface and Logic Modules
Both the Servo and Motor Hats had readily available Python libraries that enabled control with a few simple Import and Sudo statements on the Pi.
The Infrared Sensor Beams and LEDs required me to create a custom board to simplify wiring and also supply the IR Beam required 5V and the LED required 12V. I used the GPIO library to directly communicate to the standard GPIO as well as the PCM pins I used to control and input feedback.
Once the interface with the hardware was set, I build purpose-specific logic to control each sensor, motor and servo.
An example of this is the logic that controls the Sorting Servo: the main module selects the bin number to sort, and the logic module is where the calibration between bin number and the specific input frequency is chosen to direct the servo to the right bin.
This logic layer provides the main module with just the output/input needed to run the Lego Sorting logic in 1 or 2 function calls for each hardware interface.
Camera and Image Detection Modules
As I started performing the initial runs on the mechanical separation, I did notice that the efficiency of brick separation was hovering around 95%.
Using OpenCV (and getting a lot of advice from the PyImageSearch blog) I built a robust Image Detection module that enabled me to identify if the mechanical separation was accurate and also avoid false triggers.
This is an extremely powerful module in the system, and one that presents you with great learning opportunities. This is were a lot of where the key activities of ‘data set generation’ and ‘data set cleaning’ happen. I will leave most of the details of my build out of this blog, so you can have an opportunity to try it for yourself through the Training and Evaluation image sets here.
The Neural Network Modules
The Image Logging Module
To achieve capturing a wide range of image sets, the Sorter also acts as an automated Image capture system. The Logging Module is a key component and a differentiator to the Cucumber Farmer, as it enables the automated capture of image training sets. By feeding the sorter a single type of piece, this module would store the image with its proper filename, tag and csv directory to be used for training.
The Neural Network training and deployment
My initial objective was to replicate the Cucumber Farmer approach, which is illustrated to rely on hybrid computing (Arduino / Pi / Google Cloud). This setup is straightforward and possible, but it would have required a considerable time to setup the http request and responses, in conjunction with authentication.
While a hybrid approach provides flexible computing power beyond the Pi, I decided to follow the shortest path, by running the training in my Desktop GPU and directly deploying the Model to the Raspberry Pi. For this, I relied heavily on the Google Codelabs Tensorflow examples and made a few modifications to the training and image labeling modules to execute towards my use case.
Running the model on the Pi consistently takes 3.8 seconds to classify each block. We can certainly qualify it as slow, but it proved ‘good-enough’ towards the object.
It is worth nothing that as a comparison, running the classification on my desktop takes 2.7 seconds to classify each block (desktop setup is Tensorflow with GPU install, Intel i7, CUDA Compute 3.5 enabled GPU (a GeForce GT 730) .
Getting Closer to 100% Classifier Accuracy
In it’s current state as of August 2018, the system and software are capable of achieving a 89% total Yield (Mechanical separation plus Classification Accuracy).
Watching the system perform closely, the biggest opportunities to get close to 100% classification accuracy is focused on a few key modules:
- Harden the Image Detection module to enable better detection of mechanical separation errors, which still account to significant cause for classification gaps (e.g. sending image of two overlapping bricks generates a wrong output).
- Adding additional cameras would instantly multiply the availability of training sets and would make for more powerful classification based on multiple capture angles.
- Increase Training Set to allow the model to have a more robust set to identify bricks from a single camera angle.
- Iterate the Neural Network Models by using different techniques, pre-trained models, higher resolution models and custom design of the neural network layers to optimize for this specific use case.
5 Parts to This Blog
This part of a 5 Blog Series to cover the mechanical and software design for the Lego Sorter, as well as sharing the training set and some evaluation sets:
1. Lego Sorter using TensorFlow on Raspberry
2. Mechanical Separation (Design, Motors and Sensors)
3. Overview of the Software stack
4. Using Inception V3 to Identify LEGO vs. Generic Bricks
5. Try It Yourself: 2 Big Data Sets so you can Replicate this Project