Published November 1, 2020 © GPL3+

Make Candy Sorting Easy with Machine Learning

Automatically go through your stash of candy and have it sorted into piles with the power of machine learning.

IntermediateFull instructions provided6 hours1,954

Make Candy Sorting Easy with Machine Learning

Things used in this project

Hardware components

Raspberry Pi 4 Model B

Raspberry Pi Camera Module V2

Arduino Mega 2560

Stepper Motor, Mini Step

Driver DRV8825 for Stepper Motors for Theremino System

Power MOSFET N-Channel

Software apps and online services

Arduino IDE

Microsoft VS Code

Autodesk Fusion

Hand tools and fabrication machines

3D Printer (generic)

Story

Inspiration

Recently, I've been experimenting with training TensorFlow models on image datasets using Edge Impulse. So while thinking about Halloween, there's a common problem that nearly everyone runs into. Why does it take so long to sort through and organize all of the candy? Because picking up each piece individually and chucking it into its respective pile consumes far too much effort, I wanted to automate the process using the power of machine learning.

Principle of Operation

The theory goes as follows: the user begins by placing a piece of candy onto a platform, where it is then scanned by a camera. That image gets classified by a model which produces a label of what it thinks the candy is. Based on the result, a pile is chosen, and then platform then moves to that location. Once it has arrived, some kind of pushing mechanism shoves the candy into the pile, and the platform then returns and the cycle resets.

Setting Up Edge Impulse

To gather training data and produce a TensorFlow model, I went with Edge Impulse. I began by creating a new project called "candySorter" and then heading to the device page. Edge Impulse has a nice feature where you can import many types of data an automatically infer their labels based on just the filename, so I wrote a simple Python script that takes a picture using the Pi's camera and saves it with the desired label when a button is pressed.

1 / 2

Training and Deploying a Model

Now that there is a large amount of image data, it's time to train a TensorFlow model on it. For input, I went with an image block that scales the original image to one that is 96x96 pixels large. The processing block just sets the color depth to RGB instead of monochrome, and the training block utilizes transfer learning with the MobileNetV2 model. Output is one of the following labels: kitkat, sour patch, twizzler, or reeses.

After training had finished, I deployed the model as a standalone web assembly library that can be called externally from a Node.js source file.

Processing Images

The system works by using a Raspberry Pi 4 to capture an image with the Raspberry Pi Camera module, but the classifier can't understand that format, it needs an array of pixel values of a certain size instead. So the capture gets saved to a BytesIO stream and resized to 96x96 pixels large. It is then converted into a string of flattened pixel values, which takes in a tuple of RGB values and transforms it into a single hex value and appends it to a list.

1 / 2

Classifying an Object

The Node.js source file is called by the subprocess.run() function, and it passes the color data string as an argument to the program where it is then parsed and given to the model. After the model comes up with what it thinks the object is, a JSON string is outputted over stdout, where it is then piped into a Python variable and parsed as a dictionary object. Finally, the maximum value is found and the associated label is extracted and used to determine into which bin the candy should be placed. This bin number is then sent over USB to the awaiting Arduino Mega 2560.

The Hardware

First and foremost, there is the x-axis roller, which gets moved into position by a single NEMA17 stepper motor that is driven by a DRV8825 stepper motor driver module.

My 3D printer is broken, so I had to make this one from wood instead :/

When the Arduino Mega 2560 starts up, it homes the axis by slowly sliding the roller until it hits the limit switch, at which point it's at 0. When there's new serial data, it gets read and parsed into an integer, which says the bin the candy should be placed into. The bins are spaced 10cm apart, so this number is multiplied by 100mm and the stepper moves the corresponding number of steps. Once at the destination, there's a solenoid on the top that activates and quickly pushes the candy below.

1 / 2

Sorting Candy into Piles

With the ML model loaded onto the Pi, I plugged everything in. the actual sorting mechanism whirred to life and home on its own, and the Pi started showing a preview of what it was seeing. From there, I simply placed a piece of candy onto the platform and pressed the button to take a picture of it. That image was resized and sent to the model, which determined its label. Finally, the bin number was sent over and the candy was deposited into the correct spot. In the future, I would like to automate this even more by adding a hopper that can have candy just dumped in the top and then deposit each piece individually over time.

