in 04_detect_segment/utils_box.py [0:0]
def n_largest_rois_in_cell(tile, rois, rois_n, grid_n, n, comparator="largest_w", expand=1.0):
# handle the case of rois_n == 0 by creating one dummy empty roi, otherwise the code will not work with rois_n=0
rois, rois_n = tf.cond(tf.equal(rois_n, 0),
true_fn=lambda: (tf.constant([[0.0, 0.0, 0.0, 0.0]]), tf.constant(1)),
false_fn=lambda: (rois, rois_n))
grid, cell_w = gen_grid_for_tile(tile, grid_n)
# grid shape [grid_n, grid_n, 2]
# rois shape [rois_n, 3]
rois = x1y1x2y2_to_cxcyw(rois)
cross_rois = reshape_rois(rois, grid_n) # shape [grid_n, grid_n, rois_n, 3]]
cross_rois_cx, cross_rois_cy, cross_rois_w = tf.unstack(cross_rois, axis=-1) # shape [grid_n, grid_n, rois_n]]
has_center = center_in_grid_cell(grid, grid_n, cell_w, rois, expand=expand)
grid_centers = (grid + grid + cell_w) / 2.0 # shape [grid_n, grid_n, 2]
g_cx, g_cy = tf.unstack(grid_centers, axis=-1) # shape [grid_n, grid_n]
g_cx = tf.expand_dims(g_cx, axis=-1) # force broadcasting on correct axis
g_cy = tf.expand_dims(g_cy, axis=-1)
# iterate on largest a fixed number of times to get N largest
n_largest = []
zeros = tf.zeros(shape=[grid_n, grid_n, 3])
for i in range(n):
any_roi_in_cell = tf.reduce_any(has_center, axis=2) # shape [grid_n, grid_n]
if comparator=="largest_w":
largest_indices = tf.argmax(tf.cast(has_center, tf.float32) * cross_rois_w, axis=2) # shape [grid_n, grid_n]
elif comparator=="furthest_from_center":
d_from_cell_center = tf.abs(cross_rois_cx - g_cx) + tf.abs(cross_rois_cy - g_cy)
largest_indices = tf.argmax(tf.cast(has_center, tf.float32) * d_from_cell_center, axis=2) # shape [grid_n, grid_n]
elif comparator=="closest_to_center":
d_from_cell_center = tf.abs(cross_rois_cx - g_cx) + tf.abs(cross_rois_cy - g_cy)
ones = tf.ones(tf.shape(d_from_cell_center))
largest_indices = tf.argmin(tf.where(has_center, d_from_cell_center, 1000*ones), axis=2) # shape [grid_n, grid_n]
# as of TF1.3 can use tf.gather(axis=2)
rs_largest_indices = tf.reshape(largest_indices, [grid_n*grid_n])
rs_largest_indices = tf.unstack(rs_largest_indices, axis=0) # list
rs_cross_rois = tf.reshape(cross_rois, [grid_n*grid_n, rois_n, 3])
rs_cross_rois = tf.unstack(rs_cross_rois, axis=0) # list
rs_largest_roi_in_cell = [tf.gather(cr, li) for cr, li in zip(rs_cross_rois, rs_largest_indices)]
largest_roi_in_cell = tf.stack(rs_largest_roi_in_cell, axis=0) # shape [grid_n * grid_n, 3]
largest_roi_in_cell = tf.reshape(largest_roi_in_cell, [grid_n, grid_n, 3]) # shape [grid_n, grid_n, 3]
# cells that do not have a roi in them, set their "largest roi in cell" to (x=0,y=0,w=0)
any_roi_in_cell = tf.tile(tf.expand_dims(any_roi_in_cell, axis=-1), [1, 1, 3]) # shape [grid_n, grid_n, 3]
largest_roi_in_cell = tf.where(any_roi_in_cell, largest_roi_in_cell, zeros) # shape [grid_n, grid_n, 3]
n_largest.append(largest_roi_in_cell)
# zero-out the largest element per cell to get the next largest on the next iteration
zero_mask = tf.logical_not(tf.cast(tf.one_hot(largest_indices, rois_n), dtype=tf.bool))
has_center = tf.logical_and(has_center, zero_mask)
n_largest = tf.stack(n_largest, axis=2) # shape [grid_n, grid_n, n, 3]
return n_largest # shape [grid_n, grid_n, n, 3]