Next, rescale the logits to the original image size and apply argmax on the class dimension:

logits = tf.transpose(logits, [0, 2, 3, 1])
upsampled_logits = tf.image.resize(
     logits,
     # We reverse the shape of image because image.size returns width and height.
     image.size[::-1],
 )
pred_seg = tf.math.argmax(upsampled_logits, axis=-1)[0]