dlunion caffe加入hard examples 训练 R-Net,O-Net

上一篇提到，dlunion作者的caffe训练MTCNN并没有加入hard examples，所里这里将加入hard examples 训练R-Net和O-Net.

生成R-Net的训练数据

执行

python create_hard1.py

create_hard1.py的代码如下：

import sys import tools import caffe import cv2 import numpy as np import os from utils import * deploy = '../12/det1.prototxt' caffemodel = '../12/det1.caffemodel' net_12 = caffe.Net(deploy, caffemodel, caffe.TEST) def view_bar(num, total): rate = float(num) / total rate_num = int(rate * 100) r = '\r[%s%s]%d%% (%d/%d)' % ("#"*rate_num, " "*(100-rate_num), rate_num, num, total) sys.stdout.write(r) sys.stdout.flush() def detectFace(img_path, threshold): img = cv2.imread(img_path) caffe_img = img.copy()-128 origin_h, origin_w, ch = caffe_img.shape scales = tools.calculateScales(img) out = [] for scale in scales: hs = int(origin_h*scale) ws = int(origin_w*scale) scale_img = cv2.resize(caffe_img, (ws, hs)) scale_img = np.swapaxes(scale_img, 0, 2) net_12.blobs['data'].reshape(1, 3, ws, hs) net_12.blobs['data'].data[...] = scale_img caffe.set_device(0) caffe.set_mode_gpu() out_ = net_12.forward() out.append(out_) image_num = len(scales) rectangles = [] for i in range(image_num): cls_prob = out[i]['prob1'][0][1] roi = out[i]['conv4-2'][0] out_h, out_w = cls_prob.shape out_side = max(out_h, out_w) rectangle = tools.detect_face_12net(cls_prob, roi, out_side, 1/scales[i], origin_w, origin_h, threshold[0]) rectangles.extend(rectangle) return rectangles anno_file = './wider_face_train.txt' im_dir = "./WIDER_train/images/" neg_save_dir = "./24/negative" pos_save_dir = "./24/positive" part_save_dir = "./24/part" image_size = 24 f1 = open('./24/pos_24.txt', 'a') f2 = open('./24/neg_24.txt', 'a') f3 = open('./24/part_24.txt', 'a') threshold = [0.6, 0.6, 0.7] with open(anno_file, 'r') as f: annotations = f.readlines() num = len(annotations) print("%d pics in total" % num) p_idx = 0 # positive n_idx = 0 # negative d_idx = 0 # dont care image_idx = 0 for annotation in annotations: annotation = annotation.strip().split(' ') bbox = list(map(float, annotation[1:])) gts = np.array(bbox, dtype=np.float32).reshape(-1, 4) img_path = im_dir + annotation[0] + '.jpg' rectangles = detectFace(img_path, threshold) img = cv2.imread(img_path) image_idx += 1 view_bar(image_idx,num) for box in rectangles: x_left, y_top, x_right, y_bottom, _ = box crop_w = x_right - x_left + 1 crop_h = y_bottom - y_top + 1 # ignore box that is too small or beyond image border if crop_w < image_size or crop_h < image_size : continue # compute intersection over union(IoU) between current box and all gt boxes Iou = IoU(box, gts) cropped_im = img[y_top:y_bottom + 1, x_left:x_right + 1] resized_im = cv2.resize(cropped_im, (image_size, image_size), interpolation=cv2.INTER_LINEAR) # save negative images and write label if np.max(Iou) < 0.3: # Iou with all gts must below 0.3 save_file = os.path.join(neg_save_dir, "%s.jpg" % n_idx) f2.write("%s/negative/%s" % (image_size, n_idx) + ' 0\n') cv2.imwrite(save_file, resized_im) n_idx += 1 else: # find gt_box with the highest iou idx = np.argmax(Iou) assigned_gt = gts[idx] x1, y1, x2, y2 = assigned_gt # compute bbox reg label offset_x1 = (x1 - x_left) / float(crop_w) offset_y1 = (y1 - y_top) / float(crop_h) offset_x2 = (x2 - x_right) / float(crop_w) offset_y2 = (y2 - y_bottom) / float(crop_h) # save positive and part-face images and write labels if np.max(Iou) >= 0.65: save_file = os.path.join(pos_save_dir, "%s.jpg" % p_idx) f1.write("%s/positive/%s" % (image_size, p_idx) + ' 1 %.2f %.2f %.2f %.2f\n' % (offset_x1, offset_y1, offset_x2, offset_y2)) cv2.imwrite(save_file, resized_im) p_idx += 1 elif np.max(Iou) >= 0.4: save_file = os.path.join(part_save_dir, "%s.jpg" % d_idx) f3.write("%s/part/%s" % (image_size, d_idx) + ' -1 %.2f %.2f %.2f %.2f\n' % (offset_x1, offset_y1, offset_x2, offset_y2)) cv2.imwrite(save_file, resized_im) d_idx += 1 f1.close() f2.close() f3.close()

