Tensorflow：按最大值过滤3D索引重复项

Question

I am trying to create a filter mask that removes duplicate Indices from an vector by comparing which of their respective values is greater.

My current approach is:

1. Transform 3-D Index to 1-D
2. Check the 1-D Index for uniqueness
3. Calculate the maximum values of each unique index
4. Compare the maximum values with the original values. If the same value exists, keep that 3-D Index.

I want to get an filter array so I can apply a boolean_mask to other tensors as well. For this example the mask should look the following way: [False True True True True] .

My current code kind of works unless the values themselves are also duplicated. However this seems to be the case when I am using it therefore I need to find a better solution to it.

Here is an examplary of how my Code looks

import tensorflow as tf

# Dummy Input values with same Structure as the real
x_cells   = tf.constant([1,2,3,4,1], dtype=tf.int32)   # Index_1
y_cells   = tf.constant([4,4,4,4,4], dtype=tf.int32)   # Index_2
iou_index = tf.constant([1,2,3,4,1], dtype=tf.int32) # Index_3
iou_max   = tf.constant([1.,2.,3.,4.,5.], dtype=tf.float32) # Values

# my Output should be a mask that is [False True True True True]
# So if i filter this i get e.g. x_cells = [2,3,4,1] or iou_max = [2.,3.,4.,5.]

max_dim_y = tf.constant(10)
max_dim_x = tf.constant(20)
num_anchors = 5
stride = 32

---

# 1. Transforming the 3D-Index to 1D
tmp = tf.stack([x_cells, y_cells, iou_index], axis=1)
indices = tf.matmul(tmp, [[max_dim_y * num_anchors],     [num_anchors],[1]])

# 2. Looking for unique / duplicate indices
y, idx = tf.unique(tf.squeeze(indices))

# 3. Calculating the maximum values of each unique index.
# An function like unsorted_segment_argmax() would be awesome here
num_segments = tf.shape(y)[0]
ious = tf.unsorted_segment_max(iou_max, idx, num_segments)

iou_max_length = tf.shape(iou_max)[0]
ious_length = tf.shape(ious)[0]

# 4. Compare all max values to original values.
iou_max_tiled = tf.tile(iou_max, [ious_length])
iou_reshaped = tf.reshape(iou_max_tiled, [ious_length, iou_max_length])
iou_max_reshaped = tf.transpose(iou_reshaped)
filter_mask = tf.reduce_any(tf.equal(iou_max_reshaped, ious), -1)
filter_mask = tf.reshape(filter_mask, shape=[-1])

---

This code above will fail if we simply change the value of the iou_max Variable in the beginning to:

x_cells = tf.constant([1,2,3,4,1], dtype=tf.int32)
y_cells = tf.constant([4,4,4,4,4], dtype=tf.int32)
iou_index = tf.constant([1,2,3,4,1], dtype=tf.int32)
iou_max = tf.constant([2.,2.,3.,4.,5.], dtype=tf.float32)

---

Answer 1

My current workaround changed point 4 of my question:

Basically I changed that I compare the tuples instead of the single values. This leads to me beeing able to logically check if both, index AND value are in the remaining values from 3.

# 4. Compare a Max Value and Indices with original values
rem_index_val_pair = tf.stack([ious, tf.cast(y, dtype=tf.float32)], axis=1)
orig_val_index_pair = tf.stack([iou_max, tf.cast(indices, dtype=tf.float32)], axis=1)

orig_val_index_pair_t = tf.tile(orig_val_index_pair, [1, ious_length])
orig_val_index_pair_s = tf.reshape(orig_val_index_pair_t, [iou_max_length, ious_length, 2])
filter_mask_1 = tf.equal(orig_val_index_pair_s, rem_index_val_pair)
filter_mask_2 = tf.reduce_all(filter_mask_1, -1)
filter_mask_3 = tf.reduce_any(filter_mask_2, -1)

---

# The orig_val_index_pair_s looks like the following
a =  [[[  2.  71.][  2.  71.][  2.  71.][  2.  71.]
     [[  2. 122.][  2. 122.][  2. 122.][  2. 122.]]
     [[  3. 173.][  3. 173.][  3. 173.][  3. 173.]]
     [[  4. 224.][  4. 224.][  4. 224.][  4. 224.]]
     [[  5.  71.][  5.  71.][  5.  71.][  5.  71.]]]
# I then compare it to the rem_max_val_pair which looks like this.
b =  [[  5.  71.][  2. 122.][  3. 173.][  4. 224.]]

# Using equal(a,b) will now compare each of the values resulting in:
c = [[[False  True][ True False][False False][False False]]
     [[False False][ True  True][False False][False False]]
     [[False False][False False][ True  True][False False]]
     [[False False][False False][False False][ True  True]]
     [[ True  True][False False][False False][False False]]]

# Using tf.reduce_all(c, -1) I can filter the bool pairs with a logical And. 
# (This kicks out my false positives from before).
# Afterwards I can check if the line has any true value by tf.reduce_any().

---

IMO this solution still is a dirty workaround. So if you have any better solution proposals please share them. :)