def solarize(image, threshold=128):
    threshold = tf.saturate_cast(threshold, image.dtype)
    return tf.where((image < threshold), image, (255 - image))
def solarize(image, threshold=128):
    threshold = tf.saturate_cast(threshold, image.dtype)
    return tf.where((image < threshold), image, (255 - image))