如何使用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()