fix input to nms

toshihiroryuu · Aug 24, 2017 · cd06e8b · cd06e8b
1 parent 682cff4
commit cd06e8b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/keras_rcnn/layers/detection.py b/keras_rcnn/layers/detection.py
@@ -34,7 +34,8 @@ def call(self, x, **kwargs):
         # unscale back to raw image space
 
         boxes = rois / metadata[0][2]
-
+        import pdb
+        pdb.set_trace()
         # Apply bounding-box regression deltas
         pred_boxes = keras_rcnn.backend.bbox_transform_inv(boxes, pred_deltas)
 
@@ -61,7 +62,7 @@ def call(self, x, **kwargs):
         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)
+        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)
 

diff --git a/tests/layers/test_detection.py b/tests/layers/test_detection.py
@@ -21,5 +21,9 @@ def test_call(self):
 
         metadata = keras.backend.variable([[224, 224, 1.5]])
 
-        proposal_target.call([proposals, pred_boxes, pred_scores, metadata])
+        boxes, classes = proposal_target.call([proposals, pred_boxes, pred_scores, metadata])
+
+        assert keras.backend.eval(boxes).shape == keras.backend.eval(pred_boxes).shape
+
+        assert keras.backend.eval(classes).shape == keras.backend.eval(boxes).shape