Finished product first:
The video above shows real-time objection detection using a neural network model called Single Shot MultiBox Detector (SSD). The neural network model was trained on (and runs with) Tensorflow.
The “AI” model works much better than older Computer Vision (OpenCV) techniques. For example, here’s an example of a previous attempt to track objects with the color “tomato”. The way that worked was by filtering a camera frame for a specific color, then masking out that color to find the contour of an object.
There were many problems with this approach. For example, color-based detection is sensitive to lighting changes. It performs differently depending on whether it’s day or night time. Also, we get false detections because many other objects in the scene may have the same color.
But, one advantage of classic OpenCV techniques is they’re simple and fast. Fast enough to run in real time on a live camera feed on a tiny Raspberry Pi processor.
On the other hand, AI models are more accurate and powerful. They can detect many different objects in various orientations (even if they’re partially hidden or chopped out of frame). They’re less sensitive to lighting changes.
But the downside is they’re expensive to run. The Raspberry Pi doesn’t have quite enough horsepower. I tried it and it’s possible… but I wasn’t happy with the performance. I found some interesting USB AI accelerators like Google’s Coral but they’re pricey.
So I decided to go with the “Mars Rover” approach and use a remote inference server (running on a PC in my living room). This buys us “infinite” computing power for a fixed cost of network latency. As long as your network is fast, this scheme works well.
One thing going for us is even though the Raspberry Pi camera can capture HD images, the SSD AI model only needs 300×300 pixel inputs. So we can downsize the images before sending them over the network to improve performance.
Here’s the Raspberry Pi code. All it does is capture frames, encode them as Base64 strings and makes requests to our remote inference server:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
#/usr/bin/python import base64 import cv2 import json import picamera from picamera.array import PiRGBArray import requests import time import zmq def process_target(image, x1, y1, x2, y2, label='target'): font = cv2.FONT_HERSHEY_SIMPLEX cv2.putText(image, "%s detected" % label, (20,30), font, 0.5,(255,255,255),1,cv2.LINE_AA) X = int((x2 + x1)/2) Y = int((y2 + y1)/2) ######################################################## # TODO: do something with the detected target at (X,Y) ######################################################## if X < 280 and X < 340 and Y > 230 and Y < 250: # ON TARGET - red cross, red rectangle cv2.line(image,(300,240),(340,240),(64,64,255),1) cv2.line(image,(320,220),(320,260),(64,64,255),1) cv2.rectangle(image,(int(x1),int(y1)),(int(x2),int(y2)),(64,64,255),1) else: # white rectangle cv2.rectangle(image,(int(x1),int(y1)),(int(x2),int(y2)),(255,255,255),1) def detect_objects(image, jpgstr, scale=2): URL = "http://192.168.1.187:9002/predict" response = requests.post(URL, data={"frame": jpgstr}) ret = json.loads(response.text) if 'success' not in ret or ret['success'] != True: return # green cross cv2.line(image,(300,240),(340,240),(128,255,128),1) cv2.line(image,(320,220),(320,260),(128,255,128),1) # get largest box targets = ["person", "dog", "cat", "book", "teddy bear", "sports ball", "banana"] priorityTargets = ["dog", "cat", "book", "teddy bear", "sports ball", "banana"] target = None boxes = ret['boxes'] for b in boxes: if b['label'] not in targets: continue # annotate boxsize rect = b['bbox'] b['boxsize'] = (rect[2]-rect[0]) * (rect[3]-rect[1]) if target == None: target = b elif b['boxsize'] > target['boxsize'] \ or (b['label'] in priorityTargets and target['label'] not in priorityTargets): target = b if target: rect = target['bbox'] process_target(image, rect[0]*scale, rect[1]*scale, rect[2]*scale, rect[3]*scale, target['label']) cv2.imshow('frame', image) cv2.waitKey(1) def camera_loop(): zcontext = zmq.Context() zsock = zcontext.socket(zmq.PUB) zsock.bind('tcp://*:5555') camera = picamera.PiCamera() camera.resolution = (640, 480) camera.framerate = 7 rawCapture = PiRGBArray(camera, size=(640, 480)) for frame in camera.capture_continuous(rawCapture, format="bgr", use_video_port=True): try: image = frame.array scale = 2 halfsizeImage = cv2.resize(image, (int(image.shape[1]/scale), int(image.shape[0]/scale)), cv2.INTER_AREA) encoded, buffer = cv2.imencode('.jpg', halfsizeImage) jpgstr = base64.b64encode(buffer) zsock.send(jpgstr) detect_objects(image, jpgstr, scale) rawCapture.truncate(0) except Exception as e: rawCapture.truncate(0) |
Now for the object detection neural network. Google provides a set of pre-trained models for object detection in their Model Zoo. The model I picked was ssdlite_mobilenet_v2_coco. You’ll also need this file to convert the detection IDs to text labels: coco-labels-paper.txt
I made a Python module to load and run the pre-trained Tensorflow object detection model. The class Predictor in coco_predictor.py below loads the saved model and provides a predict() function for running inference on new input images.
coco_predictor.py
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
#!/usr/bin/python # https://github.com/opencv/opencv/wiki/TensorFlow-Object-Detection-API import cv2 as cv import tensorflow as tf class Predictor: def __init__(self, savedModel, labels): # load labels f = open(labels, "r") self.labels = [l.strip() for l in f.readlines()] f.close() # Read the graph. with tf.gfile.FastGFile(savedModel, 'rb') as f: self.savedModel = savedModel self.graphDef = tf.GraphDef() self.graphDef.ParseFromString(f.read()) # Restore session self.sess = tf.Session() self.sess.graph.as_default() tf.import_graph_def(self.graphDef, name='') def predict(self, img, thresh): rows = img.shape[0] cols = img.shape[1] inp = cv.resize(img, (300, 300)) inp = inp[:, :, [2, 1, 0]] # BGR2RGB # Run the model out = self.sess.run(\ [self.sess.graph.get_tensor_by_name('num_detections:0'), self.sess.graph.get_tensor_by_name('detection_scores:0'), self.sess.graph.get_tensor_by_name('detection_boxes:0'), self.sess.graph.get_tensor_by_name('detection_classes:0')], feed_dict={'image_tensor:0': inp.reshape(1, inp.shape[0], inp.shape[1], 3)}) detections = [] num_detections = int(out[0][0]) for i in range(num_detections): classId = int(out[3][0][i]) score = float(out[1][0][i]) bbox = [float(v) for v in out[2][0][i]] if score <= thresh: continue x = bbox[1] * cols y = bbox[0] * rows right = bbox[3] * cols bottom = bbox[2] * rows detections.append({ 'label': self.labels[classId-1], 'score': score, 'bbox': [x, y, right, bottom] }) return detections |
Finally, the Flask inference server that runs on a PC. The Raspberry Pi sends it images and it replies with detections:
coco_flask_server.py
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
#!/usr/bin/python import base64 import cv2 import coco_predictor import logging import sys import numpy as np logging.basicConfig(stream=sys.stdout, level=logging.DEBUG) model = "models/ssdlite_mobilenet_v2_coco_2018_05_09/frozen_inference_graph.pb" labels = "models/coco-labels-paper.txt" predictor = coco_predictor.Predictor(model, labels) def mlexec(query): boxes = [] try: img = query['img'] boxes = predictor.predict(img, 0.3) except Exception as e: logging.error(str(e)) response = { 'success': True, 'boxes': boxes } return response logging.info("Start COCO server...") from flask import Flask, jsonify, request application = Flask(__name__) @application.route('/') def index(): return 'Hello!' @application.route('/predict', methods=['GET', 'POST']) def predict(): if request.method == 'GET': return ':)' response = { 'success': False } try: frame = base64.b64decode(request.form['frame']) npimg = np.fromstring(frame, dtype=np.uint8) img = cv2.imdecode(npimg, 1) query = { 'img': img } response = mlexec(query) except Exception as e: logging.error(str(e)) return jsonify(response) if __name__ == '__main__': application.run(host="0.0.0.0", port=9002, debug=True) |
Oh, and if you’re curious about the Robot Kit I’m using, it’s made by Adeept. They call it the Mars Rover PiCar-B:
The hardware is excellent and well-designed. All the pieces fit together perfectly and mount onto a solid acrylic chassis with bolts. There is a steering rack and a RWD drivetrain that sends power from one motor to both rear wheels. You can pretty much follow the instructions and everything fits together like a Lego set. And when you’re not using it in robot mode, it’s a good looking “case” that lets you use it like a regular Raspberry Pi on your desk.
The software, however, is so-so. You definitely need programming experience. Some things might not work right out of the box, and some things need fiddling around with. On the plus side, they do provide tons of working code for all the various sensors, servos and components. So as long as you are comfortable with Python, you can use their code as excellent references. You can pick and choose and mix and match what you need.
I like this robot kit because it uses a Raspberry Pi (which I already have). So it uses standard Linux stuff. You can use state-of-the-art software and write regular programs like a civilized person. And because it uses a Pi, you get WiFi, Ethernet, Bluetooth, USB, HDMI, etc for free. No painful caveman Arduino programming. It runs on CR 18650 lithium ion batteries (which I also have from salvaging old laptop batteries). It can also run on the regular Pi USB power source.
“Behind the scenes”: Driving school. Programming the head.
I hope to make it autonomous and self-recharging one day!
Happy hacking!
aaron@secretsciencelab.com