Доступ к пикселям изображения в тензорном

Question

Я пытаюсь написать тензорную функцию потока для нормализации данных изображения перед тренировкой. Функция примет маску изображения и преобразует все значения пикселей ниже 30 в «0», а все остальное в «1».

Чтобы сделать это вне тензора, мы можем использовать numpy:

arr = np.array(mask)
arr[arr < 30] = 0
arr[arr >= 30] = 1
mask.putdata(arr)

Но при попытке сделать это в функции tf:

def normalize(input_image, input_mask):
  # Transform the actual image to float values between 0 -> 1
  input_image = tf.cast(input_image, tf.float32) / 255.0 

  # Modify the mask to only contain values 0 and 1
  arr = np.array(input_mask)
  arr[arr < 30] = 0
  arr[arr >= 30] = 1
  input_mask.putdata(arr)

  return input_image, input_mask

@tf.function
def load_image_train(input_image, input_mask):
  input_image = tf.image.resize(input_image, (320, 320))
  input_mask = tf.image.resize(input_mask, (320, 320))

  if tf.random.uniform(()) > 0.5:
    input_image = tf.image.flip_left_right(input_image)
    input_mask = tf.image.flip_left_right(input_mask)

  input_image, input_mask = normalize(input_image, input_mask)

  return input_image, input_mask

train = labeled_ds.map(load_image_train, num_parallel_calls=AUTOTUNE)

я получаю ошибку NotImplementedError: Cannot convert a symbolic Tensor (cond/Identity_1:0) to a numpy array.

Как я могу сделать это правильно?