python - 如何使用 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 上运行的东西。

该代码只是一个示例，我需要一些不仅适用于圆形而且适用于任何其他形状的东西。

谢谢！

标签： pythonnumpytensorflow

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()