Merge pull request broadinstitute#96 from jhung0/final_detection_layer

Final detection layer
toshihiroryuu · Aug 25, 2017 · c7b51b1 · c7b51b1
2 parents 6d6b0ff + c3e7077
commit c7b51b1
Show file tree

Hide file tree

Showing 7 changed files with 209 additions and 86 deletions.
diff --git a/keras_rcnn/backend/common.py b/keras_rcnn/backend/common.py
@@ -193,3 +193,54 @@ def smooth_l1_loss(y_true, y_pred):
     x = keras.backend.switch(p, q, r)
 
     return keras.backend.sum(x)
+
+
+def bbox_transform_inv(boxes, deltas):
+    def shape_zero():
+        x = keras.backend.int_shape(deltas)[-1]
+
+        return keras.backend.zeros_like(x, dtype=keras.backend.floatx())
+
+    def shape_non_zero():
+        a = boxes[:, 2] - boxes[:, 0] + 1.0
+        b = boxes[:, 3] - boxes[:, 1] + 1.0
+
+        ctr_x = boxes[:, 0] + 0.5 * a
+        ctr_y = boxes[:, 1] + 0.5 * b
+
+        dx = deltas[:, 0::4]
+        dy = deltas[:, 1::4]
+        dw = deltas[:, 2::4]
+        dh = deltas[:, 3::4]
+
+        pred_ctr_x = dx * a[:, keras_rcnn.backend.newaxis] + ctr_x[:, keras_rcnn.backend.newaxis]
+        pred_ctr_y = dy * b[:, keras_rcnn.backend.newaxis] + ctr_y[:, keras_rcnn.backend.newaxis]
+
+        pred_w = keras.backend.exp(dw) * a[:, keras_rcnn.backend.newaxis]
+        pred_h = keras.backend.exp(dh) * b[:, keras_rcnn.backend.newaxis]
+
+
+        indices = keras.backend.tile(keras.backend.arange(0, keras.backend.shape(deltas)[0]), [4])
+        indices = keras.backend.reshape(indices, (-1, 1))
+        indices = keras.backend.tile(indices, [1, keras.backend.shape(deltas)[-1] // 4])
+        indices = keras.backend.reshape(indices, (-1, 1))
+        indices_coords = keras.backend.tile(keras.backend.arange(0, keras.backend.shape(deltas)[1], step = 4), [keras.backend.shape(deltas)[0]])
+        indices_coords = keras.backend.concatenate([indices_coords, indices_coords + 1, indices_coords + 2, indices_coords + 3], 0)
+        indices = keras.backend.concatenate([indices, keras.backend.expand_dims(indices_coords)], axis=1)
+
+
+        updates = keras.backend.concatenate([keras.backend.reshape(pred_ctr_x - 0.5 * pred_w, (-1,)),
+                                             keras.backend.reshape(pred_ctr_y - 0.5 * pred_h, (-1,)),
+                                             keras.backend.reshape(pred_ctr_x + 0.5 * pred_w, (-1,)),
+                                             keras.backend.reshape(pred_ctr_y + 0.5 * pred_h, (-1,))], axis=0)
+        pred_boxes = keras_rcnn.backend.scatter_add_tensor(keras.backend.zeros_like(deltas), indices, updates)
+        return pred_boxes
+
+
+    zero_boxes = keras.backend.equal(keras.backend.shape(deltas)[0], 0)
+
+    pred_boxes = keras.backend.switch(zero_boxes, shape_zero, shape_non_zero)
+
+    return pred_boxes
+
+
diff --git a/keras_rcnn/layers/__init__.py b/keras_rcnn/layers/__init__.py
@@ -3,3 +3,5 @@
 from .object_detection import (AnchorTarget, ObjectProposal, ProposalTarget)
 
 from .pooling import RegionOfInterest
+
+from .detection import Detection
diff --git a/keras_rcnn/layers/detection.py b/keras_rcnn/layers/detection.py
@@ -0,0 +1,87 @@
+import keras.engine.topology
+import keras.backend
+import keras_rcnn.backend
+import keras_rcnn.layers.object_detection._object_proposal
+
+class Detection(keras.engine.topology.Layer):
+    """
+    Get final detections + labels by unscaling back to image space, applying regression deltas,
+    choosing box coordinates, and removing extra detections via NMS
+
+    # Arguments
+    threshold: objects with maximum score less than threshold are thrown out
+    test_nms: A float representing the threshold for deciding whether boxes overlap too much with respect to IoU
+
+    """
+    def __init__(self, threshold = 0.05, test_nms = 0.5, **kwargs):
+        self.threshold = threshold
+
+        self.TEST_NMS = test_nms
+
+        super(Detection, self).__init__(**kwargs)
+
+    def build(self, input_shape):
+
+        super(Detection, self).build(input_shape)
+
+    def call(self, x, **kwargs):
+        """
+        # Inputs
+        rois: output of proposal target (1, N, 4)
+        pred_deltas: predicted deltas (1, N, 4*classes)
+        pred_scores: score distributions (1, N, classes)
+        metadata: image information (1, 3)
+
+        # Returns
+        pred_boxes: final predicted boxes of the predicted class (1, N, 4)
+        pred_scores: score distribution over all classes (1, N, classes), note the box only corresponds to the most
+            probable class, not the other classes
+        """
+        rois, pred_deltas, pred_scores, metadata = x[0], x[1], x[2], x[3]
+
+        rois = rois[0, :, :]
+        pred_deltas = pred_deltas[0, :, :]
+        pred_scores = pred_scores[0, :, :]
+
+        # unscale back to raw image space
+
+        boxes = rois / metadata[0][2]
+
+        # Apply bounding-box regression deltas
+        pred_boxes = keras_rcnn.backend.bbox_transform_inv(boxes, pred_deltas)
+
+        pred_boxes = keras_rcnn.backend.clip(pred_boxes, metadata[0][:2])
+
+        # Final detections
+
+        # for each object, get the top class score and corresponding bbox, apply nms
+        pred_classes = keras.backend.argmax(pred_scores, axis=1)
+        pred_classes = keras.backend.cast(pred_classes, 'int32')
+
+        # keep detections above threshold
+
+        indices_threshold = keras_rcnn.backend.where(keras.backend.greater(keras.backend.max(pred_scores, axis=1), self.threshold))
+        indices_threshold = keras.backend.reshape(indices_threshold, (-1,))
+        pred_scores = keras.backend.gather(pred_scores, indices_threshold)
+        pred_boxes = keras.backend.gather(pred_boxes, indices_threshold)
+
+        indices = keras.backend.arange(0, keras.backend.shape(pred_scores)[0])
+        pred_scores_classes = keras_rcnn.backend.gather_nd(pred_scores, keras.backend.concatenate([keras.backend.expand_dims(indices), keras.backend.expand_dims(pred_classes)], axis=1))
+        indices_boxes = keras.backend.concatenate([4 * pred_classes, 4 * pred_classes + 1, 4 * pred_classes + 2, 4 * pred_classes + 3], 0)
+        indices = keras.backend.tile(indices, [4])
+
+        pred_boxes = keras_rcnn.backend.gather_nd(pred_boxes, keras.backend.concatenate([keras.backend.expand_dims(indices), keras.backend.expand_dims(indices_boxes)], axis=1))
+        pred_boxes = keras.backend.reshape(pred_boxes, (-1, 4))
+
+        indices = keras_rcnn.backend.non_maximum_suppression(pred_boxes, pred_scores_classes, keras.backend.shape(pred_boxes)[0], self.TEST_NMS)
+        pred_scores = keras.backend.gather(pred_scores, indices)
+        pred_boxes = keras.backend.gather(pred_boxes, indices)
+
+        return [keras.backend.expand_dims(pred_boxes, 0), keras.backend.expand_dims(pred_scores, 0)]
+
+
+    def compute_output_shape(self, input_shape):
+        return [(1, None, 4), (1, None, input_shape[2][2])]
+
+    def compute_mask(self, inputs, mask=None):
+        return 2 * [None]
diff --git a/keras_rcnn/layers/object_detection/_object_proposal.py b/keras_rcnn/layers/object_detection/_object_proposal.py
@@ -58,7 +58,7 @@ def call(self, inputs, **kwargs):
         deltas = keras.backend.reshape(deltas, (-1, 4))
         scores = keras.backend.reshape(scores, (-1, 1))
 
-        deltas = bbox_transform_inv(anchors, deltas)
+        deltas = keras_rcnn.backend.bbox_transform_inv(anchors, deltas)
 
         # 2. clip predicted boxes to image
         proposals = keras_rcnn.backend.clip(deltas, image_shape)
@@ -98,55 +98,6 @@ def compute_output_shape(self, input_shape):
         return None, self.maximum_proposals, 4
 
 
-def bbox_transform_inv(shifted, boxes):
-    def shape_zero():
-        x = keras.backend.int_shape(boxes)[-1]
-
-        return keras.backend.zeros_like(x, dtype=keras.backend.floatx())
-
-    def shape_non_zero():
-        a = shifted[:, 2] - shifted[:, 0] + 1.0
-        b = shifted[:, 3] - shifted[:, 1] + 1.0
-
-        ctr_x = shifted[:, 0] + 0.5 * a
-        ctr_y = shifted[:, 1] + 0.5 * b
-
-        dx = boxes[:, 0::4]
-        dy = boxes[:, 1::4]
-        dw = boxes[:, 2::4]
-        dh = boxes[:, 3::4]
-
-        pred_ctr_x = dx * a[:, keras_rcnn.backend.newaxis] + ctr_x[:, keras_rcnn.backend.newaxis]
-        pred_ctr_y = dy * b[:, keras_rcnn.backend.newaxis] + ctr_y[:, keras_rcnn.backend.newaxis]
-
-        pred_w = keras.backend.exp(dw) * a[:, keras_rcnn.backend.newaxis]
-        pred_h = keras.backend.exp(dh) * b[:, keras_rcnn.backend.newaxis]
-
-
-        indices = keras.backend.tile(keras.backend.arange(0, keras.backend.shape(boxes)[0]), [4])
-        indices = keras.backend.reshape(indices, (-1, 1))
-        indices = keras.backend.tile(indices, [1, keras.backend.shape(boxes)[-1] // 4])
-        indices = keras.backend.reshape(indices, (-1, 1))
-        indices_coords = keras.backend.tile(keras.backend.arange(0, keras.backend.shape(boxes)[1], step = 4), [keras.backend.shape(boxes)[0]])
-        indices_coords = keras.backend.concatenate([indices_coords, indices_coords + 1, indices_coords + 2, indices_coords + 3], 0)
-        indices = keras.backend.concatenate([indices, keras.backend.expand_dims(indices_coords)], axis=1)
-
-
-        updates = keras.backend.concatenate([keras.backend.reshape(pred_ctr_x - 0.5 * pred_w, (-1,)),
-                                             keras.backend.reshape(pred_ctr_y - 0.5 * pred_h, (-1,)),
-                                             keras.backend.reshape(pred_ctr_x + 0.5 * pred_w, (-1,)),
-                                             keras.backend.reshape(pred_ctr_y + 0.5 * pred_h, (-1,))], axis=0)
-        pred_boxes = keras_rcnn.backend.scatter_add_tensor(keras.backend.zeros_like(boxes), indices, updates)
-        return pred_boxes
-
-
-    zero_boxes = keras.backend.equal(keras.backend.shape(boxes)[0], 0)
-
-    pred_boxes = keras.backend.switch(zero_boxes, shape_zero, shape_non_zero)
-
-    return pred_boxes
-
-
 def filter_boxes(proposals, minimum):
     """
     Filters proposed RoIs so that all have width and height at least as big as minimum

diff --git a/tests/backend/test_common.py b/tests/backend/test_common.py
@@ -166,3 +166,40 @@ def test_shift():
     assert keras.backend.int_shape(y) == (1764, 4)
 
     assert y.dtype == keras.backend.floatx()
+
+
+def test_bbox_transform_inv():
+    anchors = 9
+    features = (14, 14)
+    shifted = keras_rcnn.backend.shift(features, 16)
+    deltas = numpy.zeros((features[0] * features[1] * anchors, 4))
+    deltas = keras.backend.variable(deltas)
+    pred_boxes = keras_rcnn.backend.bbox_transform_inv(shifted, deltas)
+    assert keras.backend.eval(pred_boxes).shape == (1764, 4)
+
+    shifted = numpy.zeros((5, 4))
+    deltas = numpy.reshape(numpy.arange(12*5), (5, -1))
+    deltas = keras.backend.variable(deltas)
+    pred_boxes = keras_rcnn.backend.bbox_transform_inv(shifted, deltas)
+    expected = numpy.array(
+        [[ -3.19452805e+00,  -8.54276846e+00,   4.19452805e+00,
+           1.15427685e+01,  -1.97214397e+02,  -5.42816579e+02,
+           2.06214397e+02,   5.53816579e+02,  -1.10047329e+04,
+          -2.99275709e+04,   1.10217329e+04,   2.99465709e+04],
+        [ -6.01289642e+05,  -1.63449519e+06,   6.01314642e+05,
+           1.63452219e+06,  -3.28299681e+07,  -8.92411330e+07,
+           3.28300011e+07,   8.92411680e+07,  -1.79245640e+09,
+          -4.87240170e+09,   1.79245644e+09,   4.87240174e+09],
+        [ -9.78648047e+10,  -2.66024120e+11,   9.78648047e+10,
+           2.66024120e+11,  -5.34323729e+12,  -1.45244248e+13,
+           5.34323729e+12,   1.45244248e+13,  -2.91730871e+14,
+          -7.93006726e+14,   2.91730871e+14,   7.93006726e+14],
+        [ -1.59279659e+16,  -4.32967002e+16,   1.59279659e+16,
+           4.32967002e+16,  -8.69637471e+17,  -2.36391973e+18,
+           8.69637471e+17,   2.36391973e+18,  -4.74805971e+19,
+          -1.29065644e+20,   4.74805971e+19,   1.29065644e+20],
+        [ -2.59235276e+21,  -7.04674541e+21,   2.59235276e+21,
+           7.04674541e+21,  -1.41537665e+23,  -3.84739263e+23,
+           1.41537665e+23,   3.84739263e+23,  -7.72769468e+24,
+          -2.10060520e+25,   7.72769468e+24,   2.10060520e+25]], dtype=numpy.float32)
+    numpy.testing.assert_array_almost_equal(keras.backend.eval(pred_boxes)[0], expected[0], 0, verbose=True)
diff --git a/tests/layers/object_detection/test_object_proposal.py b/tests/layers/object_detection/test_object_proposal.py
@@ -20,42 +20,6 @@ def test_call(self):
         object_proposal.call([metadata, deltas, scores])
 
 
-def test_bbox_transform_inv():
-    anchors = 9
-    features = (14, 14)
-    shifted = keras_rcnn.backend.shift(features, 16)
-    deltas = numpy.zeros((features[0] * features[1] * anchors, 4))
-    deltas = keras.backend.variable(deltas)
-    pred_boxes = keras_rcnn.layers.object_detection._object_proposal.bbox_transform_inv(shifted, deltas)
-    assert keras.backend.eval(pred_boxes).shape == (1764, 4)
-
-    shifted = numpy.zeros((5, 4))
-    deltas = numpy.reshape(numpy.arange(12*5), (5, -1))
-    deltas = keras.backend.variable(deltas)
-    pred_boxes = keras_rcnn.layers.object_detection._object_proposal.bbox_transform_inv(shifted, deltas)
-    expected = numpy.array(
-        [[ -3.19452805e+00,  -8.54276846e+00,   4.19452805e+00,
-           1.15427685e+01,  -1.97214397e+02,  -5.42816579e+02,
-           2.06214397e+02,   5.53816579e+02,  -1.10047329e+04,
-          -2.99275709e+04,   1.10217329e+04,   2.99465709e+04],
-        [ -6.01289642e+05,  -1.63449519e+06,   6.01314642e+05,
-           1.63452219e+06,  -3.28299681e+07,  -8.92411330e+07,
-           3.28300011e+07,   8.92411680e+07,  -1.79245640e+09,
-          -4.87240170e+09,   1.79245644e+09,   4.87240174e+09],
-        [ -9.78648047e+10,  -2.66024120e+11,   9.78648047e+10,
-           2.66024120e+11,  -5.34323729e+12,  -1.45244248e+13,
-           5.34323729e+12,   1.45244248e+13,  -2.91730871e+14,
-          -7.93006726e+14,   2.91730871e+14,   7.93006726e+14],
-        [ -1.59279659e+16,  -4.32967002e+16,   1.59279659e+16,
-           4.32967002e+16,  -8.69637471e+17,  -2.36391973e+18,
-           8.69637471e+17,   2.36391973e+18,  -4.74805971e+19,
-          -1.29065644e+20,   4.74805971e+19,   1.29065644e+20],
-        [ -2.59235276e+21,  -7.04674541e+21,   2.59235276e+21,
-           7.04674541e+21,  -1.41537665e+23,  -3.84739263e+23,
-           1.41537665e+23,   3.84739263e+23,  -7.72769468e+24,
-          -2.10060520e+25,   7.72769468e+24,   2.10060520e+25]], dtype=numpy.float32)
-    numpy.testing.assert_array_almost_equal(keras.backend.eval(pred_boxes)[0], expected[0], 0, verbose=True)
-
 def test_filter_boxes():
     proposals = numpy.array(
         [[0, 2, 3, 10],

diff --git a/tests/layers/test_detection.py b/tests/layers/test_detection.py
@@ -0,0 +1,31 @@
+import keras.backend
+import keras.utils
+import numpy
+
+import keras_rcnn.layers
+
+
+class TestDetection:
+    def test_call(self):
+        num_classes = 3
+        proposal_target = keras_rcnn.layers.Detection()
+
+        pred_boxes = numpy.random.random((1, 100, 4 * num_classes))
+        pred_boxes = keras.backend.variable(pred_boxes)
+
+        proposals = numpy.random.choice(range(0, 224), (1, 100, 4))
+        proposals = keras.backend.variable(proposals)
+
+        pred_scores = numpy.random.random((1, 100, num_classes))
+        pred_scores = keras.backend.variable(pred_scores)
+
+        metadata = keras.backend.variable([[224, 224, 1.5]])
+
+        boxes, classes = proposal_target.call([proposals, pred_boxes, pred_scores, metadata])
+
+        assert keras.backend.eval(classes).shape[:2] == keras.backend.eval(boxes).shape[:2]
+
+        assert keras.backend.eval(boxes).shape[-1] == 4
+
+        assert keras.backend.eval(classes).shape[-1] == num_classes
+
Original file line number	Diff line number	Diff line change
Expand Up		@@ -3,3 +3,5 @@
		from .object_detection import (AnchorTarget, ObjectProposal, ProposalTarget)

		from .pooling import RegionOfInterest

		from .detection import Detection