用作者提供的net1->net2生成MTCNN的训练样本(positive,negative,part,landmark)



本代码基于作者提供的python版本代码修改,参考:

https://github.com/DuinoDu/mtcnn/blob/master/demo.py  (作者提供)

https://github.com/dlunion/mtcnn/blob/master/train/gen_48net_data2.py


1,生成positive,negative,part三种样本,用作者的net1->net2生成bbox, 根据预测的bbox和ground truth计算IOU:

positive: IOU >= 0.65;

negative: IOU < 0.3;

part: 0.4 <= IOU < 0.65


代码如下:



 
 
  1. #!/usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. import _init_paths
  4. import caffe
  5. import cv2
  6. import numpy as np
  7. #from python_wrapper import *
  8. import os
  9. def bbreg(boundingbox, reg):
  10. reg = reg.T
  11. # calibrate bouding boxes
  12. if reg.shape[ 1] == 1:
  13. print "reshape of reg"
  14. pass # reshape of reg
  15. w = boundingbox[:, 2] - boundingbox[:, 0] + 1
  16. h = boundingbox[:, 3] - boundingbox[:, 1] + 1
  17. bb0 = boundingbox[:, 0] + reg[:, 0]*w
  18. bb1 = boundingbox[:, 1] + reg[:, 1]*h
  19. bb2 = boundingbox[:, 2] + reg[:, 2]*w
  20. bb3 = boundingbox[:, 3] + reg[:, 3]*h
  21. boundingbox[:, 0: 4] = np.array([bb0, bb1, bb2, bb3]).T
  22. #print "bb", boundingbox
  23. return boundingbox
  24. def pad(boxesA, w, h):
  25. boxes = boxesA.copy() # shit, value parameter!!!
  26. tmph = boxes[:, 3] - boxes[:, 1] + 1
  27. tmpw = boxes[:, 2] - boxes[:, 0] + 1
  28. numbox = boxes.shape[ 0]
  29. dx = np.ones(numbox)
  30. dy = np.ones(numbox)
  31. edx = tmpw
  32. edy = tmph
  33. x = boxes[:, 0: 1][:, 0]
  34. y = boxes[:, 1: 2][:, 0]
  35. ex = boxes[:, 2: 3][:, 0]
  36. ey = boxes[:, 3: 4][:, 0]
  37. tmp = np.where(ex > w)[ 0]
  38. if tmp.shape[ 0] != 0:
  39. edx[tmp] = -ex[tmp] + w -1 + tmpw[tmp]
  40. ex[tmp] = w -1
  41. tmp = np.where(ey > h)[ 0]
  42. if tmp.shape[ 0] != 0:
  43. edy[tmp] = -ey[tmp] + h -1 + tmph[tmp]
  44. ey[tmp] = h -1
  45. tmp = np.where(x < 1)[ 0]
  46. if tmp.shape[ 0] != 0:
  47. dx[tmp] = 2 - x[tmp]
  48. x[tmp] = np.ones_like(x[tmp])
  49. tmp = np.where(y < 1)[ 0]
  50. if tmp.shape[ 0] != 0:
  51. dy[tmp] = 2 - y[tmp]
  52. y[tmp] = np.ones_like(y[tmp])
  53. # for python index from 0, while matlab from 1
  54. dy = np.maximum( 0, dy -1)
  55. dx = np.maximum( 0, dx -1)
  56. y = np.maximum( 0, y -1)
  57. x = np.maximum( 0, x -1)
  58. edy = np.maximum( 0, edy -1)
  59. edx = np.maximum( 0, edx -1)
  60. ey = np.maximum( 0, ey -1)
  61. ex = np.maximum( 0, ex -1)
  62. return [dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph]
  63. def IoU(box, boxes):
  64. """Compute IoU between detect box and gt boxes
  65. Parameters:
  66. ----------
  67. box: numpy array , shape (5, ): x1, y1, x2, y2, score
  68. input box
  69. boxes: numpy array, shape (n, 4): x1, y1, x2, y2
  70. input ground truth boxes
  71. Returns:
  72. -------
  73. ovr: numpy.array, shape (n, )
  74. IoU
  75. """
  76. box_area = (box[ 2] - box[ 0] + 1) * (box[ 3] - box[ 1] + 1)
  77. area = (boxes[:, 2] - boxes[:, 0] + 1) * (boxes[:, 3] - boxes[:, 1] + 1)
  78. xx1 = np.maximum(box[ 0], boxes[:, 0])
  79. yy1 = np.maximum(box[ 1], boxes[:, 1])
  80. xx2 = np.minimum(box[ 2], boxes[:, 2])
  81. yy2 = np.minimum(box[ 3], boxes[:, 3])
  82. # compute the width and height of the bounding box
  83. w = np.maximum( 0, xx2 - xx1 + 1)
  84. h = np.maximum( 0, yy2 - yy1 + 1)
  85. inter = w * h
  86. ovr = inter / (box_area + area - inter)
  87. return ovr
  88. def rerec(bboxA):
  89. # convert bboxA to square
  90. w = bboxA[:, 2] - bboxA[:, 0]
  91. h = bboxA[:, 3] - bboxA[:, 1]
  92. l = np.maximum(w,h).T
  93. bboxA[:, 0] = bboxA[:, 0] + w* 0.5 - l* 0.5
  94. bboxA[:, 1] = bboxA[:, 1] + h* 0.5 - l* 0.5
  95. bboxA[:, 2: 4] = bboxA[:, 0: 2] + np.repeat([l], 2, axis = 0).T
  96. return bboxA
  97. def nms(boxes, threshold, type):
  98. """nms
  99. :boxes: [:,0:5]
  100. :threshold: 0.5 like
  101. :type: 'Min' or others
  102. :returns: TODO
  103. """
  104. if boxes.shape[ 0] == 0:
  105. return np.array([])
  106. x1 = boxes[:, 0]
  107. y1 = boxes[:, 1]
  108. x2 = boxes[:, 2]
  109. y2 = boxes[:, 3]
  110. s = boxes[:, 4]
  111. area = np.multiply(x2-x1+ 1, y2-y1+ 1)
  112. I = np.array(s.argsort()) # read s using I
  113. pick = [];
  114. while len(I) > 0:
  115. xx1 = np.maximum(x1[I[ -1]], x1[I[ 0: -1]])
  116. yy1 = np.maximum(y1[I[ -1]], y1[I[ 0: -1]])
  117. xx2 = np.minimum(x2[I[ -1]], x2[I[ 0: -1]])
  118. yy2 = np.minimum(y2[I[ -1]], y2[I[ 0: -1]])
  119. w = np.maximum( 0.0, xx2 - xx1 + 1)
  120. h = np.maximum( 0.0, yy2 - yy1 + 1)
  121. inter = w * h
  122. if type == 'Min':
  123. o = inter / np.minimum(area[I[ -1]], area[I[ 0: -1]])
  124. else:
  125. o = inter / (area[I[ -1]] + area[I[ 0: -1]] - inter)
  126. pick.append(I[ -1])
  127. I = I[np.where( o <= threshold)[ 0]]
  128. return pick
  129. def generateBoundingBox(map, reg, scale, t):
  130. stride = 2
  131. cellsize = 12
  132. map = map.T
  133. dx1 = reg[ 0,:,:].T
  134. dy1 = reg[ 1,:,:].T
  135. dx2 = reg[ 2,:,:].T
  136. dy2 = reg[ 3,:,:].T
  137. (x, y) = np.where(map >= t)
  138. yy = y
  139. xx = x
  140. score = map[x,y]
  141. reg = np.array([dx1[x,y], dy1[x,y], dx2[x,y], dy2[x,y]])
  142. if reg.shape[ 0] == 0:
  143. pass
  144. boundingbox = np.array([yy, xx]).T
  145. bb1 = np.fix((stride * (boundingbox) + 1) / scale).T # matlab index from 1, so with "boundingbox-1"
  146. bb2 = np.fix((stride * (boundingbox) + cellsize - 1 + 1) / scale).T # while python don't have to
  147. score = np.array([score])
  148. boundingbox_out = np.concatenate((bb1, bb2, score, reg), axis= 0)
  149. return boundingbox_out.T
  150. def drawBoxes(im, boxes):
  151. x1 = boxes[:, 0]
  152. y1 = boxes[:, 1]
  153. x2 = boxes[:, 2]
  154. y2 = boxes[:, 3]
  155. for i in range(x1.shape[ 0]):
  156. cv2.rectangle(im, (int(x1[i]), int(y1[i])), (int(x2[i]), int(y2[i])), ( 0, 255, 0), 1)
  157. return im
  158. def drawlandmark(im, points):
  159. for i in range(points.shape[ 0]):
  160. for j in range( 5):
  161. cv2.circle(im, (int(points[i][j]), int(points[i][j+ 5])), 2, ( 255, 0, 0))
  162. return im
  163. from time import time
  164. _tstart_stack = []
  165. def tic():
  166. _tstart_stack.append(time())
  167. def toc(fmt="Elapsed: %s s"):
  168. print fmt % (time()-_tstart_stack.pop())
  169. def detect_face(img, minsize, PNet, RNet, threshold, fastresize, factor):
  170. img2 = img.copy()
  171. factor_count = 0
  172. total_boxes = np.zeros(( 0, 9), np.float)
  173. points = []
  174. h = img.shape[ 0]
  175. w = img.shape[ 1]
  176. minl = min(h, w)
  177. img = img.astype(float)
  178. m = 12.0/minsize
  179. minl = minl*m
  180. # create scale pyramid
  181. scales = []
  182. while minl >= 12:
  183. scales.append(m * pow(factor, factor_count))
  184. minl *= factor
  185. factor_count += 1
  186. # first stage
  187. for scale in scales:
  188. hs = int(np.ceil(h*scale))
  189. ws = int(np.ceil(w*scale))
  190. if fastresize:
  191. im_data = (img -127.5)* 0.0078125 # [0,255] -> [-1,1]
  192. im_data = cv2.resize(im_data, (ws,hs)) # default is bilinear
  193. else:
  194. im_data = cv2.resize(img, (ws,hs)) # default is bilinear
  195. im_data = (im_data -127.5)* 0.0078125 # [0,255] -> [-1,1]
  196. #im_data = imResample(img, hs, ws); print "scale:", scale
  197. im_data = np.swapaxes(im_data, 0, 2)
  198. im_data = np.array([im_data], dtype = np.float)
  199. PNet.blobs[ 'data'].reshape( 1, 3, ws, hs)
  200. PNet.blobs[ 'data'].data[...] = im_data
  201. out = PNet.forward()
  202. boxes = generateBoundingBox(out[ 'prob1'][ 0, 1,:,:], out[ 'conv4-2'][ 0], scale, threshold[ 0])
  203. if boxes.shape[ 0] != 0:
  204. pick = nms(boxes, 0.5, 'Union')
  205. if len(pick) > 0 :
  206. boxes = boxes[pick, :]
  207. if boxes.shape[ 0] != 0:
  208. total_boxes = np.concatenate((total_boxes, boxes), axis= 0)
  209. #np.save('total_boxes_101.npy', total_boxes)
  210. #####
  211. # 1 #
  212. #####
  213. # print "[1]:",total_boxes.shape[0]
  214. #print total_boxes
  215. #return total_boxes, []
  216. numbox = total_boxes.shape[ 0]
  217. if numbox > 0:
  218. # nms
  219. pick = nms(total_boxes, 0.7, 'Union')
  220. total_boxes = total_boxes[pick, :]
  221. # print "[2]:",total_boxes.shape[0]
  222. # revise and convert to square
  223. regh = total_boxes[:, 3] - total_boxes[:, 1]
  224. regw = total_boxes[:, 2] - total_boxes[:, 0]
  225. t1 = total_boxes[:, 0] + total_boxes[:, 5]*regw
  226. t2 = total_boxes[:, 1] + total_boxes[:, 6]*regh
  227. t3 = total_boxes[:, 2] + total_boxes[:, 7]*regw
  228. t4 = total_boxes[:, 3] + total_boxes[:, 8]*regh
  229. t5 = total_boxes[:, 4]
  230. total_boxes = np.array([t1,t2,t3,t4,t5]).T
  231. total_boxes = rerec(total_boxes) # convert box to square
  232. # print "[4]:",total_boxes.shape[0]
  233. total_boxes[:, 0: 4] = np.fix(total_boxes[:, 0: 4])
  234. # print "[4.5]:",total_boxes.shape[0]
  235. #print total_boxes
  236. [dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph] = pad(total_boxes, w, h)
  237. numbox = total_boxes.shape[ 0]
  238. if numbox > 0:
  239. # second stage
  240. # construct input for RNet
  241. tempimg = np.zeros((numbox, 24, 24, 3)) # (24, 24, 3, numbox)
  242. for k in range(numbox):
  243. tmp = np.zeros((int(tmph[k]) + 1, int(tmpw[k]) + 1, 3))
  244. tmp[int(dy[k]):int(edy[k])+ 1, int(dx[k]):int(edx[k])+ 1] = img[int(y[k]):int(ey[k])+ 1, int(x[k]):int(ex[k])+ 1]
  245. #print "y,ey,x,ex", y[k], ey[k], x[k], ex[k]
  246. #print "tmp", tmp.shape
  247. tempimg[k,:,:,:] = cv2.resize(tmp, ( 24, 24))
  248. #print tempimg.shape
  249. #print tempimg[0,0,0,:]
  250. tempimg = (tempimg -127.5)* 0.0078125 # done in imResample function wrapped by python
  251. # RNet
  252. tempimg = np.swapaxes(tempimg, 1, 3)
  253. #print tempimg[0,:,0,0]
  254. RNet.blobs[ 'data'].reshape(numbox, 3, 24, 24)
  255. RNet.blobs[ 'data'].data[...] = tempimg
  256. out = RNet.forward()
  257. score = out[ 'prob1'][:, 1]
  258. #print 'score', score
  259. pass_t = np.where(score>threshold[ 1])[ 0]
  260. #print 'pass_t', pass_t
  261. score = np.array([score[pass_t]]).T
  262. total_boxes = np.concatenate( (total_boxes[pass_t, 0: 4], score), axis = 1)
  263. # print "[5]:",total_boxes.shape[0]
  264. #print total_boxes
  265. #print "1.5:",total_boxes.shape
  266. mv = out[ 'conv5-2'][pass_t, :].T
  267. #print "mv", mv
  268. if total_boxes.shape[ 0] > 0:
  269. pick = nms(total_boxes, 0.7, 'Union')
  270. # print 'pick', pick
  271. if len(pick) > 0:
  272. total_boxes = total_boxes[pick, :]
  273. # print "[6]:", total_boxes.shape[0]
  274. total_boxes = bbreg(total_boxes, mv[:, pick])
  275. # print "[7]:", total_boxes.shape[0]
  276. total_boxes = rerec(total_boxes)
  277. # print "[8]:", total_boxes.shape[0]
  278. return total_boxes
  279. def main():
  280. img_dir = "/home/xiao/code/mtcnn-caffe/prepare_data/WIDER_train/images/"
  281. imglistfile = "wider_face_train.txt"
  282. with open(imglistfile, 'r') as f:
  283. annotations = f.readlines()
  284. num = len(annotations)
  285. print "%d pics in total" % num
  286. neg_save_dir = "/media/xiao/软件/mtcnn/train/48/negative/"
  287. pos_save_dir = "/media/xiao/软件/mtcnn/train/48/positive/"
  288. part_save_dir = "/media/xiao/软件/mtcnn/train/48/part/"
  289. image_size = 48
  290. f1 = open( '/media/xiao/软件/mtcnn/train/48/pos_48.txt', 'w')
  291. f2 = open( '/media/xiao/软件/mtcnn/train/48/neg_48.txt', 'w')
  292. f3 = open( '/media/xiao/软件/mtcnn/train/48/part_48.txt', 'w')
  293. p_idx = 0 # positive
  294. n_idx = 0 # negative
  295. d_idx = 0 # dont care
  296. image_idx = 0
  297. minsize = 20
  298. caffe_model_path = "./model"
  299. threshold = [ 0.6, 0.7, 0.7]
  300. factor = 0.709
  301. caffe.set_mode_gpu()
  302. PNet = caffe.Net(caffe_model_path+ "/det1.prototxt", caffe_model_path+ "/det1.caffemodel", caffe.TEST)
  303. RNet = caffe.Net(caffe_model_path+ "/det2.prototxt", caffe_model_path+ "/det2.caffemodel", caffe.TEST)
  304. for annotation in annotations:
  305. # imgpath = imgpath.split('\n')[0]
  306. annotation = annotation.strip().split( ' ')
  307. bbox = map(float, annotation[ 1:])
  308. gts = np.array(bbox, dtype=np.float32).reshape( -1, 4)
  309. img_path = img_dir + annotation[ 0] + '.jpg'
  310. #print "######\n", img_path
  311. print image_idx
  312. image_idx += 1
  313. img = cv2.imread(img_path)
  314. img_matlab = img.copy()
  315. tmp = img_matlab[:,:, 2].copy()
  316. img_matlab[:,:, 2] = img_matlab[:,:, 0]
  317. img_matlab[:,:, 0] = tmp
  318. boundingboxes = detect_face(img_matlab, minsize, PNet, RNet, threshold, False, factor)
  319. #img = drawBoxes(img, boundingboxes)
  320. #cv2.imshow('img', img)
  321. #cv2.waitKey(1000)
  322. # generate positive,negative,part samples
  323. for box in boundingboxes:
  324. x_left, y_top, x_right, y_bottom, _ = box
  325. crop_w = x_right - x_left + 1
  326. crop_h = y_bottom - y_top + 1
  327. # ignore box that is too small or beyond image border
  328. if crop_w < image_size / 2 or crop_h < image_size / 2:
  329. continue
  330. if x_left < 0 or y_top < 0:
  331. continue
  332. # compute intersection over union(IoU) between current box and all gt boxes
  333. Iou = IoU(box, gts)
  334. cropped_im = img[int(y_top):int(y_bottom + 1) , int(x_left):int(x_right + 1) ]
  335. resized_im = cv2.resize(cropped_im, (image_size, image_size), interpolation=cv2.INTER_LINEAR)
  336. #try:
  337. # resized_im = cv2.resize(cropped_im, (image_size, image_size), interpolation=cv2.INTER_LINEAR)
  338. #except Exception as e:
  339. # print " 1 "
  340. # print e
  341. # save negative images and write label, 负样本
  342. if np.max(Iou) < 0.3:
  343. # Iou with all gts must below 0.3
  344. save_file = os.path.join(neg_save_dir, "%s.jpg" % n_idx)
  345. f2.write( "%s/negative/%s.jpg" % (image_size, n_idx) + ' 0')
  346. f2.write( " -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1\n")
  347. cv2.imwrite(save_file, resized_im)
  348. n_idx += 1
  349. else:
  350. # find gt_box with the highest iou
  351. idx = np.argmax(Iou)
  352. assigned_gt = gts[idx]
  353. x1, y1, x2, y2 = assigned_gt
  354. # compute bbox reg label,offset_x1,offset_y1相对于左上角; offset_x2,offset_y2相对于右上角
  355. offset_x1 = (x1 - x_left) / float(crop_w)
  356. offset_y1 = (y1 - y_top) / float(crop_h)
  357. # offset_x2 = (x2 - x_left) / float(crop_w)
  358. # offset_y2 = (y2 - y_top) / float(crop_h)
  359. offset_x2 = (x2 - x_right) / float(crop_w)
  360. offset_y2 = (y2 - y_bottom )/ float(crop_h)
  361. # save positive and part-face images and write labels, 正样本
  362. if np.max(Iou) >= 0.65:
  363. save_file = os.path.join(pos_save_dir, "%s.jpg" % p_idx)
  364. f1.write( "%s/positive/%s.jpg" % (image_size, p_idx) + ' 1 %.6f %.6f %.6f %.6f' % (offset_x1, offset_y1, offset_x2, offset_y2))
  365. f1.write( " -1 -1 -1 -1 -1 -1 -1 -1 -1 -1\n")
  366. cv2.imwrite(save_file, resized_im)
  367. p_idx += 1
  368. # part 样本
  369. elif np.max(Iou) >= 0.4:
  370. save_file = os.path.join(part_save_dir, "%s.jpg" % d_idx)
  371. f3.write( "%s/part/%s.jpg" % (image_size, d_idx) + ' -1 %.6f %.6f %.6f %.6f' % (offset_x1, offset_y1, offset_x2, offset_y2))
  372. f3.write( " -1 -1 -1 -1 -1 -1 -1 -1 -1 -1\n")
  373. cv2.imwrite(save_file, resized_im)
  374. d_idx += 1
  375. f.close()
  376. f1.close()
  377. f2.close()
  378. f3.close()
  379. if __name__ == "__main__":
  380. main()


