jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def to_4d(image: tf.Tensor) -> tf.Tensor:
  """Converts an input Tensor to 4 dimensions.

  4D image => [N, H, W, C] or [N, C, H, W]
  3D image => [1, H, W, C] or [1, C, H, W]
  2D image => [1, H, W, 1]

  Args:
    image: The 2/3/4D input tensor.

  Returns:
    A 4D image tensor.

  Raises:
    `TypeError` if `image` is not a 2/3/4D tensor.

  """
  shape = tf.shape(image)
  original_rank = tf.rank(image)
  left_pad = tf.to_int32(tf.less_equal(original_rank, 3))
  right_pad = tf.to_int32(tf.equal(original_rank, 2))
  new_shape = tf.concat(
      [
          tf.ones(shape=left_pad, dtype=tf.int32),
          shape,
          tf.ones(shape=right_pad, dtype=tf.int32),
      ],
      axis=0,
  )
  return tf.reshape(image, new_shape)
//  [python]
def to_4d(image: tf.Tensor) -> tf.Tensor:
  """Converts an input Tensor to 4 dimensions.

  4D image => [N, H, W, C] or [N, C, H, W]
  3D image => [1, H, W, C] or [1, C, H, W]
  2D image => [1, H, W, 1]

  Args:
    image: The 2/3/4D input tensor.

  Returns:
    A 4D image tensor.

  Raises:
    `TypeError` if `image` is not a 2/3/4D tensor.

  """
  shape = tf.shape(image)
  original_rank = tf.rank(image)
  left_pad = tf.cast(tf.less_equal(original_rank, 3), dtype=tf.int32)
  right_pad = tf.cast(tf.equal(original_rank, 2), dtype=tf.int32)
  new_shape = tf.concat(
      [
          tf.ones(shape=left_pad, dtype=tf.int32),
          shape,
          tf.ones(shape=right_pad, dtype=tf.int32),
      ],
      axis=0,
  )
  return tf.reshape(image, new_shape)
//  [python]
def rotate(image: tf.Tensor, degrees: float) -> tf.Tensor:
  """Rotates the image by degrees either clockwise or counterclockwise.

  Args:
    image: An image Tensor of type uint8.
    degrees: Float, a scalar angle in degrees to rotate all images by. If
      degrees is positive the image will be rotated clockwise otherwise it will
      be rotated counterclockwise.

  Returns:
    The rotated version of image.

  """
  # Convert from degrees to radians.
  degrees_to_radians = math.pi / 180.0
  radians = t.cast(degrees * degrees_to_radians, t.float32)

  original_ndims = tf.rank(image)
  image = to_4d(image)

  image_height = t.cast(tf.shape(image)[1], t.float32)
  image_width = t.cast(tf.shape(image)[2], t.float32)
  transforms = _convert_angles_to_transform(
      angles=radians, image_width=image_width, image_height=image_height)
  # In practice, we should randomize the rotation degrees by flipping
  # it negatively half the time, but that's done on 'degrees' outside
  # of the function.
  image = transform(image, transforms=transforms)
  return from_4d(image, original_ndims)
//  [python]
def rotate(image: tf.Tensor, degrees: float) -> tf.Tensor:
  """Rotates the image by degrees either clockwise or counterclockwise.

  Args:
    image: An image Tensor of type uint8.
    degrees: Float, a scalar angle in degrees to rotate all images by. If
      degrees is positive the image will be rotated clockwise otherwise it will
      be rotated counterclockwise.

  Returns:
    The rotated version of image.

  """
  # Convert from degrees to radians.
  degrees_to_radians = math.pi / 180.0
  radians = tf.cast(degrees * degrees_to_radians, tf.float32)

  original_ndims = tf.rank(image)
  image = to_4d(image)

  image_height = tf.cast(tf.shape(image)[1], tf.float32)
  image_width = tf.cast(tf.shape(image)[2], tf.float32)
  transforms = _convert_angles_to_transform(
      angles=radians, image_width=image_width, image_height=image_height)
  # In practice, we should randomize the rotation degrees by flipping
  # it negatively half the time, but that's done on 'degrees' outside
  # of the function.
  image = transform(image, transforms=transforms)
  return from_4d(image, original_ndims)
//  [python]
def solarize_add(image: tf.Tensor,
                 addition: int = 0,
                 threshold: int = 128) -> tf.Tensor:
  # For each pixel in the image less than threshold
  # we add 'addition' amount to it and then clip the
  # pixel value to be between 0 and 255. The value
  # of 'addition' is between -128 and 128.
  added_image = t.cast(image, tf.int64) + addition
  added_image = t.cast(tf.clip_by_value(added_image, 0, 255), tf.uint8)
  return tf.where(image < threshold, added_image, image)
//  [python]
def solarize_add(image: tf.Tensor,
                 addition: int = 0,
                 threshold: int = 128) -> tf.Tensor:
  # For each pixel in the image less than threshold
  # we add 'addition' amount to it and then clip the
  # pixel value to be between 0 and 255. The value
  # of 'addition' is between -128 and 128.
  added_image = tf.cast(image, tf.int64) + addition
  added_image = tf.cast(tf.clip_by_value(added_image, 0, 255), tf.uint8)
  return tf.where(image < threshold, added_image, image)
//  [python]
def blend(image1: tf.Tensor, image2: tf.Tensor, factor: float) -> tf.Tensor:
  """Blend image1 and image2 using 'factor'.

  Factor can be above 0.0.  A value of 0.0 means only image1 is used.
  A value of 1.0 means only image2 is used.  A value between 0.0 and
  1.0 means we linearly interpolate the pixel values between the two
  images.  A value greater than 1.0 "extrapolates" the difference
  between the two pixel values, and we clip the results to values
  between 0 and 255.

  Args:
    image1: An image Tensor of type uint8.
    image2: An image Tensor of type uint8.
    factor: A floating point value above 0.0.

  Returns:
    A blended image Tensor of type uint8.
  """
  if factor == 0.0:
    return t.convert_to_tensor(image1)
  if factor == 1.0:
    return t.convert_to_tensor(image2)

  image1 = t.cast(image1, t.float32)
  image2 = t.cast(image2, t.float32)

  difference = image2 - image1
  scaled = factor * difference

  # Do addition in float.
  temp = t.cast(image1, t.float32) + scaled

  # Interpolate
  if factor > 0.0 and factor < 1.0:
    # Interpolation means we always stay within 0 and 255.
    return t.cast(temp, t.uint8)

  # Extrapolate:
  #
  # We need to clip and then cast.
  return t.cast(tf.clip_by_value(temp, 0.0, 255.0), t.uint8)