Schematics

Code

from time import sleep
from picamera import PiCamera
from gpiozero import Button
from io import BytesIO
from PIL import Image
from signal import pause
import json
import subprocess
import serial

#one for each label
binList = {
    "kitkat": 0,
    "twizzler": 1,
    "reeses": 2,
    "sour patch": 3,
    "none": 4
}

ser = serial.Serial(port="/dev/ttyACM0", baudrate=115200)
camera = PiCamera()
stream = BytesIO()
resolution = (2592, 1944)

button = Button(2)

def tupleToValue(rgbValues: tuple):
    newValue = rgbValues[0] * 65536 + rgbValues[1] * 256 + rgbValues[2]
    print(f"Hex representation of the tuple {rgbValues} is: {hex(newValue)}")
    return hex(newValue)

def classifyAndSendValue(dataString: str):
    result = subprocess.run(['nodejs', 'run-impulse.js', dataString], stdout=subprocess.PIPE).stdout.decode('utf-8')
    print(result)
    result = json.loads(result)["results"]
    print(f"Result is {result}")
    maxIndex = -1.0
    currentLabel = ""
    for item in result:
        print(f'checking item {item["label"]}, value={item["value"]}')
        if item["value"] > maxIndex: 
            maxIndex = item["value"]
            currentLabel = item["label"]
    correspondingBin = binList[currentLabel]
    print(f"Bin is ", correspondingBin)
    toSend = (str(correspondingBin)+'\n').encode('utf-8')
    ser.write(toSend)
    print(f"Sent {toSend} over serial")

def processCamera():
    camera.capture(stream, resize=(96, 96), format='jpeg')
    stream.seek(0)
    image = Image.open(stream)
    image.convert("RGB")
    rgbData = list()
    numberList = list(image.getdata())
    for number in numberList:
        rgbData.append(tupleToValue(number))
    rgbDataString = ','.join(rgbData)
    #print("Raw features result is: ", rgbDataString)
    classifyAndSendValue(rgbDataString)
    sleep(2)

button.when_pressed = processCamera

camera.resolution = resolution
camera.framerate = 15
camera.start_preview()

sleep(2)

pause()

#include "BasicStepperDriver.h"

// X axis STEP, DIR, EN, LIMIT_SW, SOLENOID
const uint8_t pins[] = {6, 7, 8, 26, 11};

#define MICROSTEPS 16
#define RPM 120
#define STEPS_PER_MM 200    // 200 steps (including microsteps) per mm moved
#define TOTAL_LENGTH_MM 700
#define TOWARDS_0_DIR -1
#define AWAY_FROM_0_DIR -1 * TOWARDS_0_DIR

BasicStepperDriver stepper(200, pins[1], pins[0], pins[2]);

void setup()
{
    Serial.begin(115200);
    initPins();
    stepper.begin(RPM, MICROSTEPS);
    stepper.disable();

    homePosition();
    stepper.disable();
}

void loop()
{
    
}

void serialEvent()
{
    long binToPlace = Serial.parseInt();

    long targetSteps = binToPlace * STEPS_PER_MM * 100;     // 100 mm between bins

    stepper.enable();
    stepper.move(AWAY_FROM_0_DIR * targetSteps);    // get to target
    
    digitalWrite(pins[4], HIGH);
    delay(500);     // push item off
    digitalWrite(pins[4], LOW);

    stepper.move(TOWARDS_0_DIR * targetSteps);  // return to home
    stepper.disable();
}

void initPins()
{
    pinMode(pins[3], INPUT_PULLUP);
    pinMode(pins[4], OUTPUT);

    digitalWrite(pins[4], LOW);
}

void homePosition()
{
    stepper.enable();
    if(!digitalRead(pins[3]))
        stepper.rotate(AWAY_FROM_0_DIR * 720.0); // rotate away from home

    while(digitalRead(pins[3]))
        stepper.move(TOWARDS_0_DIR * 5);
}