yolo.py 文件源码

def build_export_output(net_out, H, W, max_number_length, C, threshold):
  B, threshold, sprs_output_box_count = max_number_length, threshold, 100
  net_out = tf.reshape(net_out, [H, W, B, -1])
  boxes, boxes_scores, classes_probs = net_out[:, :, :, :4], net_out[:, :, :, 4], net_out[:, :, :, 5:]
  row = np.concatenate([np.ones([1, W, B], dtype=np.float32) * i for i in range(H)], axis=0)
  col = np.concatenate([np.ones([H, 1, B], dtype=np.float32) * i for i in range(W)], axis=1)
  anchors_w = np.concatenate([np.ones([H, W, 1], dtype=np.float32) * anchors[2 * i + 0] for i in range(B)], axis=2)
  anchors_h = np.concatenate([np.ones([H, W, 1], dtype=np.float32) * anchors[2 * i + 1] for i in range(B)], axis=2)
  with ops.name_scope(None, 'calc_boxes_coordinates'):
    boxes = tf.concat([
      tf.expand_dims((tf.sigmoid(boxes[:, :, :, 0]) + col) / W, 3),
      tf.expand_dims((tf.sigmoid(boxes[:, :, :, 1]) + row) / H, 3),
      tf.expand_dims(tf.exp(boxes[:, :, :, 2]) * anchors_w / W, 3),
      tf.expand_dims(tf.exp(boxes[:, :, :, 3]) * anchors_h / H, 3),
    ], axis=3)
    boxes = tf.cast(boxes, tf.float32)

  with ops.name_scope(None, 'calc_boxes_scores'):
    boxes_scores = tf.sigmoid(boxes_scores)
    boxes_scores = tf.nn.softmax(classes_probs) * tf.expand_dims(boxes_scores, 3)
    boxes_scores = boxes_scores * tf.cast(boxes_scores > threshold, tf.float32)
    boxes_scores = tf.cast(boxes_scores, tf.float32)

  with ops.name_scope(None, 'non_max_suppression'):
    boxes = tf.reshape(boxes, [H * W * B, 4])
    sprs_boxes, sprs_boxes_scores = [], []
    for i in range(C):
      box_scores = tf.reshape(boxes_scores[:, :, :, i], [H * W * B])
      sprs_boxes_indices = tf.image.non_max_suppression(boxes, box_scores, sprs_output_box_count, iou_threshold=0.4)
      box_scores = box_scores * tf.scatter_nd(
        tf.reshape(sprs_boxes_indices, [-1, 1]),
        tf.ones(tf.shape(sprs_boxes_indices), dtype=tf.float32), [H * W * B])
      sprs_boxes_scores.append(tf.reshape(box_scores, [H * W * B, 1]))

  with ops.name_scope(None, 'select_boxes'):
    sprs_boxes_scores = tf.concat(sprs_boxes_scores, axis=1)
    classes = tf.argmax(sprs_boxes_scores, axis=1)
    classes_probs = tf.reduce_max(sprs_boxes_scores, axis=1)

    selected_box_mask = classes_probs > threshold

    selected_classes = tf.boolean_mask(classes, selected_box_mask)
    selected_boxes = tf.boolean_mask(boxes, selected_box_mask)
    selected_classes_probs = tf.boolean_mask(classes_probs, selected_box_mask)
    lefts = selected_boxes[:, 0] - selected_boxes[:, 2] / 2
    lefts = tf.where(lefts < 0, tf.zeros(tf.shape(lefts)), lefts)
    selected_boxes = tf.concat([
      tf.expand_dims(lefts, 1),
      tf.expand_dims(selected_boxes[:, 1] - selected_boxes[:, 3] / 2, 1),
      tf.expand_dims(selected_boxes[:, 2], 1),
      tf.expand_dims(selected_boxes[:, 3], 1),
    ], axis=1)
    selected_lefts = selected_boxes[:, 0]

  with ops.name_scope(None, 'sort_boxes'):
    sorted_lefts, sorted_lefts_indices = tf.nn.top_k(selected_lefts * -1, tf.shape(selected_lefts)[0])

    sorted_classes = tf.gather(selected_classes, sorted_lefts_indices)
    sorted_boxes = tf.gather(selected_boxes, sorted_lefts_indices)
    sorted_classes_probs = tf.gather(selected_classes_probs, sorted_lefts_indices)
  return sorted_lefts * -1, sorted_boxes, sorted_classes, sorted_classes_probs