人工智能
直播中

IC爬虫

11年用户 172经验值
擅长:1344
私信 关注
[经验]

【瑞芯微RK1808计算棒试用体验】+ mobileNet SSD 练个人头检测模型

`RK1808官方提供了mobileNet SSD的多分类的demo,所以基于这个算法可以较容易开发自己的检测算法,所以用网上公开的人头检测数据集训练一个人头检测模型。

caffe 版的mobilenetSSD开源地址:https://github.com/chuanqi305/MobileNet-SSD.git
原版的SSD开源地址:https://github.com/weiliu89/caffe.git
1.下载上面两个开源代码,先确保编译原版的SSD caffe版本正常
2.将mobileNetSSD 的这个项目代码,cp到原版的SSD的example目录下
3.所用人头检测数据集为hollywoodHead
4.将数据集数据只作为了VOC格式
5.根据原版SSD的lmdb训练数据制作方法将hoollywoodHead的数据制作成train和test两个lmdb数据,方便后续训练加载
6.使用如下命令将你制作的lmdb数据链接到mobileNetSSD要求的目录下:
  1. ln -s PATH_TO_YOUR_TRAIN_LMDB trainval_lmdb
  2. ln -s PATH_TO_YOUR_TEST_LMDB test_lmdb
7.在mobileNetSSD目录下制作标签文件,我的标签文件如下:
lablemap.png
9.运行目录下的gen_model.sh生成用于training 的prototxt文件
10.下载预训练模型,地址:https://drive.google.com/file/d/0B3gersZ2cHIxVFI1Rjd5aDgwOG8/view
11.运行训练脚本train.sh,迭代了120000次,得到如下人头检测模型:
iter120000.png

12.运行使用web 摄像头的测试脚本,脚本代码如下:
  1. import numpy as np  
  2. import sys,os  
  3. import cv2
  4. # caffe_root = '/media/isp/DATA/workspace/opencvLab/caffe/'
  5. # sys.path.insert(0, caffe_root + 'python')  
  6. import caffe

  7. net_file= './example/MobileNetSSD_deploy.prototxt'  
  8. caffe_model='./snapshot/mobilenet_iter_120000.caffemodel'  
  9. test_dir = "images"

  10. if not os.path.exists(caffe_model):
  11.     print(caffe_model + " does not exist")
  12.     exit()
  13. if not os.path.exists(net_file):
  14.     print(net_file + " does not exist")
  15.     exit()
  16. net = caffe.Net(net_file,caffe_model,caffe.TEST)  

  17. CLASSES = ('background',
  18.            'head')


  19. def preprocess(src):
  20.     img = cv2.resize(src, (300,300))
  21.     img = img - 127.5
  22.     img = img * 0.007843
  23.     return img

  24. def postprocess(img, out):   
  25.     h = img.shape[0]
  26.     w = img.shape[1]
  27.     box = out['detection_out'][0,0,:,3:7] * np.array([w, h, w, h])

  28.     cls = out['detection_out'][0,0,:,1]
  29.     conf = out['detection_out'][0,0,:,2]
  30.     return (box.astype(np.int32), conf, cls)

  31. def detect(imgSRC):
  32.     img = preprocess(imgSRC)
  33.    
  34.     img = img.astype(np.float32)
  35.     img = img.transpose((2, 0, 1))

  36.     net.blobs['data'].data[...] = img
  37.     out = net.forward()  
  38.     box, conf, cls = postprocess(imgSRC, out)

  39.     for i in range(len(box)):
  40.         if conf[i] > 0.4:
  41.             p1 = (box[i][0], box[i][1])
  42.             p2 = (box[i][2], box[i][3])
  43.             cv2.rectangle(imgSRC, p1, p2, (0,255,0))
  44.             p3 = (max(p1[0], 15), max(p1[1], 15))
  45.             title = "%s:%.2f" % (CLASSES[int(cls[i])], conf[i])
  46.             cv2.putText(imgSRC, title, p3, cv2.FONT_ITALIC, 0.6, (0, 255, 0), 1)
  47.     cv2.imshow("SSD", imgSRC)

  48.     k = cv2.waitKey(1) & 0xff
  49.         #Exit if ESC pressed
  50.     if k == 27 : return False
  51.     return True

  52. cap = cv2.VideoCapture(0)
  53. while(True):
  54.     ret, frame = cap.read()
  55.     detect(frame)
这个我训练好的模型:
mobilenet_iter_120000.caffemodel.zip (19.81 MB)
(下载次数: 42, 2019-10-25 17:24 上传)

测试结果:

` head_test.png

更多回帖

发帖
×
20
完善资料,
赚取积分