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

训练大概要2个多小时,完成后,我们可以在笔记本上先测试一下。

1.生成deployment model

python merge_bn.py --model example/MobileNetSSD_deploy.prototxt --weights snapshot/mobilenet_iter_20000.caffemodel

生成no_bn.prototxt和no_bn.caffemodel

name: "MobileNet-SSD"
input: "data"
input_shape {
  dim: 1
  dim: 3
  dim: 300
  dim: 300
}
layer {
  name: "conv0"
  type: "Convolution"
  bottom: "data"
  top: "conv0"
...

layer {
  name: "detection_out"
  type: "DetectionOutput"
  bottom: "mbox_loc"
  bottom: "mbox_conf_flatten"
  bottom: "mbox_priorbox"
  top: "detection_out"
  include {
    phase: TEST
  }
  detection_output_param {
    num_classes: 5
    share_location: true
    background_label_id: 0
    nms_param {
      nms_threshold: 0.44999998807907104
      top_k: 100
    }
    code_type: CENTER_SIZE
    keep_top_k: 100
    confidence_threshold: 0.25
  }
}

2. 测试

# inference.py
import numpy as np
import sys,os
import cv2
caffe_root = '/home/young/caffe/'
sys.path.insert(0, caffe_root + 'python')
import caffe

net_file= 'no_bn.prototxt'
caffe_model='no_bn.caffemodel'

if not os.path.exists(caffe_model):
  print(caffe_model + " does not exist")
  exit()
if not os.path.exists(net_file):
  print(net_file + " does not exist")
  exit()
net = caffe.Net(net_file,caffe_model,caffe.TEST)

CLASSES = ('background', 'bunny', 'doll', 'doraemon', 'snoopy')
np.random.seed(3)
colors = 255 * np.random.rand(len(CLASSES), 3)

def preprocess(src):
  img = cv2.resize(src, (300,300))
  img = img - 127.5
  img = img * 0.007843
return img

def postprocess(img, out):
  h = img.shape[0]
  w = img.shape[1]
  box = out['detection_out'][0,0,:,3:7] * np.array([w, h, w, h])
  cls = out['detection_out'][0,0,:,1]
  conf = out['detection_out'][0,0,:,2]
return (box.astype(np.int32), conf, cls)

def detect(imgfile): 
  origimg = cv2.imread(imgfile)
  img = preprocess(origimg)
  img = img.astype(np.float32)
  img = img.transpose((2, 0, 1))
  net.blobs['data'].data[...] = img
  out = net.forward()
  box, conf, cls = postprocess(origimg, out)

  for i in range(len(box)):
    p1 = (box[i][0], box[i][1])
    p2 = (box[i][2], box[i][3])
    cv2.rectangle(origimg, p1, p2, (0,255,0),2)
    p3 = (max(p1[0], 15), max(p1[1]-10, 10))
    title = "%s:%.2f" % (CLASSES[int(cls[i])], conf[i])
    cv2.putText(origimg, title, p3, cv2.FONT_ITALIC, 0.8, colors[i], 2)
  
  cv2.imshow("Caffe SSD", origimg)
  k = cv2.waitKey(0) & 0xff 
  if k == 27 : return False
  return True

for f in os.listdir(test_dir):
  if detect(test_dir + "/" + f) == False:
    break

将测试图片拷贝到images目录下

测试训练集里没有的图片