Image Recognition Robot with PiStorms and Pi Camera

Select the Icons

 
Before you write the program, you should make a new folder on your Raspberry Pi. This folder will contain your image recognition program, as well as all of the icons you wish for the program to identify.

An example of making a new folder for our icon recognition program:

sudo mkdir /home/pi/PiStorms/programs/09-IconTracker

The program we will make today will require .png files and they will need a square frame around them. Below are example icons that you may wish to use. These icons are used by the Polylander's image recognition programs. See the video for more information.

Example .png images:

adult.png

circle-arrow-left.png
(You can also use this as “circle-arrow-up.png” and “circle-arrow-right.png” by simply rotating the image)

map-marker-alt.png

 
When the program runs, it will check for all of the other .png files that are in that same folder as the program and these images will be the ones that the program is trained to recognize. When it sees an image, the program will return the name of that icon (ex: “adult.png”).
 

Writing the Icon Recognition Program

 
After you have the folder set up with the icons you would like to identify, we can now write the Python program that will identify the icon and return its name! Follow the steps below to write the program.

First, we need to import the libraries necessary to run a PiStorms program, as well as the libraries needed for the Pi Camera and some additional setup code:

1. from picamera.array import PiRGBArray
2. from picamera import PiCamera
3. from PiStorms import PiStorms
4. import cv2
5. import sys, os
6. import imutils
7. import numpy as np
8. import argparse, time
9.  
10. import inspect
11. currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
12. parentdir = os.path.dirname(currentdir)
13. sys.path.insert(0,parentdir)
14.  
15. psm = PiStorms()

Second, we're going to write the code for the “training” function. What this does is checks for .png files that are in the same folder as the program. Then, it takes these images, changes them to grayscale and converts them into an array in order to store them in a form the program can recognize.
class icontracker:

1. training = {}
2. version = ''
3. # Generate the images dictonary and store it in memory
4. # Each image is 200x200
5. def __init__(self):
6. self.training = self.get_training()
7. self.version = '1.00'
9. def get_training(self):
10. path = currentdir
11. training = {}
12. files = os.listdir(path)
13. files_png = [i for i in files if i.endswith('.png')]
14. print "files_png: ", files_png
15. for file in files_png :
16. imR = cv2.imread(currentdir+"/"+file)
17. # perform the actual resizing of the image and show it
18. imR = self.preprocess(imR)
19. training[file] = imR
21. return training

Next, we'll define the function for capturing an image from the Pi Camera.

1. # Captures a single image from the camera and returns it in PIL format
2. def get_image(self,camera):
3. # read is the easiest way to get a full image out of a VideoCapture object.
4. retval, ig = camera.read()
5. return im

Now, we can start moving into our image recognition code. Let's define a few more functions:

1. ###############################################################################
2. # Image Matching
3. ###############################################################################
4. def preprocess(self,img):
5. gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
6. blur = cv2.GaussianBlur(gray,(5,5),5 )
7. thresh = cv2.adaptiveThreshold(blur,255,1,1,11,1)
8. return thresh
10. def imgdiff(self,img1,img2):
12. img1 = cv2.GaussianBlur(img1,(5,5),5)
13. img2 = cv2.GaussianBlur(img2,(5,5),5)
14. diff = cv2.absdiff(img1,img2)
15. diff = cv2.GaussianBlur(diff,(5,5),5)
16. flag, diff = cv2.threshold(diff, 200, 255, cv2.THRESH_BINARY)
17. return np.sum(diff)
18.  
19. def find_closest_card(self,training,img):
20. features = preprocess(img)
21. return sorted(training.values(), key=lambda x:imgdiff(x[1],features))[0][0]

In order to know if we have found our icon, the robot needs to see the square border around the image. This function will see if the icon in front of it has a square border.

