为机器人换个好使的脑子(十八)

终于,我们可以让六足机器人张脑子了!

1. 生成NCS需要的graph

$ cd MobileNet-SSD
$ mkdir models
$ cp no_bn.prototxt no_bn.caffemodel models
$ mvNCCompile models/no_bn.prototxt \
    -w models/no_bn.caffemodel \
    -s 12 -is 300 300 -o models/mobilenet_graph

2. 测试

#coding=utf8
# NCSDK API V2
from mvnc import mvncapi as mvnc
import numpy as np
import cv2
import time

GRAPH = 'models/mobilenet_graph'
VIDEO = '/home/sean/Downloads/robot.mp4'
fourcc = cv2.VideoWriter_fourcc(*'XVID')
CLASSES = ('background', 'bunny', 'doll', 'doraemon', 'snoopy')
input_size = (300, 300)
np.random.seed(3)
colors = 255 * np.random.rand(len(CLASSES), 3)

def preprocess(src):
    img = cv2.resize(src, input_size)
    img = img - 127.5
    img = img / 127.5
    return img.astype(np.float32)

# discover our device
devices = mvnc.enumerate_devices()
device = mvnc.Device(devices[0])
device.open()

# load graph onto the device
graph = mvnc.Graph('graph1')
with open(GRAPH, 'rb') as f:
    graph_file = f.read()
input_fifo, output_fifo = graph.allocate_with_fifos(device, graph_file)

capture = cv2.VideoCapture(VIDEO)
_, image = capture.read()
height, width = image.shape[:2]
out = cv2.VideoWriter('robot.avi',fourcc, 10.0, (width,height))
while capture.isOpened():
    stime = time.time()
    _, image = capture.read()
    image_pro = preprocess(image)
    graph.queue_inference_with_fifo_elem(input_fifo, output_fifo, image_pro, 'object1')
    output, user_obj = output_fifo.read_elem()
    valid_boxes = int(output[0])
    for i in range(7, 7 * (1 + valid_boxes), 7):
        if not np.isfinite(sum(output[i + 1: i + 7])):
            continue
        clss = CLASSES[int(output[i + 1])]
        conf = output[i + 2]
        color = colors[int(output[i + 1])]
        x1 = max(0, int(output[i + 3] * width))
        y1 = max(0, int(output[i + 4] * height))
        x2 = min(width, int(output[i + 5] * width))
        y2 = min(height, int(output[i + 6] * height))
        label = '{}: {:.0f}%'.format(clss, conf * 100)
        image = cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
        y = y1 - 5 if y1 - 15 > 15 else y1 + 18
        image = cv2.putText(image, label, (x1, y), cv2.FONT_HERSHEY_SIMPLEX, 0.8,color, 2)

    fps = 'FPS = {:.1f}'.format(1 / (time.time() - stime))
    image = cv2.putText(image, fps, (width-220, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.2,(0,255,0), 3)    
    cv2.imshow('frame', image)
    out.write(image)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

capture.release()
cv2.destroyAllWindows()
input_fifo.destroy()
output_fifo.destroy()
graph.destroy()
device.close()
device.destroy()

耶!从六足机器人的摄像头,我们可以正确地识别娃娃了。