//  [python]
def blend(image1: tf.Tensor, image2: tf.Tensor, factor: float) -> tf.Tensor:
  """Blend image1 and image2 using 'factor'.

  Factor can be above 0.0.  A value of 0.0 means only image1 is used.
  A value of 1.0 means only image2 is used.  A value between 0.0 and
  1.0 means we linearly interpolate the pixel values between the two
  images.  A value greater than 1.0 "extrapolates" the difference
  between the two pixel values, and we clip the results to values
  between 0 and 255.

  Args:
    image1: An image Tensor of type uint8.
    image2: An image Tensor of type uint8.
    factor: A floating point value above 0.0.

  Returns:
    A blended image Tensor of type uint8.
  """
  if factor == 0.0:
    return tf.convert_to_tensor(image1)
  if factor == 1.0:
    return tf.convert_to_tensor(image2)

  image1 = tf.cast(image1, tf.float32)
  image2 = tf.cast(image2, tf.float32)

  difference = image2 - image1
  scaled = factor * difference

  # Do addition in float.
  temp = tf.cast(image1, tf.float32) + scaled

  # Interpolate
  if factor > 0.0 and factor < 1.0:
    # Interpolation means we always stay within 0 and 255.
    return tf.cast(temp, tf.uint8)

  # Extrapolate:
  #
  # We need to clip and then cast.
  return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
//  [python]
def contrast(image: tf.Tensor, factor: float) -> tf.Tensor:
  """Equivalent of PIL Contrast."""
  degenerate = tf.image.rgb_to_grayscale(image)
  # Cast before calling tf.histogram.
  degenerate = tf.to_int32(degenerate)

  # Compute the grayscale histogram, then compute the mean pixel value,
  # and create a constant image size of that value.  Use that as the
  # blending degenerate target of the original image.
  hist = tf.histogram_fixed_width(degenerate, [0, 255], nbins=256)
  mean = tf.reduce_sum(tf.to_float(hist)) / 256.0
  degenerate = tf.ones_like(degenerate, dtype=tf.float32) * mean
  degenerate = tf.clip_by_value(degenerate, 0.0, 255.0)
  degenerate = tf.image.grayscale_to_rgb(tf.to_uint8(degenerate))
  return blend(degenerate, image, factor)
//  [python]
def contrast(image: tf.Tensor, factor: float) -> tf.Tensor:
  """Equivalent of PIL Contrast."""
  degenerate = tf.image.rgb_to_grayscale(image)
  # Cast before calling tf.histogram.
  degenerate = tf.cast(degenerate, tf.int32)

  # Compute the grayscale histogram, then compute the mean pixel value,
  # and create a constant image size of that value.  Use that as the
  # blending degenerate target of the original image.
  hist = tf.histogram_fixed_width(degenerate, [0, 255], nbins=256)
  mean = tf.reduce_sum(tf.cast(hist, tf.float32)) / 256.0
  degenerate = tf.ones_like(degenerate, dtype=tf.float32) * mean
  degenerate = tf.clip_by_value(degenerate, 0.0, 255.0)
  degenerate = tf.image.grayscale_to_rgb(tf.cast(degenerate, tf.uint8))
  return blend(degenerate, image, factor)
//  [python]
def _randomly_negate_tensor(tensor):
  """With 50% prob turn the tensor negative."""
  should_flip = tf.floor(tf.random.uniform([]) + 0.5) > 0
  final_tensor = tf.cond(should_flip, lambda: tensor, lambda: -tensor)
  return final_tensor
//  [python]
def _randomly_negate_tensor(tensor):
  """With 50% prob turn the tensor negative."""
  should_flip = tf.cast(tf.floor(tf.random.uniform([]) + 0.5), tf.bool)
  final_tensor = tf.cond(should_flip, lambda: tensor, lambda: -tensor)
  return final_tensor
//  [python]
def _apply_func_with_prob(func: Any, image: tf.Tensor, args: Any, prob: float):
  """Apply `func` to image w/ `args` as input with probability `prob`."""
  assert isinstance(args, tuple)

  # Apply the function with probability `prob`.
  should_apply_op = tf.floor(tf.random.uniform([], dtype=tf.float32) + prob) > 0
  augmented_image = tf.cond(should_apply_op, lambda: func(image, *args),
                            lambda: image)
  return augmented_image
//  [python]
def _apply_func_with_prob(func: Any, image: tf.Tensor, args: Any, prob: float):
  """Apply `func` to image w/ `args` as input with probability `prob`."""
  assert isinstance(args, tuple)

  # Apply the function with probability `prob`.
  should_apply_op = tf.cast(
      tf.floor(tf.random.uniform([], dtype=tf.float32) + prob), tf.bool)
  augmented_image = tf.cond(should_apply_op, lambda: func(image, *args),
                            lambda: image)
  return augmented_image
