def get_num_image_channels(module_or_spec, signature=None, input_name=None):
    """Returns expected num_channels dimensions of an image input.

  This is for advanced users only who expect to handle modules with
  image inputs that might not have the 3 usual RGB channels.

  TODO(b/139530454): This does not work yet with TF2.

  Args:
    module_or_spec: a Module or ModuleSpec that accepts image inputs.
    signature: a string with the key of the signature in question.
      If None, the default signature is used.
    input_name: a string with the input name for images. If None, the
      conventional input name `images` for the default signature is used.

  Returns:
    An integer with the number of input channels to the module.

  Raises:
    ValueError: If the channel information is missing or malformed.
  """
    if input_name is None:
        input_name = 'images'
    input_info_dict = module_or_spec.get_input_info_dict(signature)
    try:
        shape = input_info_dict[input_name].get_shape()
    except KeyError:
        raise ValueError("Module is missing input '%s' in signature '%s'." %
            (input_name, signature or 'default'))
    try:
        _, _, _, num_channels = shape.as_list()
        if num_channels is None:
            raise ValueError
    except ValueError:
        raise ValueError(
            'Shape of module input is %s, expected [batch_size, height, width, num_channels] with known num_channels'
             % shape)
    return num_channels
