class_docstr: Merge labels of instances with similar bounding boxes.
  
  This is useful when data contains multiple non-disjoint labels per instance
  (e.g. in federated datasets, when bounding boxes for specified labels are
  annotated independently of each other).
  
  Box similarity is assessed using the IoU (Intersetion over Union) metric. See
  `transforms.box_iou` for details.
  
  Attributes:
    iou_threshold: Box IoU threshold used to determine whether two instances
      should be merged.
    eps: Small float number used for numerical stability when computing IoU.
    label_feature_keys: Sequence of instance label modalities that should be
      merged. Note that these should be multi-label modalities.
class_docstr_tok: ["Merge","labels","of","instances","with","similar","bounding","boxes",".","This","is","useful","when","data","contains","multiple","non-disjoint","labels","per","instance","(","e.g",".","in","federated","datasets",",","when","bounding","boxes","for","specified","labels","are","annotated","independently","of","each","other",")",".","Box","similarity","is","assessed","using","the","IoU","(","Intersetion","over","Union",")","metric",".","See","`","transforms.box_iou","`","for","details",".","Attributes",":","iou_threshold",":","Box","IoU","threshold","used","to","determine","whether","two","instances","should","be","merged",".","eps",":","Small","float","number","used","for","numerical","stability","when","computing","IoU",".","label_feature_keys",":","Sequence","of","instance","label","modalities","that","should","be","merged",".","Note","that","these","should","be","multi-label","modalities","."]
docstr: Iteratively merge instances with bbox IoU above a preset threshold.
docstr_tok: ["Iteratively","merge","instances","with","bbox","IoU","above","a","preset","threshold","."]
code: def __call__(self, features: Features) -> Features:
      """Iteratively merge instances with bbox IoU above a preset threshold."""
  
      def _compute_masked_iou(boxes):
          iou, _ = transforms.box_iou(boxes, boxes, eps=self.eps)
          # Mask comparison to self.
          iou = tf.where(
              tf.eye(tf.shape(iou)[0], dtype=tf.bool), 0.0, iou
          )
          return iou
  
      def _has_boxes_to_merge(boxes, labels, iou, rng):
          del boxes, labels, rng
          max_iou = tf.reduce_max(iou)
          return tf.greater_equal(max_iou, self.iou_threshold)
  
      def _merge_boxes(boxes, labels, iou, rng):
          ind = tf.unravel_index(  # Indices [i, j] of the two most similar boxes.
              tf.cast(tf.argmax(tf.reshape(iou, (-1,))), tf.int32),
              tf.shape(iou),
          )
  
          # Merge labels of the two boxes; do not pay attention to possible
          # label collisions.
          labels_merged = tf.nest.map_structure(
              lambda lab: tf.gather(lab, ind).flat_values, labels
          )
  
          # Pick one box randomly.
          rngs = tf.random.experimental.stateless_split(rng)
          rng, new_rng = rngs[0], rngs[1]
          boxes_merged = tf.gather(boxes, ind)
          choice = tf.greater(tf.random.stateless_uniform([], rng), 0.5)
          boxes_merged = tf.where(
              choice, boxes_merged[0], boxes_merged[1]
          )
  
          # Updated labels and boxes based on the merge. Tensors will be merged
          # using a gather. Prepare gather indices below.
          num_instances = tf.shape(boxes)[0]
          i, j = tf.reduce_min(ind), tf.reduce_max(ind)
          ind_all = tf.range(num_instances)
          ind_no_j = ind_all[ind_all != j]
          ind = tf.where(ind_no_j == i, num_instances, ind_no_j)
  
          def _merge_tensors(tensor, merged_value):
              tensor = tf.concat(
                  [tensor, tf.expand_dims(merged_value, 0)], axis=0
              )
              return tf.gather(tensor, ind)
  
          boxes = _merge_tensors(boxes, boxes_merged)
          labels = tf.nest.map_structure(
              _merge_tensors, labels, labels_merged
          )
  
          # Finally, update iou for the next iteration.
          iou = _compute_masked_iou(boxes)
  
          return (boxes, labels, iou, new_rng)
  
      if SEED_KEY not in features:
          raise ValueError("Merged box choice requires a random seed.")
  
      rng = features[preprocess_spec.SEED_KEY]
  
      boxes = features[modalities.BOXES]
      max_instances = tf.shape(boxes)[0]
      not_padding = tf.reduce_any(boxes != -1, axis=-1)
      boxes = boxes[not_padding]
  
      # Prepare (ragged) label tensors. This simplifies label merging logic.
      def to_ragged(x: tf.Tensor) -> tf.RaggedTensor:
          x = x[not_padding]
          lengths = tf.reduce_sum(
              tf.cast(
                  tf.not_equal(
                      x, transforms.get_padding_value(x.dtype)
                  ),
                  tf.int32,
              ),
              axis=-1,
          )
          return tf.RaggedTensor.from_tensor(x, lengths)
  
      orig_labels = {}
      for name in self.label_feature_keys:
          if name in features:
              orig_labels[name] = features[name]
      labels = tf.nest.map_structure(to_ragged, orig_labels)
  
      state = (boxes, labels, _compute_masked_iou(boxes), rng)
      boxes, labels, _, rng = tf.while_loop(
          _has_boxes_to_merge,
          _merge_boxes,
          state,
          parallel_iterations=1,
          back_prop=False,
      )
  
      features_new = dict(features)
  
      features_new[preprocess_spec.SEED_KEY] = rng
  
      # Pad boxes to original shape.
      features_new[modalities.BOXES] = tf.pad(
          boxes,
          [(0, max_instances - tf.shape(boxes)[0]), (0, 0)],
          constant_values=-1,
      )
  
      # Convert ragged to normal tensor.
      def to_tensor(labels, orig_labels):
          return labels.to_tensor(
              default_value=transforms.get_padding_value(
                  orig_labels.dtype
              ),
              shape=tf.shape(orig_labels),
          )
  
      labels = tf.nest.map_structure(to_tensor, labels, orig_labels)
      features_new.update(labels)
      return features_new
  
imports: import abc
  import dataclasses
  from typing import Optional, Sequence, Tuple, Union
  from clu import preprocess_spec
  from scenic.projects.owl_vit.preprocessing import modalities
  from scenic.projects.owl_vit.preprocessing import transforms
  import tensorflow
label_desc: Use image node_id 1 for calling the MergeOverlappingInstances obj's underlying member method code with example usage: obj.__call__(features) and returns: features_new, iou, tf, boxes,  labels,  iou,  new_rng, tf, labels, tf