预测部分由两步组成： 1.由特征图获得全部的预测框 2.NMS对所有的预测框进行处理
该部分计算每个预测框的中心点坐标、wh、是否有目标的概率，目标类别的概率四个值。
def decode_feature(feature, anchor, width, height): shape = tf.shape(feature) shape = tf.cast(shape, tf.float32) # [batch_size, 13, 13, 3, 5+class_num] yi_pred = tf.reshape(feature, [shape[0], shape[1], shape[2], 3, -1]) # shape : [batch_size,13,13,3,2] [batch_size,13,13,3,2] [batch_size,13,13,3,1] [batch_size,13,13,3, class_num] xy, wh, conf, prob = tf.split(yi_pred, [2, 2, 1, -1], axis=-1) ''' compute offset of x and y ''' offset_x = tf.range(shape[2], dtype=tf.float32) #width offset_y = tf.range(shape[1], dtype=tf.float32) # height offset_x, offset_y = tf.meshgrid(offset_x, offset_y) offset_x = tf.reshape(offset_x, (-1, 1)) offset_y = tf.reshape(offset_y, (-1, 1)) offset_xy = tf.concat([offset_x, offset_y], axis=-1) # [13, 13, 1, 2] offset_xy = tf.reshape(offset_xy, [shape[1], shape[2], 1, 2]) xy = tf.math.sigmoid(xy) + offset_xy xy = xy / [shape[2], shape[1]] wh = tf.math.exp(wh) * anchor wh = wh / [width, height] conf = tf.math.sigmoid(conf) prob = tf.math.sigmoid(prob) return xy, wh, conf, prob该部分对boxes、conf、prob分别reshape为[batch_size, 13x13x3, 4]、[batch_size, 13x13x3, 1]、[batch_size, 13x13x3, class_num]，并对三个featuremap进行合并。
def reshape(xy, wh, conf, prob): # [1, 13, 13, 3, 1] x_min = xy[..., 0: 1] - wh[..., 0: 1] / 2.0 x_max = xy[..., 0: 1] + wh[..., 0: 1] / 2.0 y_min = xy[..., 1: 2] - wh[..., 1: 2] / 2.0 y_max = xy[..., 1: 2] + wh[..., 1: 2] / 2.0 # [1, 13, 13, 3, 4] boxes = tf.concat([x_min, y_min, x_max, y_max], -1) shape = tf.shape(boxes) # [1, 13*13*3, 4] boxes = tf.reshape(boxes, (shape[0], shape[1] * shape[2]* shape[3], shape[4])) # [1, 13 * 13 * 3, 1] conf = tf.reshape(conf, (shape[0], shape[1] * shape[2]* shape[3], 1)) # [1, 13*13*3, class_num] prob = tf.reshape(prob, (shape[0], shape[1] * shape[2]* shape[3], -1)) return boxes, conf, prob # reshape # [batch_size, 13*13*3, 4], [batch_size, 13*13*3, 1], [batch_size, 13*13*3, class_num] boxes_1, conf_1, prob_1 = reshape(xy1, wh1, conf1, prob1) boxes_2, conf_2, prob_2 = reshape(xy2, wh2, conf2, prob2) boxes_3, conf_3, prob_3 = reshape(xy3, wh3, conf3, prob3) # gather # [1, 13*13*3, 4] & [1, 26*26*3, 4] ==> [1, V, 4] boxes = tf.concat([boxes_1, boxes_2, boxes_3], 1) conf = tf.concat([conf_1, conf_2, conf_3], 1) prob = tf.concat([prob_1, prob_2, prob_3], 1) return boxes, conf, prob针对每一类别，按score进行非极大值抑制，根据设定的score、iou、maxboxes阈值进行筛选，最终得到每一类别的预测框。
def nms(boxes, scores, class_num, max_boxes=50, score_thresh=0.5, iou_threshold=0.5): boxes_list, label_list, score_list = [], [], [] max_boxes = tf.constant(max_boxes, dtype='int32') # [V, 4] boxes = tf.reshape(boxes, [-1, 4]) # [V, class_num] score = tf.reshape(scores, [-1, class_num]) mask = tf.greater_equal(score, tf.constant(score_thresh)) for i in range(class_num): # Step 3: Apply the mask to scores, boxes and pick them out filter_boxes = tf.boolean_mask(boxes, mask[:,i]) filter_score = tf.boolean_mask(score[:,i], mask[:,i]) nms_indices = tf.image.non_max_suppression(boxes=filter_boxes, scores=filter_score, max_output_size=max_boxes, iou_threshold=iou_threshold, name='nms_indices') label_list.append(tf.ones_like(tf.gather(filter_score, nms_indices), 'int32')*i) boxes_list.append(tf.gather(filter_boxes, nms_indices)) score_list.append(tf.gather(filter_score, nms_indices)) # stack boxes = tf.concat(boxes_list, axis=0) score = tf.concat(score_list, axis=0) label = tf.concat(label_list, axis=0) return boxes, score, label