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

接下来,最后的扫尾工作:我们要让六足机器人找到兔子。

这个比较简单,我们把识别出兔子的框中心与摄像头画面中心对准,就行了。

代码:

# -*- coding: utf-8 -*-
# NCSDK API V2
from mvnc import mvncapi as mvnc
import numpy as np
import cv2
import time
import sonar
import time
import spider
import threading

GRAPH = 'graph'
CLASSES = ('background', 'bunny', 'doll', 'doraemon', 'snoopy')
distance = 0.0
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)

# get the distance
def updateDistance():
  global distance
  while True:
    try:
      distance = sonar.getDistance()
    except Exception as e:
    print('Error: %s' %e)
    distance = 0.0
    time.sleep(0.2)

th_distance = threading.Thread(target=updateDistance)
th_distance.setDaemon(True)
th_distance.start()

# 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(0)
_, image = capture.read()
height, width = image.shape[:2]

spider.runActionGroup('standUp', 1)

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])]
    # get the bonding box
    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))
    # display class label and bonding box
    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)

    # follow the white rabbit: align the center of bonding box with the screen center
    if clss == 'bunny':
      x_diff = (x1+x2)/2 - width/2
      y_diff = (y1+y2)/2 - height/2
      if distance > 10:
        if x_diff > 20:
          spider.runActionGroup('turn-right', 0)
        elif x_diff < -20:
          spider.runActionGroup('turn-left', 0)
        elif y_diff > 20:
          spider.runActionGroup('head_up', 0)
        elif y_diff < -20:
          spider.runActionGroup('head_down', 0)
        else:
          spider.runActionGroup('walk', 0)
      else:
        spider.runActionGroup('stop', 0)

  cv2.imshow('frame', image)
  print('FPS = {:.1f}'.format(1 / (time.time() - stime)))
  if cv2.waitKey(1) & 0xFF == ord('q'):
    break

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