//  [python]
def autocontrast(image: tf.Tensor) -> tf.Tensor:
  """Implements Autocontrast function from PIL using TF ops.

  Args:
    image: A 3D uint8 tensor.

  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

  def scale_channel(image: tf.Tensor) -> tf.Tensor:
    """Scale the 2D image using the autocontrast rule."""
    # A possibly cheaper version can be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = t.cast(t.reduce_min(image), t.float32)
    hi = t.cast(t.reduce_max(image), t.float32)

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = t.cast(im, t.float32) * scale + offset
      im = t.clip_by_value(im, 0.0, 255.0)
      return t.cast(im, t.uint8)

    result = t.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = t.stack([s1, s2, s3], 2)
  return image
//  [python]
def autocontrast(image: tf.Tensor) -> tf.Tensor:
  """Implements Autocontrast function from PIL using TF ops.

  Args:
    image: A 3D uint8 tensor.

  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

  def scale_channel(image: tf.Tensor) -> tf.Tensor:
    """Scale the 2D image using the autocontrast rule."""
    # A possibly cheaper version can be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = tf.cast(tf.reduce_min(image), tf.float32)
    hi = tf.cast(tf.reduce_max(image), tf.float32)

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = tf.cast(im, tf.float32) * scale + offset
      im = tf.clip_by_value(im, 0.0, 255.0)
      return tf.cast(im, tf.uint8)

    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def sharpness(image: tf.Tensor, factor: float) -> tf.Tensor:
  """Implements Sharpness function from PIL using TF ops."""
  orig_image = image
  image = t.cast(image, tf.float32)
  # Make image 4D for conv operation.
  image = tf.expand_dims(image, 0)
  # SMOOTH PIL Kernel.
  kernel = tf.constant([[1, 1, 1], [1, 5, 1], [1, 1, 1]],
                       dtype=tf.float32,
                       shape=[3, 3, 1, 1]) / 13.
  # Tile across channel dimension.
  kernel = tf.tile(kernel, [1, 1, 3, 1])
  strides = [1, 1, 1, 1]
  degenerate = tf.nn.depthwise_conv2d(
      image, kernel, strides, padding='VALID', dilations=[1, 1])
  degenerate = tf.clip_by_value(degenerate, 0.0, 255.0)
  degenerate = tf.squeeze(t.cast(degenerate, tf.uint8), [0])

  # For the borders of the resulting image, fill in the values of the
  # original image.
  mask = tf.ones_like(degenerate)
  padded_mask = tf.pad(mask, [[1, 1], [1, 1], [0, 0]])
  padded_degenerate = tf.pad(degenerate, [[1, 1], [1, 1], [0, 0]])
  result = tf.where(tf.equal(padded_mask, 1), padded_degenerate, orig_image)

  # Blend the final result.
  return blend(result, orig_image, factor)
//  [python]
def sharpness(image: tf.Tensor, factor: float) -> tf.Tensor:
  """Implements Sharpness function from PIL using TF ops."""
  orig_image = image
  image = tf.cast(image, tf.float32)
  # Make image 4D for conv operation.
  image = tf.expand_dims(image, 0)
  # SMOOTH PIL Kernel.
  kernel = tf.constant([[1, 1, 1], [1, 5, 1], [1, 1, 1]],
                       dtype=tf.float32,
                       shape=[3, 3, 1, 1]) / 13.
  # Tile across channel dimension.
  kernel = tf.tile(kernel, [1, 1, 3, 1])
  strides = [1, 1, 1, 1]
  degenerate = tf.nn.depthwise_conv2d(
      image, kernel, strides, padding='VALID', dilations=[1, 1])
  degenerate = tf.clip_by_value(degenerate, 0.0, 255.0)
  degenerate = tf.squeeze(tf.cast(degenerate, tf.uint8), [0])

  # For the borders of the resulting image, fill in the values of the
  # original image.
  mask = tf.ones_like(degenerate)
  padded_mask = tf.pad(mask, [[1, 1], [1, 1], [0, 0]])
  padded_degenerate = tf.pad(degenerate, [[1, 1], [1, 1], [0, 0]])
  result = tf.where(tf.equal(padded_mask, 1), padded_degenerate, orig_image)

  # Blend the final result.
  return blend(result, orig_image, factor)
//  [python]
def _build_scalars(self):
        with tf.name_scope('scalars'):
            # accuracy should not change, so we just add it as a sanity check
            # [batch]
            predicted = tf.argmax(self._logits, axis=-1, name='predicted')
            equal_first = tf.to_float(tf.equal(predicted, self._placeholders.labels), name='acc_equal')
            acc = tf.reduce_mean(equal_first, name='acc')

        self._scalars = dict(acc=acc)
//  [python]
def _build_scalars(self):
        with tf.name_scope('scalars'):
            # accuracy should not change, so we just add it as a sanity check
            # [batch]
            predicted = tf.argmax(self._logits, axis=-1, name='predicted')
            equal_first = tf.cast(tf.equal(predicted, self._placeholders.labels), tf.float32, name='acc_equal')
            acc = tf.reduce_mean(equal_first, name='acc')

        self._scalars = dict(acc=acc)
//  [python]
def _linear_indices_and_weights(coordinate):
    lower = tf.floor(coordinate)
    upper_weight = coordinate - lower
    lower_weight = 1 - upper_weight
    index = t.cast(lower, tf.int32)
    return [(index, lower_weight), (index + 1, upper_weight)]
//  [python]
def _linear_indices_and_weights(coordinate):
    lower = tf.floor(coordinate)
    upper_weight = coordinate - lower
    lower_weight = 1 - upper_weight
    index = tf.cast(lower, tf.int32)
    return [(index, lower_weight), (index + 1, upper_weight)]
//  [python]
def equalize(image: tf.Tensor) -> tf.Tensor:
  """Implements Equalize function from PIL using TF ops."""

  def scale_channel(im, c):
    """Scale the data in the channel to implement equalize."""
    im = t.cast(im[:, :, c], tf.int32)
    # Compute the histogram of the image channel.
    histo = tf.histogram_fixed_width(im, [0, 255], nbins=256)

    # For the purposes of computing the step, filter out the nonzeros.
    nonzero = tf.where(tf.not_equal(histo, 0))
    nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])
    step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255

    def build_lut(histo, step):
      # Compute the cumulative sum, shifting by step // 2
      # and then normalization by step.
      lut = (tf.cumsum(histo) + (step // 2)) // step
      # Shift lut, prepending with 0.
      lut = tf.concat([[0], lut[:-1]], 0)
      # Clip the counts to be in range.  This is done
      # in the C code for image.point.
      return tf.clip_by_value(lut, 0, 255)

    # If step is zero, return the original image.  Otherwise, build
    # lut from the full histogram and step and then index from it.
    result = tf.cond(
        tf.equal(step, 0), lambda: im,
        lambda: tf.gather(build_lut(histo, step), im))

    return t.cast(result, tf.uint8)

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image, 0)
  s2 = scale_channel(image, 1)
  s3 = scale_channel(image, 2)
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def equalize(image: tf.Tensor) -> tf.Tensor:
  """Implements Equalize function from PIL using TF ops."""

  def scale_channel(im, c):
    """Scale the data in the channel to implement equalize."""
    im = tf.cast(im[:, :, c], tf.int32)
    # Compute the histogram of the image channel.
    histo = tf.histogram_fixed_width(im, [0, 255], nbins=256)

    # For the purposes of computing the step, filter out the nonzeros.
    nonzero = tf.where(tf.not_equal(histo, 0))
    nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])
    step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255

    def build_lut(histo, step):
      # Compute the cumulative sum, shifting by step // 2
      # and then normalization by step.
      lut = (tf.cumsum(histo) + (step // 2)) // step
      # Shift lut, prepending with 0.
      lut = tf.concat([[0], lut[:-1]], 0)
      # Clip the counts to be in range.  This is done
      # in the C code for image.point.
      return tf.clip_by_value(lut, 0, 255)

    # If step is zero, return the original image.  Otherwise, build
    # lut from the full histogram and step and then index from it.
    result = tf.cond(
        tf.equal(step, 0), lambda: im,
        lambda: tf.gather(build_lut(histo, step), im))

    return tf.cast(result, tf.uint8)

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image, 0)
  s2 = scale_channel(image, 1)
  s3 = scale_channel(image, 2)
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    example = tf.io.parse_single_example(record, self.name_to_features_spec())

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in example:
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.int32(t)
      example[name] = t

    return example
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    example = tf.io.parse_single_example(record, self.name_to_features_spec())

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in example:
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example
//  [python]
def __convert_rgb2od(self, I, Io, beta):
        I = tf.transpose(I, [1, 2, 0])

        # calculate optical density
        OD = -tf.math.log((tf.reshape(I, [tf.math.reduce_prod(I.shape[:-1]), I.shape[-1]]) + 1) / Io)

        # remove transparent pixels
        ODhat = OD[~tf.math.reduce_any(OD < beta, axis=1)]

        return OD, ODhat
//  [python]
def __convert_rgb2od(self, I, Io, beta):
        I = tf.transpose(I, [1, 2, 0])

        # calculate optical density
        OD = -tf.math.log((tf.cast(tf.reshape(I, [tf.math.reduce_prod(I.shape[:-1]), I.shape[-1]]), tf.float32) + 1) / Io)

        # remove transparent pixels
        ODhat = OD[~tf.math.reduce_any(OD < beta, axis=1)]

        return OD, ODhat
//  [python]
def distort(self, image: tf.Tensor) -> tf.Tensor:
    """Applies the RandAugment policy to `image`.

    Args:
      image: `Tensor` of shape [height, width, 3] representing an image.

    Returns:
      The augmented version of `image`.
    """
    input_image_type = image.dtype

    if input_image_type != tf.uint8:
      image = tf.clip_by_value(image, 0.0, 255.0)
      image = t.cast(image, dtype=tf.uint8)

    replace_value = [128] * 3
    min_prob, max_prob = 0.2, 0.8

    for _ in range(self.num_layers):
      op_to_select = tf.random.uniform([],
                                       maxval=len(self.available_ops) + 1,
                                       dtype=tf.int32)

      branch_fns = []
      for (i, op_name) in enumerate(self.available_ops):
        prob = tf.random.uniform([],
                                 minval=min_prob,
                                 maxval=max_prob,
                                 dtype=tf.float32)
        func, _, args = _parse_policy_info(op_name, prob, self.magnitude,
                                           replace_value, self.cutout_const,
                                           self.translate_const)
        branch_fns.append((
            i,
            # pylint:disable=g-long-lambda
            lambda selected_func=func, selected_args=args: selected_func(
                image, *selected_args)))
        # pylint:enable=g-long-lambda

      image = tf.switch_case(
          branch_index=op_to_select,
          branch_fns=branch_fns,
          default=lambda: tf.identity(image))

    image = t.cast(image, dtype=input_image_type)
    return image
//  [python]
def distort(self, image: tf.Tensor) -> tf.Tensor:
    """Applies the RandAugment policy to `image`.

    Args:
      image: `Tensor` of shape [height, width, 3] representing an image.

    Returns:
      The augmented version of `image`.
    """
    input_image_type = image.dtype

    if input_image_type != tf.uint8:
      image = tf.clip_by_value(image, 0.0, 255.0)
      image = tf.cast(image, dtype=tf.uint8)

    replace_value = [128] * 3
    min_prob, max_prob = 0.2, 0.8

    for _ in range(self.num_layers):
      op_to_select = tf.random.uniform([],
                                       maxval=len(self.available_ops) + 1,
                                       dtype=tf.int32)

      branch_fns = []
      for (i, op_name) in enumerate(self.available_ops):
        prob = tf.random.uniform([],
                                 minval=min_prob,
                                 maxval=max_prob,
                                 dtype=tf.float32)
        func, _, args = _parse_policy_info(op_name, prob, self.magnitude,
                                           replace_value, self.cutout_const,
                                           self.translate_const)
        branch_fns.append((
            i,
            # pylint:disable=g-long-lambda
            lambda selected_func=func, selected_args=args: selected_func(
                image, *selected_args)))
        # pylint:enable=g-long-lambda

      image = tf.switch_case(
          branch_index=op_to_select,
          branch_fns=branch_fns,
          default=lambda: tf.identity(image))

    image = tf.cast(image, dtype=input_image_type)
    return image
//  [python]
def split_time(dataset, idx):
  batch_idx = t.cast(idx, t.int64)
  train_data_time = dataset.take(batch_idx)
  test_data_time = dataset.skip(batch_idx).take(1)

  return train_data_time, test_data_time
//  [python]
def split_time(dataset, idx):
  batch_idx = tf.cast(idx, tf.int64)
  train_data_time = dataset.take(batch_idx)
  test_data_time = dataset.skip(batch_idx).take(1)

  return train_data_time, test_data_time
//  [python]
def _define_optimizer_and_cost(self):
        """This function defines the optimizer and cost function of the LSTM model

        Returns:
            AdamOptimizer, Tensor: The optimizer function to reduce loss and the loss values
        """
        if self.use_crf_layer:
            flattened_labels = tf.to_int32(tf.argmax(self.label_tf, axis=2))
            log_likelihood, _ = tf.contrib.crf.crf_log_likelihood(
                self.lstm_output_tf, flattened_labels, self.batch_sequence_lengths_tf
            )
            cost_tf = tf.reduce_mean(-log_likelihood, name="cost_tf")
        else:
            masked_logits = tf.boolean_mask(
                tf.reshape(self.lstm_output_tf, [-1, self.output_dimension]),
                self.batch_sequence_mask_tf,
            )

            masked_labels = tf.boolean_mask(
                tf.reshape(self.label_tf, [-1, self.output_dimension]),
                self.batch_sequence_mask_tf,
            )

            softmax_loss_tf = tf.nn.softmax_cross_entropy_with_logits(
                logits=masked_logits, labels=masked_labels, name="softmax_loss_tf"
            )

            cost_tf = tf.reduce_mean(softmax_loss_tf, name="cost_tf")

        optimizer_tf = tf.train.AdamOptimizer(
            learning_rate=float(self.learning_rate)
        ).minimize(cost_tf)

        return optimizer_tf, cost_tf
//  [python]
def _define_optimizer_and_cost(self):
        """This function defines the optimizer and cost function of the LSTM model

        Returns:
            AdamOptimizer, Tensor: The optimizer function to reduce loss and the loss values
        """
        if self.use_crf_layer:
            flattened_labels = tf.cast(tf.argmax(self.label_tf, axis=2), tf.int32)
            log_likelihood, _ = tf.contrib.crf.crf_log_likelihood(
                self.lstm_output_tf, flattened_labels, self.batch_sequence_lengths_tf
            )
            cost_tf = tf.reduce_mean(-log_likelihood, name="cost_tf")
        else:
            masked_logits = tf.boolean_mask(
                tf.reshape(self.lstm_output_tf, [-1, self.output_dimension]),
                self.batch_sequence_mask_tf,
            )

            masked_labels = tf.boolean_mask(
                tf.reshape(self.label_tf, [-1, self.output_dimension]),
                self.batch_sequence_mask_tf,
            )

            softmax_loss_tf = tf.nn.softmax_cross_entropy_with_logits(
                logits=masked_logits, labels=masked_labels, name="softmax_loss_tf"
            )

            cost_tf = tf.reduce_mean(softmax_loss_tf, name="cost_tf")

        optimizer_tf = tf.train.AdamOptimizer(
            learning_rate=float(self.learning_rate)
        ).minimize(cost_tf)

        return optimizer_tf, cost_tf
//  [python]
def read_single_example_and_decode(filename_queue):

    tfrecord_options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)

    reader = tf.TFRecordReader(options=tfrecord_options)
    # reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)

    features = tf.parse_single_example(
        serialized=serialized_example,
        features={
            'img_name': tf.FixedLenFeature([], tf.string),
            'img_height': tf.FixedLenFeature([], tf.int64),
            'img_width': tf.FixedLenFeature([], tf.int64),
            'img': tf.FixedLenFeature([], tf.string),
            'gtboxes_and_label': tf.FixedLenFeature([], tf.string),
            'num_objects': tf.FixedLenFeature([], tf.int64)
        }
    )
    img_name = features['img_name']
    img_height = tf.to_int32(features['img_height'])
    img_width = tf.to_int32(features['img_width'])
    img = tf.decode_raw(features['img'], tf.uint8)

    img = tf.reshape(img, shape=[img_height, img_width, 3])
    gtboxes_and_label = tf.decode_raw(features['gtboxes_and_label'], tf.int32)
    gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 9])
    num_objects = tf.to_int32(features['num_objects'])
    return img_name, img, gtboxes_and_label, num_objects
//  [python]
def read_single_example_and_decode(filename_queue):

    tfrecord_options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)

    reader = tf.TFRecordReader(options=tfrecord_options)
    # reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)

    features = tf.parse_single_example(
        serialized=serialized_example,
        features={
            'img_name': tf.FixedLenFeature([], tf.string),
            'img_height': tf.FixedLenFeature([], tf.int64),
            'img_width': tf.FixedLenFeature([], tf.int64),
            'img': tf.FixedLenFeature([], tf.string),
            'gtboxes_and_label': tf.FixedLenFeature([], tf.string),
            'num_objects': tf.FixedLenFeature([], tf.int64)
        }
    )
    img_name = features['img_name']
    img_height = tf.cast(features['img_height'], tf.int32)
    img_width = tf.cast(features['img_width'], tf.int32)
    img = tf.decode_raw(features['img'], tf.uint8)

    img = tf.reshape(img, shape=[img_height, img_width, 3])
    gtboxes_and_label = tf.decode_raw(features['gtboxes_and_label'], tf.int32)
    gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 9])
    num_objects = tf.cast(features['num_objects'], tf.int32)
    return img_name, img, gtboxes_and_label, num_objects
//  [python]
def normalize(self, I, Io=240, alpha=1, beta=0.15, stains=True):
        ''' Normalize staining appearence of H&E stained images

        Example use:
            see test.py

        Input:
            I: RGB input image: tensor of shape [C, H, W] and type uint8
            Io: (optional) transmitted light intensity
            alpha: percentile
            beta: transparency threshold
            stains: if true, return also H & E components

        Output:
            Inorm: normalized image
            H: hematoxylin image
            E: eosin image

        Reference:
            A method for normalizing histology slides for quantitative analysis. M.
            Macenko et al., ISBI 2009
        '''
        c, h, w = I.shape

        HE, C, maxC = self.__compute_matrices(I, Io, alpha, beta)

        # normalize stain concentrations
        C *= tf.expand_dims((self.maxCRef / maxC), axis=-1)

        # recreate the image using reference mixing matrix
        Inorm = Io * tf.math.exp(-tf.linalg.matmul(self.HERef, C))
        Inorm = tf.clip_by_value(Inorm, 0, 255)
        Inorm = t.cast(tf.reshape(tf.transpose(Inorm), shape=(h, w, c)), t.int32)

        H, E = None, None

        if stains:
            H = tf.math.multiply(Io, tf.math.exp(tf.linalg.matmul(-tf.expand_dims(self.HERef[:, 0], axis=-1), tf.expand_dims(C[0, :], axis=0))))
            H = tf.clip_by_value(H, 0, 255)
            H = t.cast(tf.reshape(tf.transpose(H), shape=(h, w, c)), t.int32)

            E = tf.math.multiply(Io, tf.math.exp(tf.linalg.matmul(-tf.expand_dims(self.HERef[:, 1], axis=-1), tf.expand_dims(C[1, :], axis=0))))
            E = tf.clip_by_value(E, 0, 255)
            E = t.cast(tf.reshape(tf.transpose(E), shape=(h, w, c)), t.int32)

        return Inorm, H, E
//  [python]
def normalize(self, I, Io=240, alpha=1, beta=0.15, stains=True):
        ''' Normalize staining appearence of H&E stained images

        Example use:
            see test.py

        Input:
            I: RGB input image: tensor of shape [C, H, W] and type uint8
            Io: (optional) transmitted light intensity
            alpha: percentile
            beta: transparency threshold
            stains: if true, return also H & E components

        Output:
            Inorm: normalized image
            H: hematoxylin image
            E: eosin image

        Reference:
            A method for normalizing histology slides for quantitative analysis. M.
            Macenko et al., ISBI 2009
        '''
        c, h, w = I.shape

        HE, C, maxC = self.__compute_matrices(I, Io, alpha, beta)

        # normalize stain concentrations
        C *= tf.expand_dims((self.maxCRef / maxC), axis=-1)

        # recreate the image using reference mixing matrix
        Inorm = Io * tf.math.exp(-tf.linalg.matmul(self.HERef, C))
        Inorm = tf.clip_by_value(Inorm, 0, 255)
        Inorm = tf.cast(tf.reshape(tf.transpose(Inorm), shape=(h, w, c)), tf.int32)

        H, E = None, None

        if stains:
            H = tf.math.multiply(Io, tf.math.exp(tf.linalg.matmul(-tf.expand_dims(self.HERef[:, 0], axis=-1), tf.expand_dims(C[0, :], axis=0))))
            H = tf.clip_by_value(H, 0, 255)
            H = tf.cast(tf.reshape(tf.transpose(H), shape=(h, w, c)), tf.int32)

            E = tf.math.multiply(Io, tf.math.exp(tf.linalg.matmul(-tf.expand_dims(self.HERef[:, 1], axis=-1), tf.expand_dims(C[1, :], axis=0))))
            E = tf.clip_by_value(E, 0, 255)
            E = tf.cast(tf.reshape(tf.transpose(E), shape=(h, w, c)), tf.int32)

        return Inorm, H, E
//  [python]
def compute_length_after_conv(max_time_steps, ctc_time_steps, input_length):
  """Computes the time_steps/ctc_input_length after convolution.

  Suppose that the original feature contains two parts:
  1) Real spectrogram signals, spanning input_length steps.
  2) Padded part with all 0s.
  The total length of those two parts is denoted as max_time_steps, which is
  the padded length of the current batch. After convolution layers, the time
  steps of a spectrogram feature will be decreased. As we know the percentage
  of its original length within the entire length, we can compute the time steps
  for the signal after conv as follows (using ctc_input_length to denote):
  ctc_input_length = (input_length / max_time_steps) * output_length_of_conv.
  This length is then fed into ctc loss function to compute loss.

  Args:
    max_time_steps: max_time_steps for the batch, after padding.
    ctc_time_steps: number of timesteps after convolution.
    input_length: actual length of the original spectrogram, without padding.

  Returns:
    the ctc_input_length after convolution layer.
  """
  ctc_input_length = t.cast(t.multiply(
      input_length, ctc_time_steps), dtype=t.float32)
  return t.cast(t.floordiv(
      ctc_input_length, t.cast(max_time_steps, dtype=t.float32)), dtype=t.int32)
//  [python]
def compute_length_after_conv(max_time_steps, ctc_time_steps, input_length):
  """Computes the time_steps/ctc_input_length after convolution.

  Suppose that the original feature contains two parts:
  1) Real spectrogram signals, spanning input_length steps.
  2) Padded part with all 0s.
  The total length of those two parts is denoted as max_time_steps, which is
  the padded length of the current batch. After convolution layers, the time
  steps of a spectrogram feature will be decreased. As we know the percentage
  of its original length within the entire length, we can compute the time steps
  for the signal after conv as follows (using ctc_input_length to denote):
  ctc_input_length = (input_length / max_time_steps) * output_length_of_conv.
  This length is then fed into ctc loss function to compute loss.

  Args:
    max_time_steps: max_time_steps for the batch, after padding.
    ctc_time_steps: number of timesteps after convolution.
    input_length: actual length of the original spectrogram, without padding.

  Returns:
    the ctc_input_length after convolution layer.
  """
  ctc_input_length = tf.cast(tf.multiply(
      input_length, ctc_time_steps), dtype=tf.float32)
  return tf.cast(tf.math.floordiv(
      ctc_input_length, tf.cast(max_time_steps, dtype=tf.float32)), dtype=tf.int32)
//  [python]
def test_with_var_read_x64(self, with_jit=True):
    outer_var_array = np.array([3., 4.], dtype=np.float64)
    outer_var = tf.Variable(outer_var_array)

    def fun_tf(x):
      return x * tf.cast(outer_var, x.dtype) + 1.

    x = np.array([2., 5.,], dtype=np.float32)
    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
    self.assertAllClose(x * outer_var_array + 1., res, check_dtypes=False)
//  [python]
def test_with_tensor_capture_x64(self, with_jit=True):
    outer_tensor = tf.constant(3., dtype=np.float64)

    def fun_tf(x):
      return x * t.cast(outer_tensor * 3.14, tf.float32) + 1.

    x = np.float32(2.)
    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
    self.assertAllClose(x * 3. * 3.14 + 1., res, check_dtypes=False)
//  [python]
def test_with_tensor_capture_x64(self, with_jit=True):
    outer_tensor = tf.constant(3., dtype=np.float64)

    def fun_tf(x):
      return x * tf.cast(outer_tensor * 3.14, tf.float32) + 1.

    x = np.float32(2.)
    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
    self.assertAllClose(x * 3. * 3.14 + 1., res, check_dtypes=False)
//  [python]
def pad(
        self,
        sequence: tf.RaggedTensor,
    ) -> TF_TOKENIZER_INPUT:

        r"""
        Add padding of a sequence and supports different strategies for doing so.

        ::note: passing in padding will have no effect when loading a saved tokenizer.

        `strategy` supports different options:
            - `max_length`: Pad all batches to tokenizer `max_seq_length`.
            - `longest`: Pad all to same length as the longest sequence in the batch.
            - `none`: Do not pad at all.
        """
        padding = self.padding
        max_length = self.max_length
        _input_ids, _attention_mask = TF_TENSORS, TF_TENSORS  # type: ignore

        if padding == PaddingStrategies.MAX_LENGTH:
            _input_ids, _attention_mask = text.pad_model_inputs(
                input=sequence, max_seq_length=max_length
            )
        elif padding == PaddingStrategies.LONGEST:
            seq_len_per_batch: tf.Tensor = sequence.nested_row_lengths()[0]
            longest_seq_in_batch = tf.reduce_max(seq_len_per_batch)
            _input_ids, _attention_mask = text.pad_model_inputs(
                input=sequence, max_seq_length=longest_seq_in_batch
            )
        elif padding == PaddingStrategies.NONE:
            _input_ids: tf.RaggedTensor = sequence
            _attention_mask: tf.RaggedTensor = tf.ones_like(sequence)
        else:
            raise NotValidPaddingStrategy(padding=padding)

        # Since Huggingface tensors are of same type
        input_ids = t.cast(_input_ids, dtype=tf.int64)
        attention_mask = t.cast(_attention_mask, dtype=tf.int64)
        return input_ids, attention_mask
//  [python]
def pad(
        self,
        sequence: tf.RaggedTensor,
    ) -> TF_TOKENIZER_INPUT:

        r"""
        Add padding of a sequence and supports different strategies for doing so.

        ::note: passing in padding will have no effect when loading a saved tokenizer.

        `strategy` supports different options:
            - `max_length`: Pad all batches to tokenizer `max_seq_length`.
            - `longest`: Pad all to same length as the longest sequence in the batch.
            - `none`: Do not pad at all.
        """
        padding = self.padding
        max_length = self.max_length
        _input_ids, _attention_mask = TF_TENSORS, TF_TENSORS  # type: ignore

        if padding == PaddingStrategies.MAX_LENGTH:
            _input_ids, _attention_mask = text.pad_model_inputs(
                input=sequence, max_seq_length=max_length
            )
        elif padding == PaddingStrategies.LONGEST:
            seq_len_per_batch: tf.Tensor = sequence.nested_row_lengths()[0]
            longest_seq_in_batch = tf.reduce_max(seq_len_per_batch)
            _input_ids, _attention_mask = text.pad_model_inputs(
                input=sequence, max_seq_length=longest_seq_in_batch
            )
        elif padding == PaddingStrategies.NONE:
            _input_ids: tf.RaggedTensor = sequence
            _attention_mask: tf.RaggedTensor = tf.ones_like(sequence)
        else:
            raise NotValidPaddingStrategy(padding=padding)

        # Since Huggingface tensors are of same type
        input_ids = tf.cast(_input_ids, dtype=tf.int64)
        attention_mask = tf.cast(_attention_mask, dtype=tf.int64)
        return input_ids, attention_mask
//  [python]
def iou(self, boxes, clusters):
    """Computes iou."""
    # broadcast the clusters to the same shape as the boxes
    n = tf.shape(boxes)[0]
    k = tf.shape(clusters)[0]
    boxes = tf.repeat(boxes, k, axis=0)
    boxes = tf.reshape(boxes, (n, k, -1))
    boxes = tf.to_float(boxes)

    clusters = tf.tile(clusters, [n, 1])
    clusters = tf.reshape(clusters, (n, k, -1))
    clusters = tf.to_float(clusters)

    # compute the IOU
    return _iou(boxes, clusters)
//  [python]
def iou(self, boxes, clusters):
    """Computes iou."""
    # broadcast the clusters to the same shape as the boxes
    n = tf.shape(boxes)[0]
    k = tf.shape(clusters)[0]
    boxes = tf.repeat(boxes, k, axis=0)
    boxes = tf.reshape(boxes, (n, k, -1))
    boxes = tf.cast(boxes, tf.float32)

    clusters = tf.tile(clusters, [n, 1])
    clusters = tf.reshape(clusters, (n, k, -1))
    clusters = tf.cast(clusters, tf.float32)

    # compute the IOU
    return _iou(boxes, clusters)
//  [python]
def get_loss(self, logits, labels, pre_cal_weights, mask_boolen=None):
        # calculate the weighted cross entropy according to the inverse frequency
        class_weights = tf.convert_to_tensor(pre_cal_weights, dtype=tf.float32)
        one_hot_labels = tf.one_hot(labels, depth=self.config.num_classes)
        weights = tf.reduce_sum(class_weights * one_hot_labels, axis=1)
        unweighted_losses = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=one_hot_labels)
        weighted_losses = unweighted_losses * weights
        
        if mask_boolen is not None:
            weighted_losses = tf.reduce_sum(tf.to_float(mask_boolen) * weighted_losses) / tf.reduce_sum(tf.to_float(mask_boolen))
            output_loss = weighted_losses + \
                                tf.to_float(tf.greater_equal(self.training_epoch, self.rampup)) * (
                                            self.loss_siamese + self.loss_inexact + self.loss_smooth)
        else:
            output_loss = tf.reduce_mean(weighted_losses)
        return output_loss
//  [python]
def get_loss(self, logits, labels, pre_cal_weights, mask_boolen=None):
#         pdb.set_trace()
        # calculate the weighted cross entropy according to the inverse frequency
        class_weights = tf.convert_to_tensor(pre_cal_weights, dtype=tf.float32)
        one_hot_labels = tf.one_hot(labels, depth=self.config.num_classes)
        weights = tf.reduce_sum(class_weights * one_hot_labels, axis=1)
        unweighted_losses = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=one_hot_labels)
        weighted_losses = unweighted_losses * weights
        
        if mask_boolen is not None:
            weighted_losses = tf.reduce_sum(tf.cast(mask_boolen, tf.float32) * weighted_losses) / tf.reduce_sum(tf.cast(mask_boolen, tf.float32))
            output_loss = weighted_losses + \
                                tf.cast(tf.greater_equal(self.training_epoch, self.rampup), dtype=tf.float32) * (
                                            self.loss_siamese + self.loss_inexact + self.loss_smooth)
        else:
            output_loss = tf.reduce_mean(weighted_losses)
        return output_loss
//  [python]
def deterministic_preprocess(serialized_example):
  # First parse the example protobuf.
  image_buffer, label, bbox = parse_example_proto(serialized_example)
  
  # Then decode and perform mean substraction.
  image = decode(image_buffer, NUM_CHANNELS)
  image = t.cast(image, t.float32)
  image = _mean_image_subtraction(image, CHANNEL_MEANS, 3)
  
  # Subtract one so that labels are in [0, 1000), and cast to float32 for
  # Keras model.
  label = t.cast(
      t.cast(t.reshape(label, shape=[1]), dtype=t.int32) - 1,
      dtype=t.float32)
  return image, label, bbox
//  [python]
def deterministic_preprocess(serialized_example):
  # First parse the example protobuf.
  image_buffer, label, bbox = parse_example_proto(serialized_example)
  
  # Then decode and perform mean substraction.
  image = decode(image_buffer, NUM_CHANNELS)
  image = tf.cast(image, tf.float32)
  image = _mean_image_subtraction(image, CHANNEL_MEANS, 3)
  
  # Subtract one so that labels are in [0, 1000), and cast to float32 for
  # Keras model.
  label = tf.cast(
      tf.cast(tf.reshape(label, shape=[1]), dtype=tf.int32) - 1,
      dtype=tf.float32)
  return image, label, bbox
//  [python]
def log_differences(self, n_vars, tensor_array, steps, name_scope):
        with tf.name_scope(name_scope):
            for i in range(1, self.train_rounds, 1):
                if i < steps:
                    new = tensor_array.read(i)
                    old = tensor_array.read(i - 1)
                    new = tf.round(tf.sigmoid(new))
                    old = tf.round(tf.sigmoid(old))
                    matching = tf.equal(new, old)
                    matching = tf.reduce_sum(matching, axis=0) / tf.cast(n_vars, tf.float32)
                    same = tf.reduce_mean(matching)

                    tf.summary.scalar(f"step_{i}_vs_{i - 1}", same)
//  [python]
def log_differences(self, n_vars, tensor_array, steps, name_scope):
        with tf.name_scope(name_scope):
            for i in range(1, self.train_rounds, 1):
                if i < steps:
                    new = tensor_array.read(i)
                    old = tensor_array.read(i - 1)
                    new = tf.round(tf.sigmoid(new))
                    old = tf.round(tf.sigmoid(old))
                    matching = tf.cast(tf.equal(new, old), tf.float32)
                    matching = tf.reduce_sum(matching, axis=0) / tf.cast(n_vars, tf.float32)
                    same = tf.reduce_mean(matching)

                    tf.summary.scalar(f"step_{i}_vs_{i - 1}", same)
//  [python]
def WeakSupLoss(self):
        '''
        Define additional losses for weakly supervised segmentation
        Inputs:
            self.logits = tf.reshape(self.logits, [-1, config.num_classes])
            self.labels = tf.reshape(self.labels, [-1])
            self.mask_boolen = tf.reshape(self.mask_boolen, [-1])
        Returns:

        '''
        
        ## MIL Branch
        one_hot_labels = tf.one_hot(self.batch_labels, depth=self.config.num_classes, axis=-1)
#         one_hot_labels = tf.reshape(one_hot_labels, [self.inputs.shape[0],self.inputs.shape[1],self.config.num_classes])
        L_gt = tf.reduce_max(one_hot_labels[0::2], axis=1)  # inexact labels B*13
        L_pred = tf.reduce_max(self.batch_logits[0::2], axis=1)
        loss_ineaxct = tf.nn.sigmoid_cross_entropy_with_logits(labels=L_gt, logits=L_pred)  # B*K, independent loss for each class
        self.loss_inexact = 1* tf.reduce_mean(loss_ineaxct)
        
        ## Siamese Branch
        self.prob_logits = tf.nn.softmax(self.batch_logits, axis=-1)
        self.loss_siamese = 1* tf.reduce_mean(tf.reduce_sum((self.prob_logits[0::2] - self.prob_logits[1::2]) ** 2, axis=-1))
        
#         pdb.set_trace()
        ## Smooth Branch
#         self.loss_smooth = 10* SmoothConstraint.Loss_SpatialColorSmooth_add_SelfContain(self.prob_logits[0::2], self.inputs['features'][0::2, :, 0:3], M=self.batch_masks[0::2], Y=self.batch_labels[0::2], gamma=1e-1)
        self.loss_smooth = 10 * SmoothConstraint.Loss_SpatialColorSmooth_add_SelfContain_mustLink(self.prob_logits[0::2], self.inputs['features'][0::2, :, 0:6], M=self.batch_masks[0::2], Y=self.batch_labels[0::2], gamma=1e-1)