squeeze

toshihiroryuu · Aug 17, 2017 · b8018e2 · b8018e2
1 parent e3d6846
commit b8018e2
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 2 deletions.
diff --git a/keras_rcnn/backend/tensorflow_backend.py b/keras_rcnn/backend/tensorflow_backend.py
@@ -94,3 +94,10 @@ def crop_and_resize(image, boxes, size):
     boxes = keras.backend.reshape(boxes, [-1, 4])
 
     return tensorflow.image.crop_and_resize(image, boxes, box_ind, size)
+
+
+def squeeze(a, axis=None):
+    """
+    Remove single-dimensional entries from the shape of an array.
+    """
+    return tensorflow.squeeze(a, axis)
diff --git a/keras_rcnn/layers/object_detection/_anchor_target.py b/keras_rcnn/layers/object_detection/_anchor_target.py
@@ -256,7 +256,7 @@ def unmap(data, count, inds_inside, fill=0):
 
         inds_nd = keras.backend.concatenate([inds_ii, inds_coords], 1)
 
-    inverse_ret = tensorflow.squeeze(tensorflow.gather_nd(-1 * ret, inds_nd))
+    inverse_ret = keras_rcnn.backend.squeeze(tensorflow.gather_nd(-1 * ret, inds_nd))
 
     ret = keras_rcnn.backend.scatter_add_tensor(ret, inds_nd, inverse_ret + data)
 

diff --git a/tests/backend/test_tensorflow_backend.py b/tests/backend/test_tensorflow_backend.py
@@ -2,7 +2,7 @@
 import numpy
 import tensorflow
 
-import keras_rcnn.backend
+import keras_rcnn.backend.tensorflow_backend
 import keras_rcnn.backend.common
 
 
@@ -35,3 +35,19 @@ def test_crop_and_resize():
     slices = keras_rcnn.backend.crop_and_resize(image, boxes, size)
 
     assert keras.backend.eval(slices).shape == (2, 7, 7, 3)
+
+
+def test_squeeze():
+    x = [[[0], [1], [2]]]
+
+    x = keras.backend.variable(x)
+
+    assert keras.backend.int_shape(x) == (1, 3, 1)
+
+    y = keras_rcnn.backend.tensorflow_backend.squeeze(x)
+
+    assert keras.backend.int_shape(y) == (3,)
+
+    y = keras_rcnn.backend.tensorflow_backend.squeeze(x, axis=0)
+
+    assert keras.backend.int_shape(y) == (3, 1)