test passes

toshihiroryuu · Aug 24, 2017 · e55e598 · e55e598
1 parent b648e2b
commit e55e598
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 12 deletions.
diff --git a/keras_rcnn/layers/box_regression.py b/keras_rcnn/layers/box_regression.py
@@ -26,6 +26,7 @@ def call(self, x, **kwargs):
         metadata: image information (1, 3)
         """
         rois, pred_boxes, pred_scores, metadata = x[0], x[1], x[2], x[3]
+
         rois = rois[0, :, :]
         pred_boxes = pred_boxes[0, :, :]
         pred_scores = pred_scores[0, :, :]
@@ -35,25 +36,30 @@ def call(self, x, **kwargs):
         boxes = rois / metadata[0][2]
 
         # Apply bounding-box regression deltas
+        pred_boxes = keras_rcnn.backend.bbox_transform_inv(boxes, pred_boxes)
 
-        pred_boxes = keras_rcnn.layers.object_detection._object_proposal.bbox_transform_inv(boxes, pred_boxes)
         pred_boxes = keras_rcnn.backend.clip(pred_boxes, metadata[0][:2])
 
         # Final detections
 
-        # keep detections above threshold
-        indices_threshold = keras_rcnn.backend.where(keras.backend.greater(pred_scores, self.threshold))
-        pred_scores = keras_rcnn.backend.gather_nd(pred_scores, indices_threshold)
-        pred_boxes = keras_rcnn.backend.gather_nd(pred_boxes, indices_threshold)
-
         # 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)))
+
+        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)[1], self.TEST_NMS)
         pred_scores = keras.backend.gather(pred_scores, indices)
@@ -62,5 +68,6 @@ def call(self, x, **kwargs):
 
         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, 3)
diff --git a/tests/layers/test_box_regression.py b/tests/layers/test_box_regression.py
@@ -10,13 +10,13 @@ def test_call(self):
         classes = 3
         proposal_target = keras_rcnn.layers.BoxRegression()
 
-        proposals = numpy.random.random((1, 100, 4))
-        proposals = keras.backend.variable(proposals)
-
-        pred_boxes = numpy.random.choice(range(0, 224), (1, 300, 4 * classes))
+        pred_boxes = numpy.random.random((1, 100, 4 * classes))
         pred_boxes = keras.backend.variable(pred_boxes)
 
-        pred_scores = numpy.random.random((1, 300, classes))
+        proposals = numpy.random.choice(range(0, 224), (1, 100, 4))
+        proposals = keras.backend.variable(proposals)
+
+        pred_scores = numpy.random.random((1, 100, classes))
         pred_scores = keras.backend.variable(pred_scores)
 
         metadata = keras.backend.variable([[224, 224, 1.5]])