1. def findSquare( self,frame ):
2. gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
3. blurred = cv2.GaussianBlur(gray, (7, 7), 0)
4. edged = cv2.Canny(blurred, 60, 60)
5. # find contours in the edge map
6. (cnts, _) = cv2.findContours(edged.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
7. # loop over our contours to find hexagon
8. cnts = sorted(cnts, key = cv2.contourArea, reverse = True)[:50]
9. screenCnt = None
10. for c in cnts:
11. # approximate the contour
12. peri = cv2.arcLength(c, True)
13. approx = cv2.approxPolyDP(c, 0.004 * peri, True)
14. # if our approximated contour has four points, then
15. # we can assume that we have found our square
17. if len(approx) >= 4:
18. screenCnt = approx
19. x,y,w,h = cv2.boundingRect(c)
20. cv2.drawContours(image, [approx], -1, (0, 0, 255), 1)
21. #create the mask and remove rest of the background
22. mask = np.zeros(image.shape[:2], dtype = "uint8")
23. cv2.drawContours(mask, [screenCnt], -1, 255, -1)
24. masked = cv2.bitwise_and(image, image, mask = mask)
25. #crop the masked image to to be compared to referance image
26. cropped = masked[y:y+h,x:x+w]
27. #scale the image so it is fixed size as referance image
28. cropped = cv2.resize(cropped, (200,200), interpolation =cv2.INTER_AREA)
30. return cropped

Just two more functions to go. These last two are used for locating and identifying the icon:

1. def locate_target( self,frame ):
2. #find the square logo image from our image
3. self.target = self.preprocess( self.findSquare(frame))
4. return self.target
5.  
6. def identify_target( self,frame ):
7. results = {}
8. for file in self.training :
9. results[file] = self.imgdiff(self.locate_target(frame),self.training[file])
10.  
11. x = min (((e) for e in results if e), key=results.get)
12. return x

Now it's time to put these functions together in a cohesive program so we can identify the icons! First, we need to initialize the Pi Camera, setting up resolution and frame rate.

1. if __name__ == '__main__':
3. icon = icontracker()
4. print icon.training
5. print icon.version
6. # Now we can initialize the camera capture object with the cv2.VideoCapture class.
7. # All it needs is the index to a camera port.
8. camera = PiCamera()
9. camera.resolution = (320, 240)
10. rawCapture = PiRGBArray(camera, size=(320, 240))
11. camera.framerate = 30

We will now “ramp” the camera, which will throw away frames as the camera is stating up.

1. # Ramp the camera - these frames will be discarded and are only used to allow v4l2
2. # to adjust light levels, if necessary
3. '''
4. for i in xrange(ramp_frames):
5. temp = icon.get_image(camera)
6. '''
7. i =0
8. lasttime = time.time()

Processing the image captured by the camera so we can compare it to the images from “training:”

1. for frame in camera.capture_continuous(rawCapture, format="bgr", use_video_port=True):
2. # grab the raw NumPy array representing the image, then initialize the timestamp
3. # and occupied/unoccupied text
5. i = i+1
6. image = frame.array
7. gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
8. print i
9.  
10. '''
11. #find the square logo image from our image
12. target = icn.preprocess( icn.findSquare(image))
13. cv2.imshow('target',target)
14. cv2.waitKey(0)
15. print 'target', target.shape
16.  
17. results = {}
18. for file in icn.training :
19. results[file] = icn.imgdiff(target,icn.training[file])
20.  
21. print min(results, key=results.get)
22. '''

Finally, for the last part of the image recognition program, we will store the name of the image the robot has identified in a variable called “img.” The robot will print the name of the image to the PiStorms screen (ex: “adult.png”), as well as a time stamp. Now, we can program cases for the robot to react to the identified image.

1. img = icon.identify_target(image)
2. print "identified: " , img
3.  
4. psm.screen.termPrintAt(7, "count: " + str(i))
5. psm.screen.termPrintAt(8, "identified: " + str(img))
6. print 1000*(time.time() - lasttime)
7. lasttime = time.time()
8. # clear the stream in preparation for the next frame
9. rawCapture.truncate(0)

Programming the Robot's Reaction to the Icons

Now, we can program our robot to react to whichever icon it has identified.

Recall in the last step, that we ended the image recognition portion of the program by storing the name of the identified image in a variable called “img.” Now, we are going to have the robot read this variable and program in reactions based on whichever image is identified.

To do this, we are going to have an if statement for each image we want the robot to react to in that program, followed by the series of commands we want the robot to execute after it has seen the icon. It should look something like this (be sure to replace “image_name.png” with the actual image name):

1. if (img == “image_name.png”):
2. # do something

Here is an example of how to do this, taken directly from the Polylander's “PolylanderIconArrow.py” program. This program identifies the arrow icons. Once an icon is identified, the Polylander displays the identified icon on the screen and steers in the direction the arrow tells it to drive:

1. if (img == "circle-arrow-left.png"):
2. # Turn left:
3. psm.screen.fillBmp(30, 0, 240, 240, path = currentdir+'/'+"circle-arrow-left.png")
4. psm.BAM1.setSpeed(100)
5. psm.BAM2.setSpeed(30)
6.  
7. if (img == "circle-arrow-right.png"):
8. # Turn right:
9. psm.screen.fillBmp(30, 0, 240, 240, path = currentdir+'/'+"circle-arrow-right.png")
10. psm.BAM1.setSpeed(30)
11. psm.BAM2.setSpeed(100)
12.  
13. if (img == "circle-arrow-up.png"):
14. # Drive forward:
15. psm.screen.fillBmp(30, 0, 240, 240, path = currentdir+'/'+"circle-arrow-up.png")
16. psm.BAM1.setSpeed(75)
17. psm.BAM2.setSpeed(75)

An important thing to note is that unless you program it to, the camera will not continue to capture and identify new images while in an if statement, so try to keep them short.

It may be wise to choose to make a case for when the robot sees nothing (which is img = “none.png”). This will be the default case, if no icon is detected. Below is an example from the Polylander program, which stops the robot and clears the screen if no image is seen:

1. if (img == "none.png"):
2. # If there is no image found, halt:
3. psm.screen.clearScreen()
4. psm.BAM1.brake()
5. psm.BAM2.brake()

Finally, this last piece of code is very important. This piece of code will quit the program if the “GO” button is pressed, providing a way to exit the program and return to the PiStorms browser. This piece of code should go just after all of your other image “if” statements.

1. # Exiting the program:
2. if (psm.isKeyPressed()):
3. psm.BBM1.brake()
4. psm.BAM1.brake()
5. psm.BAM2.brake()
6. psm.screen.clearScreen()
7. psm.screen.termPrintAt(9, "Exiting to menu")
8. time.sleep(0.5)
9. quit()

Check out the image recognition program in action on the Polylander, a robotic tank made using PiStorms:

Program Downloads

You can download the programs and icons in this zip