Как извлечь вектор точных координат изображения в тензорном потоке?

Question

Я пытаюсь написать код для модели генератора GAN с использованием керас с бэкэндом тензорного потока. Я хочу, чтобы на выходе генератора были векторные (для каждого изображения в пакете одинакового размера) значения изображения в точных координатах. Эти координаты даны как входные данные для генератора.

Я пытался использовать tf.gather_nd в качестве функции для выполнения операции, подобной крошечной, при извлечении значений из точных координат.

img является сгенерированным из шумового изображения с формой = (?, 28,28,1),

coordinates - это входной тензор формы (?, 80,2) с 80 точками, извлекаемыми из сгенерированного изображения img,

vect - выходной вектор, должен иметь размер (?, 80), где ? размер партии.

vect = Lambda(lambda x: tf.gather_nd(x, tf.cast(coordinates, 'int64')))(img)

Наконец, форма вывода этой функции (?, 80,28,1) вместо (?, 80).

Как лучше извлечь такие очки?

Answer 1

Вы можете сделать это с помощью tf.gather_nd следующим образом:

import tensorflow as tf

def extract_pixels(img, coords):
    # Number of images and pixels
    s = tf.shape(coords, out_type=coords.dtype)
    n = s[0]
    p = s[1]
    # Make gather index
    i = tf.range(n)
    ii = tf.tile(i[:, tf.newaxis, tf.newaxis], [1, p, 1])
    idx = tf.concat([ii, coords], axis=-1)
    # Gather pixel values
    pixels = tf.gather_nd(tf.squeeze(img, axis=-1), idx)
    return pixels

# ...
vect = Lambda(lambda x: extract_pixels(x, tf.cast(coordinates, 'int64')))(img)