其中会导入tools.py，这个脚本的代码如下：

import sys from operator import itemgetter import numpy as np import cv2 ''' Function: calculate Intersect of Union Input: rect_1: 1st rectangle rect_2: 2nd rectangle Output: IoU ''' def IoU(rect_1, rect_2): x11 = rect_1[0] # first rectangle top left x y11 = rect_1[1] # first rectangle top left y x12 = rect_1[2] # first rectangle bottom right x y12 = rect_1[3] # first rectangle bottom right y x21 = rect_2[0] # second rectangle top left x y21 = rect_2[1] # second rectangle top left y x22 = rect_2[2] # second rectangle bottom right x y22 = rect_2[3] # second rectangle bottom right y x_overlap = max(0, min(x12,x22) -max(x11,x21)) y_overlap = max(0, min(y12,y22) -max(y11,y21)) intersection = x_overlap * y_overlap union = (x12-x11) * (y12-y11) + (x22-x21) * (y22-y21) - intersection if union == 0: return 0 else: return float(intersection) / union ''' Function: calculate Intersect of Min area Input: rect_1: 1st rectangle rect_2: 2nd rectangle Output: IoM ''' def IoM(rect_1, rect_2): x11 = rect_1[0] # first rectangle top left x y11 = rect_1[1] # first rectangle top left y x12 = rect_1[2] # first rectangle bottom right x y12 = rect_1[3] # first rectangle bottom right y x21 = rect_2[0] # second rectangle top left x y21 = rect_2[1] # second rectangle top left y x22 = rect_2[2] # second rectangle bottom right x y22 = rect_2[3] # second rectangle bottom right y x_overlap = max(0, min(x12,x22) -max(x11,x21)) y_overlap = max(0, min(y12,y22) -max(y11,y21)) intersection = x_overlap * y_overlap rect1_area = (y12 - y11) * (x12 - x11) rect2_area = (y22 - y21) * (x22 - x21) min_area = min(rect1_area, rect2_area) return float(intersection) / min_area ''' Function: apply NMS(non-maximum suppression) on ROIs in same scale Input: rectangles: rectangles[i][0:3] is the position, rectangles[i][4] is scale, rectangles[i][5] is score Output: rectangles: same as input ''' def NMS(rectangles, threshold, type): sorted(rectangles, key=itemgetter(4), reverse=True) result_rectangles = rectangles number_of_rects = len(result_rectangles) cur_rect = 0 while cur_rect < number_of_rects : rects_to_compare = number_of_rects - cur_rect - 1 cur_rect_to_compare = cur_rect + 1 while rects_to_compare > 0: score = 0 if type == 'iou': score = IoU(result_rectangles[cur_rect], result_rectangles[cur_rect_to_compare]) else: score = IoM(result_rectangles[cur_rect], result_rectangles[cur_rect_to_compare]) if score >= threshold: del result_rectangles[cur_rect_to_compare] # delete the rectangle number_of_rects -= 1 else: cur_rect_to_compare += 1 # skip to next rectangle rects_to_compare -= 1 cur_rect += 1 # finished comparing for current rectangle return result_rectangles def detect_face_12net(cls_prob,roi,out_side,scale,width,height,threshold): in_side = 2*out_side+11 stride = 0 if out_side != 1: stride = float(in_side-12)/(out_side-1) boundingBox = [] for (x, y), prob in np.ndenumerate(cls_prob): if(prob >= threshold): original_x1 = int((stride*x + 1)*scale) original_y1 = int((stride*y + 1)*scale) original_w = int((12.0 - 1)*scale) original_h = int((12.0 - 1)*scale) original_x2 = original_x1 + original_w original_y2 = original_y1 + original_h rect = [] x1 = int(round(max(0, original_x1 + original_w * roi[0][x][y]))) y1 = int(round(max(0, original_y1 + original_h * roi[1][x][y]))) x2 = int(round(min(width, original_x2 + original_w * roi[2][x][y]))) y2 = int(round(min(height, original_y2 + original_h * roi[3][x][y]))) if x2 > x1 and y2 > y1: rect = [x1, y1, x2, y2, prob] boundingBox.append(rect) return NMS(boundingBox, 0.5, 'iou') def filter_face_24net(cls_prob,roi,rectangles,width,height,threshold): boundingBox = [] rect_num = len(rectangles) for i in range(rect_num): if cls_prob[i][1] > threshold: original_w = rectangles[i][2]-rectangles[i][0]+1 original_h = rectangles[i][3]-rectangles[i][1]+1 x1 = int(round(max(0, rectangles[i][0] + original_w * roi[i][0]))) y1 = int(round(max(0, rectangles[i][1] + original_h * roi[i][1]))) x2 = int(round(min(width, rectangles[i][2] + original_w * roi[i][2]))) y2 = int(round(min(height, rectangles[i][3] + original_h * roi[i][3]))) if x2 > x1 and y2 > y1: rect = [x1, y1, x2, y2, cls_prob[i][1]] boundingBox.append(rect) return NMS(boundingBox,0.7,'iou') def filter_face_48net(cls_prob, roi, pts, rectangles, width, height, threshold): boundingBox = [] rect_num = len(rectangles) for i in range(rect_num): if cls_prob[i][1]>threshold: rect = [rectangles[i][0], rectangles[i][1], rectangles[i][2], rectangles[i][3], cls_prob[i][1], roi[i][0], roi[i][1], roi[i][2], roi[i][3], pts[i][0], pts[i][5], pts[i][1], pts[i][6], pts[i][2], pts[i][7], pts[i][3], pts[i][8], pts[i][4], pts[i][9]] boundingBox.append(rect) rectangles = NMS(boundingBox, 0.7, 'iom') rect = [] for rectangle in rectangles: roi_w = rectangle[2]-rectangle[0]+1 roi_h = rectangle[3]-rectangle[1]+1 x1 = round(max(0, rectangle[0]+rectangle[5]*roi_w)) y1 = round(max(0, rectangle[1]+rectangle[6]*roi_h)) x2 = round(min(width, rectangle[2]+rectangle[7]*roi_w)) y2 = round(min(height, rectangle[3]+rectangle[8]*roi_h)) pt0 = rectangle[9]*roi_w + rectangle[0] - 1 pt1 = rectangle[10]*roi_h + rectangle[1] - 1 pt2 = rectangle[11]*roi_w + rectangle[0] - 1 pt3 = rectangle[12]*roi_h + rectangle[1] - 1 pt4 = rectangle[13]*roi_w + rectangle[0] - 1 pt5 = rectangle[14]*roi_h + rectangle[1] - 1 pt6 = rectangle[15]*roi_w + rectangle[0] - 1 pt7 = rectangle[16]*roi_h + rectangle[1] - 1 pt8 = rectangle[17]*roi_w + rectangle[0] - 1 pt9 = rectangle[18]*roi_h + rectangle[1] - 1 score = rectangle[4] rect_ = np.round([x1, y1, x2, y2, pt0, pt1, pt2, pt3, pt4, pt5, pt6, pt7, pt8, pt9]).astype(int) rect_ = np.append(rect_, score) rect.append(rect_) return rect def calculateScales(img): caffe_img = img.copy() h, w, ch = caffe_img.shape pr_scale = 1000.0/max(h, w) w = int(w*pr_scale) h = int(h*pr_scale) scales = [] factor = 0.7937 factor_count = 0 minl = min(h, w) while minl >= 12: scales.append(pr_scale*pow(factor, factor_count)) minl *= factor factor_count += 1 return scales

生成O-Net的hard-examples的代码类似，只需要把create_hard1.py中的12修改成24。 生成难样本之后的数据预处理和训练步骤与训练P-Net的步骤类似，具体参见上一篇。