Add tensorflow implementation of proposal_layer_top

study3h · Mar 15, 2018 · 83744d0 · 83744d0
1 parent acb3643
commit 83744d0
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 15 deletions.
diff --git a/lib/layer_utils/proposal_top_layer.py b/lib/layer_utils/proposal_top_layer.py
@@ -9,9 +9,11 @@
 
 import numpy as np
 from model.config import cfg
-from model.bbox_transform import bbox_transform_inv, clip_boxes
+from model.bbox_transform import bbox_transform_inv, clip_boxes, bbox_transform_inv_tf, clip_boxes_tf
 import numpy.random as npr
 
+import tensorflow as tf
+
 def proposal_top_layer(rpn_cls_prob, rpn_bbox_pred, im_info, _feat_stride, anchors, num_anchors):
   """A layer that just selects the top region proposals
      without using non-maximal suppression,
@@ -51,3 +53,21 @@ def proposal_top_layer(rpn_cls_prob, rpn_bbox_pred, im_info, _feat_stride, ancho
   batch_inds = np.zeros((proposals.shape[0], 1), dtype=np.float32)
   blob = np.hstack((batch_inds, proposals.astype(np.float32, copy=False)))
   return blob, scores
+
+def proposal_top_layer_tf(rpn_cls_prob, rpn_bbox_pred, im_info, _feat_stride, anchors, num_anchors):
+  rpn_top_n = cfg.TEST.RPN_TOP_N
+
+  scores = rpn_cls_prob[:, :, :, num_anchors:]
+  rpn_bbox_pred = tf.reshape(rpn_bbox_pred, shape=(-1, 4))
+  scores = tf.reshape(scores, shape=(-1,))
+
+  top_scores, top_inds = tf.nn.top_k(scores, k=rpn_top_n)
+  top_scores = tf.reshape(top_scores, shape=(-1, 1))
+  top_anchors = tf.gather(anchors, top_inds)
+  top_rpn_bbox = tf.gather(rpn_bbox_pred, top_inds)
+  proposals = bbox_transform_inv_tf(top_anchors, top_rpn_bbox)
+  proposals = clip_boxes_tf(proposals, im_info[:2])
+  proposals = tf.to_float(proposals)
+  batch_inds = tf.zeros((rpn_top_n, 1))
+  blob = tf.concat([batch_inds, proposals], 1)
+  return blob, top_scores
diff --git a/lib/nets/network.py b/lib/nets/network.py
@@ -16,7 +16,7 @@
 
 from layer_utils.snippets import generate_anchors_pre, generate_anchors_tf
 from layer_utils.proposal_layer import proposal_layer, proposal_layer_tf
-from layer_utils.proposal_top_layer import proposal_top_layer
+from layer_utils.proposal_top_layer import proposal_top_layer, proposal_top_layer_tf
 from layer_utils.anchor_target_layer import anchor_target_layer
 from layer_utils.proposal_target_layer import proposal_target_layer
 from utils.visualization import draw_bounding_boxes
@@ -87,26 +87,38 @@ def _softmax_layer(self, bottom, name):
 
   def _proposal_top_layer(self, rpn_cls_prob, rpn_bbox_pred, name):
     with tf.variable_scope(name) as scope:
-      rois, rpn_scores = tf.py_func(proposal_top_layer,
-                                    [rpn_cls_prob, rpn_bbox_pred, self._im_info,
-                                     self._feat_stride, self._anchors, self._num_anchors],
-                                    [tf.float32, tf.float32], name="proposal_top")
-      rois.set_shape([cfg.TEST.RPN_TOP_N, 5])
-      rpn_scores.set_shape([cfg.TEST.RPN_TOP_N, 1])
-
-    return rois, rpn_scores
-
-  def _proposal_layer(self, rpn_cls_prob, rpn_bbox_pred, name):
-    with tf.variable_scope(name) as scope:
-      rois, rpn_scores = proposal_layer_tf(
+      rois, rpn_scores = proposal_top_layer_tf(
         rpn_cls_prob,
         rpn_bbox_pred,
         self._im_info,
-        self._mode,
         self._feat_stride,
         self._anchors,
         self._num_anchors
       )
+      # rois, rpn_scores = tf.py_func(proposal_top_layer,
+      #                               [rpn_cls_prob, rpn_bbox_pred, self._im_info,
+      #                                self._feat_stride, self._anchors, self._num_anchors],
+      #                               [tf.float32, tf.float32], name="proposal_top")
+      rois.set_shape([cfg.TEST.RPN_TOP_N, 5])
+      rpn_scores.set_shape([cfg.TEST.RPN_TOP_N, 1])
+
+    return rois, rpn_scores
+
+  def _proposal_layer(self, rpn_cls_prob, rpn_bbox_pred, name):
+    with tf.variable_scope(name) as scope:
+      # rois, rpn_scores = proposal_layer_tf(
+      #   rpn_cls_prob,
+      #   rpn_bbox_pred,
+      #   self._im_info,
+      #   self._mode,
+      #   self._feat_stride,
+      #   self._anchors,
+      #   self._num_anchors
+      # )
+      rois, rpn_scores = tf.py_func(proposal_layer,
+                              [rpn_cls_prob, rpn_bbox_pred, self._im_info, self._mode,
+                               self._feat_stride, self._anchors, self._num_anchors],
+                              [tf.float32, tf.float32], name="proposal")
       rois.set_shape([None, 5])
       rpn_scores.set_shape([None, 1])