Commit ef334e31 authored by Yuxin Wu's avatar Yuxin Wu

[MaskRCNN] bugfix on inference

parent f0e049a5
...@@ -201,6 +201,7 @@ def fastrcnn_predictions(boxes, scores): ...@@ -201,6 +201,7 @@ def fastrcnn_predictions(boxes, scores):
Returns: n boolean, the selection Returns: n boolean, the selection
""" """
prob, box = X prob, box = X
output_shape = tf.shape(prob, out_type=tf.int64)
# filter by score threshold # filter by score threshold
ids = tf.reshape(tf.where(prob > cfg.TEST.RESULT_SCORE_THRESH), [-1]) ids = tf.reshape(tf.where(prob > cfg.TEST.RESULT_SCORE_THRESH), [-1])
prob = tf.gather(prob, ids) prob = tf.gather(prob, ids)
...@@ -216,13 +217,13 @@ def fastrcnn_predictions(boxes, scores): ...@@ -216,13 +217,13 @@ def fastrcnn_predictions(boxes, scores):
if get_tf_version_tuple() >= (1, 12): if get_tf_version_tuple() >= (1, 12):
mask = tf.sparse.SparseTensor(indices=tf.expand_dims(sorted_selection, 1), mask = tf.sparse.SparseTensor(indices=tf.expand_dims(sorted_selection, 1),
values=tf.ones_like(sorted_selection, dtype=tf.bool), values=tf.ones_like(sorted_selection, dtype=tf.bool),
dense_shape=tf.shape(prob, out_type=tf.int64)) dense_shape=output_shape)
mask = tf.sparse.to_dense(mask, default_value=False) mask = tf.sparse.to_dense(mask, default_value=False)
else: else:
# this function is deprecated by TF # this function is deprecated by TF
mask = tf.sparse_to_dense( mask = tf.sparse_to_dense(
sparse_indices=sorted_selection, sparse_indices=sorted_selection,
output_shape=tf.shape(prob, out_type=tf.int64), output_shape=output_shape,
sparse_values=True, sparse_values=True,
default_value=False) default_value=False)
return mask return mask
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment