tf2.0得到子boolmask在boolmask中的索引

在yolo中计算了单层的anchor与全局的gt间的iou score，但是我需要在其中过滤出单层的anchor对应单层的gt的iou score。目前有单一层的gt的loc_mask，以及全局的gt的glob_mask，其中loc_mask中有效区域是glob_mask的子集，因此需要找到loc_mask在glob_mask的对应索引。

问题描述

前面讲的可能比较抽象，所以给出一个例子：

iou = tf.random.normal((4, 4, 9))
loc_mask = tf.constant([[False, False, False, False, ],
                        [False, True, False, False, ],
                        [False, False, False, False, ],
                        [False, False, True, False, ]])
glob_mask = tf.constant([[False, False, True, False, ],
                         [False, True, False, True, ],
                         [True, False, False, False, ],
                         [False, True, True, False, ]])
loc_iou = tf.boolean_mask(iou, loc_mask)  # (2, 9)
glob_iou = tf.boolean_mask(iou, glob_mask)  # (6, 9)

实际上就是要找到loc_iou中元素对应glob_iou中的位置。如果没有loc_mask和glob_mask实际上很难找到。

问题解决

我想了半天，终于发现一个简单的方式，先利用where和boolmask找到有效位置，再使用gather_nd就完事了：

iou = tf.random.normal((4, 4, 9))
loc_mask = tf.constant([[False, False, False, False, ],
                        [False, True, False, False, ],
                        [False, False, False, False, ],
                        [False, False, True, False, ]])
glob_mask = tf.constant([[False, False, True, False, ],
                         [False, True, False, True, ],
                         [True, False, False, False, ],
                         [False, True, True, False, ]])
loc_iou = tf.boolean_mask(iou, loc_mask)  # (2, 9)
glob_iou = tf.boolean_mask(iou, glob_mask)  # (6, 9)

idx = tf.where(tf.boolean_mask(loc_mask, glob_mask)) # (2, 1) [[1],[5]]
mask_iou = tf.gather_nd(tf.boolean_mask(iou, glob_mask), idx, 0) # (2, 9)