Функция изменения формы в tf.while_loop(): немаксимальное подавление

Question

Я работаю над пользовательской реализацией чисто тензорного потока tf.image.non_max_suppression. Мне нужно написать собственную реализацию, чтобы немного изменить алгоритм на основе вывода.

Я пытаюсь воспроизвести алгоритм из реализации Python и хочу преобразовать его в эквивалентный код тензорного потока.

Реализация рабочего питона

# keep looping while some indexes still remain in the indexes
# list
while len(idxs) > 0:
    # grab the last index in the indexes list and add the
    # index value to the list of picked indexes
    last = len(idxs) - 1
    i = idxs[last]
    pick.append(i)

    # find the largest (x, y) coordinates for the start of
    # the bounding box and the smallest (x, y) coordinates
    # for the end of the bounding box
    xx1 = np.maximum(x1[i], x1[idxs[:last]])
    yy1 = np.maximum(y1[i], y1[idxs[:last]])
    xx2 = np.minimum(x2[i], x2[idxs[:last]])
    yy2 = np.minimum(y2[i], y2[idxs[:last]])

    # compute the width and height of the bounding box
    w = np.maximum(0, xx2 - xx1 + 1)
    h = np.maximum(0, yy2 - yy1 + 1)

    # compute the ratio of overlap
    overlap = (w * h) / area[idxs[:last]]

    # delete all indexes from the index list that have
    idxs = np.delete(idxs, np.concatenate(([last],
        np.where(overlap > overlapThresh)[0])))

Мой текущий код

import tensorflow as tf
import numpy as np    

def run_while_loop(x1, y1, x2, y2, area, idxs, pick, overlapThresh):
    # grab the last index in the indexes list and add the
    # index value to the list of picked indexes
    last = tf.size(idxs) - 1
    i = idxs[last]

    # find the largest (x, y) coordinates for the start of
    # the bounding box and the smallest (x, y) coordinates
    # for the end of the bounding box
    xx1 = tf.maximum(tf.gather(x1, i), tf.gather(x1, idxs[:last]))
    yy1 = tf.maximum(tf.gather(y1, i), tf.gather(y1, idxs[:last]))
    xx2 = tf.minimum(tf.gather(x2, i), tf.gather(x2, idxs[:last]))
    yy2 = tf.minimum(tf.gather(y2, i), tf.gather(y2, idxs[:last]))

    # compute the width and height of the bounding box
    w = tf.maximum(0, xx2 - xx1 + 1)
    h = tf.maximum(0, yy2 - yy1 + 1)

    # compute the ratio of overlap
    overlap = tf.divide(w * h, tf.gather(area, idxs[:last]))

    # get box ids to remove
    mask = overlap > overlapThresh
    removal_mask = tf.concat(mask, [False])
    idxs = tf.boolean_mask(idxs, removal_mask)
    pick.append(i)
    return tf.constant(i, dtype=tf.int32), idxs, pick

def is_not_empty(idxs, pick):
    return tf.size(idxs) > 0

def get_new_box_predictions_custom(self):

    # handle no bounding box condition -> change placeholder also

    pick = tf.Variable(tf.placeholder(tf.int32))

    x1, y1, x2, y2 = tf.split(self.pred_boxes_placeholder, [1, 1, 1, 1], axis=-1)
    self.x1 = tf.squeeze(x1)
    self.y1 = tf.squeeze(y1)
    self.x2 = tf.squeeze(x2)
    self.y2 = tf.squeeze(y2)

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    _, idxs = tf.nn.top_k(y2, k=y2.get_shape().as_list()[0])
    i, idxs, pick = tf.while_loop(is_not_empty, run_while_loop, [idxs, pick])

1. есть параметры изменения формы внутри цикла while (pick а также idxs, Как справиться с этим?

Окончательный результат должен быть pick переменная.

Любое рабочее решение поможет. Спасибо.