example-apps/PDD/pcb-defect-detection/libs/box_utils/encode_and_decode.py

   1 # -*- coding: utf-8 -*-
   2
   3 from __future__ import absolute_import
   4 from __future__ import print_function
   5 from __future__ import division
   6
   7 import tensorflow as tf
   8 import numpy as np
   9
  10
  11
  12 # def encode_boxes(ex_rois, gt_rois, scale_factor=None):
  13 #     ex_widths = ex_rois[:, 2] - ex_rois[:, 0] + 1.0
  14 #     ex_heights = ex_rois[:, 3] - ex_rois[:, 1] + 1.0
  15 #     ex_ctr_x = ex_rois[:, 0] + 0.5 * ex_widths
  16 #     ex_ctr_y = ex_rois[:, 1] + 0.5 * ex_heights
  17 #
  18 #     gt_widths = gt_rois[:, 2] - gt_rois[:, 0] + 1.0
  19 #     gt_heights = gt_rois[:, 3] - gt_rois[:, 1] + 1.0
  20 #     gt_ctr_x = gt_rois[:, 0] + 0.5 * gt_widths
  21 #     gt_ctr_y = gt_rois[:, 1] + 0.5 * gt_heights
  22 #
  23 #     targets_dx = (gt_ctr_x - ex_ctr_x) / ex_widths
  24 #     targets_dy = (gt_ctr_y - ex_ctr_y) / ex_heights
  25 #     targets_dw = np.log(gt_widths / ex_widths)
  26 #     targets_dh = np.log(gt_heights / ex_heights)
  27 #
  28 #     if scale_factor:
  29 #         targets_dx = targets_dx * scale_factor[0]
  30 #         targets_dy = targets_dy * scale_factor[1]
  31 #         targets_dw = targets_dw * scale_factor[2]
  32 #         targets_dh = targets_dh * scale_factor[3]
  33 #
  34 #     targets = np.vstack(
  35 #         (targets_dx, targets_dy, targets_dw, targets_dh)).transpose()
  36 #     return targets
  37 #
  38 #
  39 # def _concat_new_axis(t1, t2, t3, t4, axis):
  40 #     return tf.concat(
  41 #         [tf.expand_dims(t1, -1), tf.expand_dims(t2, -1),
  42 #          tf.expand_dims(t3, -1), tf.expand_dims(t4, -1)], axis=axis)
  43 #
  44 #
  45 # def decode_boxes(boxes, deltas, scale_factor=None):
  46 #     widths = boxes[:, 2] - boxes[:, 0] + 1.0
  47 #     heights = boxes[:, 3] - boxes[:, 1] + 1.0
  48 #     ctr_x = tf.expand_dims(boxes[:, 0] + 0.5 * widths, -1)
  49 #     ctr_y = tf.expand_dims(boxes[:, 1] + 0.5 * heights, -1)
  50 #
  51 #     dx = deltas[:, 0::4]
  52 #     dy = deltas[:, 1::4]
  53 #     dw = deltas[:, 2::4]
  54 #     dh = deltas[:, 3::4]
  55 #
  56 #     if scale_factor:
  57 #         dx /= scale_factor[0]
  58 #         dy /= scale_factor[1]
  59 #         dw /= scale_factor[2]
  60 #         dh /= scale_factor[3]
  61 #
  62 #     widths = tf.expand_dims(widths, -1)
  63 #     heights = tf.expand_dims(heights, -1)
  64 #
  65 #     pred_ctr_x = dx * widths + ctr_x
  66 #     pred_ctr_y = dy * heights + ctr_y
  67 #     pred_w = tf.exp(dw) * widths
  68 #     pred_h = tf.exp(dh) * heights
  69 #
  70 #     # x1
  71 #     # pred_boxes[:, 0::4] = pred_ctr_x - 0.5 * pred_w
  72 #     pred_x1 = pred_ctr_x - 0.5 * pred_w
  73 #     # y1
  74 #     # pred_boxes[:, 1::4] = pred_ctr_y - 0.5 * pred_h
  75 #     pred_y1 = pred_ctr_y - 0.5 * pred_h
  76 #     # x2
  77 #     # pred_boxes[:, 2::4] = pred_ctr_x + 0.5 * pred_w
  78 #     pred_x2 = pred_ctr_x + 0.5 * pred_w
  79 #     # y2
  80 #     # pred_boxes[:, 3::4] = pred_ctr_y + 0.5 * pred_h
  81 #     pred_y2 = pred_ctr_y + 0.5 * pred_h
  82 #
  83 #     pred_boxes = _concat_new_axis(pred_x1, pred_y1, pred_x2, pred_y2, 2)
  84 #     pred_boxes = tf.reshape(pred_boxes, (tf.shape(pred_boxes)[0], -1))
  85 #     return pred_boxes
  86
  87
  88 def decode_boxes(encoded_boxes, reference_boxes, scale_factors=None):
  89     '''
  90
  91     :param encoded_boxes:[N, 4]
  92     :param reference_boxes: [N, 4] .
  93     :param scale_factors: use for scale.
  94
  95     in the first stage, reference_boxes  are anchors
  96     in the second stage, reference boxes are proposals(decode) produced by first stage
  97     :return:decode boxes [N, 4]
  98     '''
  99
 100     t_xcenter, t_ycenter, t_w, t_h = tf.unstack(encoded_boxes, axis=1)
 101     if scale_factors:
 102         t_xcenter /= scale_factors[0]
 103         t_ycenter /= scale_factors[1]
 104         t_w /= scale_factors[2]
 105         t_h /= scale_factors[3]
 106
 107     reference_xmin, reference_ymin, reference_xmax, reference_ymax = tf.unstack(reference_boxes, axis=1)
 108     # reference boxes are anchors in the first stage
 109
 110     reference_xcenter = (reference_xmin + reference_xmax) / 2.
 111     reference_ycenter = (reference_ymin + reference_ymax) / 2.
 112     reference_w = reference_xmax - reference_xmin
 113     reference_h = reference_ymax - reference_ymin
 114
 115     predict_xcenter = t_xcenter * reference_w + reference_xcenter
 116     predict_ycenter = t_ycenter * reference_h + reference_ycenter
 117     predict_w = tf.exp(t_w) * reference_w
 118     predict_h = tf.exp(t_h) * reference_h
 119
 120     predict_xmin = predict_xcenter - predict_w / 2.
 121     predict_xmax = predict_xcenter + predict_w / 2.
 122     predict_ymin = predict_ycenter - predict_h / 2.
 123     predict_ymax = predict_ycenter + predict_h / 2.
 124
 125     return tf.transpose(tf.stack([predict_xmin, predict_ymin,
 126                                   predict_xmax, predict_ymax]))
 127
 128
 129 def encode_boxes(unencode_boxes, reference_boxes, scale_factors=None):
 130     '''
 131
 132     :param unencode_boxes: [-1, 4]
 133     :param reference_boxes: [-1, 4]
 134     :return: encode_boxes [-1, 4]
 135     '''
 136
 137     xmin, ymin, xmax, ymax = unencode_boxes[:, 0], unencode_boxes[:, 1], unencode_boxes[:, 2], unencode_boxes[:, 3]
 138
 139     reference_xmin, reference_ymin, reference_xmax, reference_ymax = reference_boxes[:, 0], reference_boxes[:, 1], \
 140                                                                      reference_boxes[:, 2], reference_boxes[:, 3]
 141
 142     x_center = (xmin + xmax) / 2.
 143     y_center = (ymin + ymax) / 2.
 144     w = xmax - xmin + 1e-8
 145     h = ymax - ymin + 1e-8
 146
 147     reference_xcenter = (reference_xmin + reference_xmax) / 2.
 148     reference_ycenter = (reference_ymin + reference_ymax) / 2.
 149     reference_w = reference_xmax - reference_xmin + 1e-8
 150     reference_h = reference_ymax - reference_ymin + 1e-8
 151
 152     # w + 1e-8 to avoid NaN in division and log below
 153
 154     t_xcenter = (x_center - reference_xcenter) / reference_w
 155     t_ycenter = (y_center - reference_ycenter) / reference_h
 156     t_w = np.log(w/reference_w)
 157     t_h = np.log(h/reference_h)
 158
 159     if scale_factors:
 160         t_xcenter *= scale_factors[0]
 161         t_ycenter *= scale_factors[1]
 162         t_w *= scale_factors[2]
 163         t_h *= scale_factors[3]
 164
 165     return np.transpose(np.stack([t_xcenter, t_ycenter, t_w, t_h], axis=0))