Как исключить все значения вне круга изображения с помощью операций TensorFlow?

Question

Это то, что я имею в виду: используя следующий код, я могу сделать так, чтобы значения за пределами круга становились равными 0. Код генерирует полностью белое изображение и устанавливает значения за пределами круга равными нулю.

import numpy as np
import matplotlib.pyplot as plt

width = 512
all_white_img = np.zeros(shape=[width, width], dtype=np.float)
all_white_img[:] = 1
plt.imshow(all_white_img, cmap='gray', vmax=1.0, vmin=0.0)
plt.show()

[X, Y] = np.mgrid[0:width, 0:width]
xpr = X - int(width) // 2
ypr = Y - int(width) // 2
radius = width // 2
reconstruction_circle = (xpr ** 2 + ypr ** 2) <= radius ** 2 #set circle
all_white_img[~reconstruction_circle] = 0.
plt.imshow(all_white_img, cmap='gray', vmax=1.0, vmin=0.0)
plt.show()

Вывод изображений:

Как сделать то же самое, эффективно используя TensorFlow?

Поскольку numpy работает наCPU и мне нужно что-то, что может работать на GPU.

Код является лишь примером, мне нужно что-то, что работает не только для кругов, но и для любых других фигур.

Спасибо!

Answer 1

Функция tf.where делает именно то, что я хочу.

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf

width = 512
sess = tf.Session()

[X, Y] = np.mgrid[0:width, 0:width]
xpr = X - int(width) // 2
ypr = Y - int(width) // 2
radius = width // 2
reconstruction_circle = (xpr ** 2 + ypr ** 2) <= radius ** 2  #set circle

all_white_img = tf.ones(shape=[width, width], dtype=tf.float32)
plt.imshow(sess.run(all_white_img), cmap='gray', vmax=1.0, vmin=0.0)
plt.show()
reconstruction_circle = tf.cast(reconstruction_circle, tf.bool)
all_white_img = tf.where(reconstruction_circle, all_white_img, tf.zeros_like(all_white_img))
plt.imshow(sess.run(all_white_img), cmap='gray', vmax=1.0, vmin=0.0)
plt.show()