code: def generate_heatmap(logits, num_classes):
      """Generates a coloured heatmap from the keypoint logits.
  
      Args:
          features: A `Tensor` of dimensions [num_batch, height, width, FLAGS.n_landmarks + 1].
      """
  
      keypoint_colours = np.array(
          [
              plt.cm.spectral(x)
              for x in np.linspace(0, 1, num_classes + 1)
          ]
      )[..., :3].astype(np.float32)
  
      prediction = tf.nn.softmax(logits)
      heatmap = tf.matmul(
          tf.reshape(prediction, (-1, num_classes + 1)),
          keypoint_colours,
      )
      heatmap = tf.reshape(
          heatmap,
          (
              tf.shape(prediction)[0],
              tf.shape(prediction)[1],
              tf.shape(prediction)[2],
              3,
          ),
      )
      return heatmap
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import tensorflow
  import menpo.io
  from menpo.image import Image
  from menpo.shape import PointCloud
  import cv2
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.contrib.framework.python.ops import variables
  from tensorflow.python.training import optimizer
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import variables
  from menpo.transform import Translation, Scale
  from menpo.shape import PointCloud
label_desc: Use image node_id 1 for calling a global function with example usage: generate_heatmap(logits, num_classes) and returns: heatmap