2,生成landmark样本,用作者的net1->net2生成bbox, 根据5个landmark是否都在bbox中作为判别条件:

代码如下:



 
 
  1. #!/usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. import _init_paths
  4. import caffe
  5. import cv2
  6. import numpy as np
  7. #from python_wrapper import *
  8. import os
  9. def bbreg(boundingbox, reg):
  10. reg = reg.T
  11. # calibrate bouding boxes
  12. if reg.shape[ 1] == 1:
  13. print "reshape of reg"
  14. pass # reshape of reg
  15. w = boundingbox[:, 2] - boundingbox[:, 0] + 1
  16. h = boundingbox[:, 3] - boundingbox[:, 1] + 1
  17. bb0 = boundingbox[:, 0] + reg[:, 0]*w
  18. bb1 = boundingbox[:, 1] + reg[:, 1]*h
  19. bb2 = boundingbox[:, 2] + reg[:, 2]*w
  20. bb3 = boundingbox[:, 3] + reg[:, 3]*h
  21. boundingbox[:, 0: 4] = np.array([bb0, bb1, bb2, bb3]).T
  22. #print "bb", boundingbox
  23. return boundingbox
  24. def pad(boxesA, w, h):
  25. boxes = boxesA.copy() # shit, value parameter!!!
  26. tmph = boxes[:, 3] - boxes[:, 1] + 1
  27. tmpw = boxes[:, 2] - boxes[:, 0] + 1
  28. numbox = boxes.shape[ 0]
  29. dx = np.ones(numbox)
  30. dy = np.ones(numbox)
  31. edx = tmpw
  32. edy = tmph
  33. x = boxes[:, 0: 1][:, 0]
  34. y = boxes[:, 1: 2][:, 0]
  35. ex = boxes[:, 2: 3][:, 0]
  36. ey = boxes[:, 3: 4][:, 0]
  37. tmp = np.where(ex > w)[ 0]
  38. if tmp.shape[ 0] != 0:
  39. edx[tmp] = -ex[tmp] + w -1 + tmpw[tmp]
  40. ex[tmp] = w -1
  41. tmp = np.where(ey > h)[ 0]
  42. if tmp.shape[ 0] != 0:
  43. edy[tmp] = -ey[tmp] + h -1 + tmph[tmp]
  44. ey[tmp] = h -1
  45. tmp = np.where(x < 1)[ 0]
  46. if tmp.shape[ 0] != 0:
  47. dx[tmp] = 2 - x[tmp]
  48. x[tmp] = np.ones_like(x[tmp])
  49. tmp = np.where(y < 1)[ 0]
  50. if tmp.shape[ 0] != 0:
  51. dy[tmp] = 2 - y[tmp]
  52. y[tmp] = np.ones_like(y[tmp])
  53. # for python index from 0, while matlab from 1
  54. dy = np.maximum( 0, dy -1)
  55. dx = np.maximum( 0, dx -1)
  56. y = np.maximum( 0, y -1)
  57. x = np.maximum( 0, x -1)
  58. edy = np.maximum( 0, edy -1)
  59. edx = np.maximum( 0, edx -1)
  60. ey = np.maximum( 0, ey -1)
  61. ex = np.maximum( 0, ex -1)
  62. return [dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph]
  63. def IoU(box, boxes):
  64. """Compute IoU between detect box and gt boxes
  65. Parameters:
  66. ----------
  67. box: numpy array , shape (5, ): x1, y1, x2, y2, score
  68. input box
  69. boxes: numpy array, shape (n, 4): x1, y1, x2, y2
  70. input ground truth boxes
  71. Returns:
  72. -------
  73. ovr: numpy.array, shape (n, )
  74. IoU
  75. """
  76. box_area = (box[ 2] - box[ 0] + 1) * (box[ 3] - box[ 1] + 1)
  77. area = (boxes[:, 2] - boxes[:, 0] + 1) * (boxes[:, 3] - boxes[:, 1] + 1)
  78. xx1 = np.maximum(box[ 0], boxes[:, 0])
  79. yy1 = np.maximum(box[ 1], boxes[:, 1])
  80. xx2 = np.minimum(box[ 2], boxes[:, 2])
  81. yy2 = np.minimum(box[ 3], boxes[:, 3])
  82. # compute the width and height of the bounding box
  83. w = np.maximum( 0, xx2 - xx1 + 1)
  84. h = np.maximum( 0, yy2 - yy1 + 1)
  85. inter = w * h
  86. ovr = inter / (box_area + area - inter)
  87. return ovr
  88. def rerec(bboxA):
  89. # convert bboxA to square
  90. w = bboxA[:, 2] - bboxA[:, 0]
  91. h = bboxA[:, 3] - bboxA[:, 1]
  92. l = np.maximum(w,h).T
  93. bboxA[:, 0] = bboxA[:, 0] + w* 0.5 - l* 0.5
  94. bboxA[:, 1] = bboxA[:, 1] + h* 0.5 - l* 0.5
  95. bboxA[:, 2: 4] = bboxA[:, 0: 2] + np.repeat([l], 2, axis = 0).T
  96. return bboxA
  97. def nms(boxes, threshold, type):
  98. """nms
  99. :boxes: [:,0:5]
  100. :threshold: 0.5 like
  101. :type: 'Min' or others
  102. :returns: TODO
  103. """
  104. if boxes.shape[ 0] == 0:
  105. return np.array([])
  106. x1 = boxes[:, 0]
  107. y1 = boxes[:, 1]
  108. x2 = boxes[:, 2]
  109. y2 = boxes[:, 3]
  110. s = boxes[:, 4]
  111. area = np.multiply(x2-x1+ 1, y2-y1+ 1)
  112. I = np.array(s.argsort()) # read s using I
  113. pick = [];
  114. while len(I) > 0:
  115. xx1 = np.maximum(x1[I[ -1]], x1[I[ 0: -1]])
  116. yy1 = np.maximum(y1[I[ -1]], y1[I[ 0: -1]])
  117. xx2 = np.minimum(x2[I[ -1]], x2[I[ 0: -1]])
  118. yy2 = np.minimum(y2[I[ -1]], y2[I[ 0: -1]])
  119. w = np.maximum( 0.0, xx2 - xx1 + 1)
  120. h = np.maximum( 0.0, yy2 - yy1 + 1)
  121. inter = w * h
  122. if type == 'Min':
  123. o = inter / np.minimum(area[I[ -1]], area[I[ 0: -1]])
  124. else:
  125. o = inter / (area[I[ -1]] + area[I[ 0: -1]] - inter)
  126. pick.append(I[ -1])
  127. I = I[np.where( o <= threshold)[ 0]]
  128. return pick
  129. def generateBoundingBox(map, reg, scale, t):
  130. stride = 2
  131. cellsize = 12
  132. map = map.T
  133. dx1 = reg[ 0,:,:].T
  134. dy1 = reg[ 1,:,:].T
  135. dx2 = reg[ 2,:,:].T
  136. dy2 = reg[ 3,:,:].T
  137. (x, y) = np.where(map >= t)
  138. yy = y
  139. xx = x
  140. score = map[x,y]
  141. reg = np.array([dx1[x,y], dy1[x,y], dx2[x,y], dy2[x,y]])
  142. if reg.shape[ 0] == 0:
  143. pass
  144. boundingbox = np.array([yy, xx]).T
  145. bb1 = np.fix((stride * (boundingbox) + 1) / scale).T # matlab index from 1, so with "boundingbox-1"
  146. bb2 = np.fix((stride * (boundingbox) + cellsize - 1 + 1) / scale).T # while python don't have to
  147. score = np.array([score])
  148. boundingbox_out = np.concatenate((bb1, bb2, score, reg), axis= 0)
  149. return boundingbox_out.T
  150. def drawBoxes(im, boxes):
  151. x1 = boxes[:, 0]
  152. y1 = boxes[:, 1]
  153. x2 = boxes[:, 2]
  154. y2 = boxes[:, 3]
  155. for i in range(x1.shape[ 0]):
  156. cv2.rectangle(im, (int(x1[i]), int(y1[i])), (int(x2[i]), int(y2[i])), ( 0, 255, 0), 1)
  157. return im
  158. def drawBoxes_align(im, boxe):
  159. x1 = boxe[ 0]
  160. y1 = boxe[ 1]
  161. x2 = boxe[ 2]
  162. y2 = boxe[ 3]
  163. cv2.rectangle(im, (int(x1), int(y1)), (int(x2), int(y2)), ( 0, 255, 0), 1)
  164. return im
  165. def drawlandmark(im, points):
  166. for i in range(points.shape[ 0]):
  167. for j in range( 5):
  168. cv2.circle(im, (int(points[i][j]), int(points[i][j+ 5])), 2, ( 255, 0, 0))
  169. return im
  170. def drawlandmark_align(im, point):
  171. for j in range( 5):
  172. cv2.circle(im, (int(point[j* 2]), int(point[j* 2+ 1])), 2, ( 255, 0, 0))
  173. return im
  174. from time import time
  175. _tstart_stack = []
  176. def tic():
  177. _tstart_stack.append(time())
  178. def toc(fmt="Elapsed: %s s"):
  179. print fmt % (time()-_tstart_stack.pop())
  180. def detect_face(img, minsize, PNet, RNet, threshold, fastresize, factor):
  181. img2 = img.copy()
  182. factor_count = 0
  183. total_boxes = np.zeros(( 0, 9), np.float)
  184. points = []
  185. h = img.shape[ 0]
  186. w = img.shape[ 1]
  187. minl = min(h, w)
  188. img = img.astype(float)
  189. m = 12.0/minsize
  190. minl = minl*m
  191. # create scale pyramid
  192. scales = []
  193. while minl >= 12:
  194. scales.append(m * pow(factor, factor_count))
  195. minl *= factor
  196. factor_count += 1
  197. # first stage
  198. for scale in scales:
  199. hs = int(np.ceil(h*scale))
  200. ws = int(np.ceil(w*scale))
  201. if fastresize:
  202. im_data = (img -127.5)* 0.0078125 # [0,255] -> [-1,1]
  203. im_data = cv2.resize(im_data, (ws,hs)) # default is bilinear
  204. else:
  205. im_data = cv2.resize(img, (ws,hs)) # default is bilinear
  206. im_data = (im_data -127.5)* 0.0078125 # [0,255] -> [-1,1]
  207. #im_data = imResample(img, hs, ws); print "scale:", scale
  208. im_data = np.swapaxes(im_data, 0, 2)
  209. im_data = np.array([im_data], dtype = np.float)
  210. PNet.blobs[ 'data'].reshape( 1, 3, ws, hs)
  211. PNet.blobs[ 'data'].data[...] = im_data
  212. out = PNet.forward()
  213. boxes = generateBoundingBox(out[ 'prob1'][ 0, 1,:,:], out[ 'conv4-2'][ 0], scale, threshold[ 0])
  214. if boxes.shape[ 0] != 0:
  215. pick = nms(boxes, 0.5, 'Union')
  216. if len(pick) > 0 :
  217. boxes = boxes[pick, :]
  218. if boxes.shape[ 0] != 0:
  219. total_boxes = np.concatenate((total_boxes, boxes), axis= 0)
  220. #np.save('total_boxes_101.npy', total_boxes)
  221. #####
  222. # 1 #
  223. #####
  224. # print "[1]:",total_boxes.shape[0]
  225. #print total_boxes
  226. #return total_boxes, []
  227. numbox = total_boxes.shape[ 0]
  228. if numbox > 0:
  229. # nms
  230. pick = nms(total_boxes, 0.7, 'Union')
  231. total_boxes = total_boxes[pick, :]
  232. # print "[2]:",total_boxes.shape[0]
  233. # revise and convert to square
  234. regh = total_boxes[:, 3] - total_boxes[:, 1]
  235. regw = total_boxes[:, 2] - total_boxes[:, 0]
  236. t1 = total_boxes[:, 0] + total_boxes[:, 5]*regw
  237. t2 = total_boxes[:, 1] + total_boxes[:, 6]*regh
  238. t3 = total_boxes[:, 2] + total_boxes[:, 7]*regw
  239. t4 = total_boxes[:, 3] + total_boxes[:, 8]*regh
  240. t5 = total_boxes[:, 4]
  241. total_boxes = np.array([t1,t2,t3,t4,t5]).T
  242. total_boxes = rerec(total_boxes) # convert box to square
  243. # print "[4]:",total_boxes.shape[0]
  244. total_boxes[:, 0: 4] = np.fix(total_boxes[:, 0: 4])
  245. # print "[4.5]:",total_boxes.shape[0]
  246. #print total_boxes
  247. [dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph] = pad(total_boxes, w, h)
  248. numbox = total_boxes.shape[ 0]
  249. if numbox > 0:
  250. # second stage
  251. # construct input for RNet
  252. tempimg = np.zeros((numbox, 24, 24, 3)) # (24, 24, 3, numbox)
  253. for k in range(numbox):
  254. tmp = np.zeros((int(tmph[k]) + 1, int(tmpw[k]) + 1, 3))
  255. tmp[int(dy[k]):int(edy[k])+ 1, int(dx[k]):int(edx[k])+ 1] = img[int(y[k]):int(ey[k])+ 1, int(x[k]):int(ex[k])+ 1]
  256. #print "y,ey,x,ex", y[k], ey[k], x[k], ex[k]
  257. #print "tmp", tmp.shape
  258. tempimg[k,:,:,:] = cv2.resize(tmp, ( 24, 24))
  259. #print tempimg.shape
  260. #print tempimg[0,0,0,:]
  261. tempimg = (tempimg -127.5)* 0.0078125 # done in imResample function wrapped by python
  262. # RNet
  263. tempimg = np.swapaxes(tempimg, 1, 3)
  264. #print tempimg[0,:,0,0]
  265. RNet.blobs[ 'data'].reshape(numbox, 3, 24, 24)
  266. RNet.blobs[ 'data'].data[...] = tempimg
  267. out = RNet.forward()
  268. score = out[ 'prob1'][:, 1]
  269. #print 'score', score
  270. pass_t = np.where(score>threshold[ 1])[ 0]
  271. #print 'pass_t', pass_t
  272. score = np.array([score[pass_t]]).T
  273. total_boxes = np.concatenate( (total_boxes[pass_t, 0: 4], score), axis = 1)
  274. # print "[5]:",total_boxes.shape[0]
  275. #print total_boxes
  276. #print "1.5:",total_boxes.shape
  277. mv = out[ 'conv5-2'][pass_t, :].T
  278. #print "mv", mv
  279. if total_boxes.shape[ 0] > 0:
  280. pick = nms(total_boxes, 0.7, 'Union')
  281. # print 'pick', pick
  282. if len(pick) > 0:
  283. total_boxes = total_boxes[pick, :]
  284. # print "[6]:", total_boxes.shape[0]
  285. total_boxes = bbreg(total_boxes, mv[:, pick])
  286. # print "[7]:", total_boxes.shape[0]
  287. total_boxes = rerec(total_boxes)
  288. # print "[8]:", total_boxes.shape[0]
  289. return total_boxes
  290. def main():
  291. img_dir = "/media/xiao/学习/face_alignment/data/CelebA/Img/img_celeba.7z/img_celeba/"
  292. anno_file = "celebA_bbox_landmark.txt"
  293. with open(anno_file, 'r') as f:
  294. annotations = f.readlines()
  295. num = len(annotations)
  296. print "%d pics in total" % num
  297. # 图片大小为48*48
  298. image_size = 48
  299. # landmark_save_dir = "48/landmark/"
  300. landmark_save_dir = "/media/xiao/软件/mtcnn/train/48/landmark/"
  301. # save_dir = "./" + str(image_size)
  302. f1 = open( '/media/xiao/软件/mtcnn/train/48/landmark_48.txt', 'w')
  303. l_idx = 0 # landmark
  304. image_idx = 0
  305. minsize = 40
  306. caffe_model_path = "./model"
  307. threshold = [ 0.6, 0.7, 0.7]
  308. factor = 0.709
  309. caffe.set_mode_gpu()
  310. PNet = caffe.Net(caffe_model_path+ "/det1.prototxt", caffe_model_path+ "/det1.caffemodel", caffe.TEST)
  311. RNet = caffe.Net(caffe_model_path+ "/det2.prototxt", caffe_model_path+ "/det2.caffemodel", caffe.TEST)
  312. for annotation in annotations:
  313. # imgpath = imgpath.split('\n')[0]
  314. annotation = annotation.strip().split( ' ')
  315. im_path = annotation[ 0]
  316. # bbox = map(float, annotation[1:-10])
  317. pts = map(float, annotation[ -10:])
  318. # boxes = np.array(bbox, dtype=np.float32).reshape(-1, 4)
  319. im_path = img_dir + im_path
  320. backupPts = pts[:]
  321. #print "######\n", img_path
  322. print image_idx
  323. image_idx += 1
  324. img = cv2.imread(im_path)
  325. img_matlab = img.copy()
  326. tmp = img_matlab[:,:, 2].copy()
  327. img_matlab[:,:, 2] = img_matlab[:,:, 0]
  328. img_matlab[:,:, 0] = tmp
  329. # 用作者提供的net1->net2生成人脸框
  330. boundingboxes = detect_face(img_matlab, minsize, PNet, RNet, threshold, False, factor)
  331. #img = drawBoxes(img, boundingboxes)
  332. #cv2.imshow('img', img)
  333. #cv2.waitKey(1000)
  334. # generate landmark samples
  335. for box in boundingboxes:
  336. x_left, y_top, x_right, y_bottom, _ = box # 得到人脸bbox坐标
  337. crop_w = x_right - x_left + 1
  338. crop_h = y_bottom - y_top + 1
  339. # ignore box that is too small or beyond image border
  340. if crop_w < image_size / 2 or crop_h < image_size / 2:
  341. continue
  342. if x_left < 0 or y_top < 0:
  343. continue
  344. # 当五个landmark都在bbox中时,这样的样本作为我们的landmark训练样本
  345. if pts[ 0] < x_left or pts[ 0] > x_right:
  346. continue
  347. if pts[ 2] < x_left or pts[ 2] > x_right:
  348. continue
  349. if pts[ 4] < x_left or pts[ 4] > x_right:
  350. continue
  351. if pts[ 6] < x_left or pts[ 6] > x_right:
  352. continue
  353. if pts[ 8] < x_left or pts[ 8] > x_right:
  354. continue
  355. if pts[ 1] < y_top or pts[ 1] > y_bottom:
  356. continue
  357. if pts[ 3] < y_top or pts[ 3] > y_bottom:
  358. continue
  359. if pts[ 5] < y_top or pts[ 5] > y_bottom:
  360. continue
  361. if pts[ 7] < y_top or pts[ 7] > y_bottom:
  362. continue
  363. if pts[ 9] < y_top or pts[ 9] > y_bottom:
  364. continue
  365. # show image
  366. #img1 = drawBoxes_align(img, box)
  367. #img1 = drawlandmark_align(img1, pts)
  368. #cv2.imshow('img', img1)
  369. #cv2.waitKey(1000)
  370. # our method, x0,y0,x1,y1,x2,y2,x3,y3,x4,y4
  371. #for k in range(len(pts) / 2):
  372. # pts[k * 2] = (pts[k * 2] - x_left) / float(crop_w);
  373. # pts[k * 2 + 1] = (pts[k * 2 + 1] - y_top) / float(crop_h);
  374. #author method, x0,x1,x2,x3,x4,y0,y1,y2,y3,y4
  375. ptsss = pts[:]
  376. # landmark标注为其相对于crop区域左上角的归一化值
  377. for k in range(len(ptsss) / 2):
  378. pts[k] = (ptsss[k * 2] - x_left) / float(crop_w);
  379. pts[ 5+k] = (ptsss[k * 2 + 1] - y_top) / float(crop_h);
  380. # 将crop区域进行resize
  381. cropped_im = img[int(y_top):int(y_bottom + 1), int(x_left):int(x_right + 1)]
  382. resized_im = cv2.resize(cropped_im, (image_size, image_size), interpolation=cv2.INTER_LINEAR)
  383. # box_ = box.reshape(1, -1)
  384. # 写图片名,class坐标,bbox坐标。
  385. save_file = os.path.join(landmark_save_dir, "%s.jpg" % l_idx)
  386. f1.write(str(image_size) + "/landmark/%s.jpg" % l_idx + ' -1 -1 -1 -1 -1')
  387. # 写landmark坐标并保存图片
  388. for k in range(len(pts)):
  389. f1.write( " %f" % pts[k])
  390. f1.write( "\n")
  391. cv2.imwrite(save_file, resized_im)
  392. l_idx += 1
  393. f.close()
  394. f1.close()
  395. if __name__ == "__main__":
  396. main()




  • 1
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值