Add a single example to the evaluator. Args: prediction: Model prediction dictionary with keys 'pred_img_ids', 'pred_probs' in shap

Add gaussian noise to the given image. Args: image: A single image,must be in 0-1 range. noise_magnitude: Amount of noise to be added to th

Applies the rotary embedding to the input.

Apply label smoothing to the one-hot targets. Applies label smoothing such that the on-values are transformed from 1.0 to `1.0 - label_smoothing

Maps `ds` using the preprocess_fn and a deterministic RNG per example. Args: ds: Dataset containing Python dictionary with the features. The 'r

Performs SpecAugment on the inputs. SpecAugment is a data augmentation technique from arXiv:1904.08779, that combines three transformations: -

Performs 2D bilinear interpolation. It is assumed that the center of the top-left pixel in `im` has coordinate (0.5, 0.5). If you want a differen

Compute box IoU. Boxes in format [l, t, w, h]. Args: boxes1: array in shape n x 4 boxes2: array in shape m x 4 Returns: iou: array in

Bounding box losses: L1 regression loss and GIoU loss.. Args: outputs: dict; Model predictions. For the purpose of this loss, outputs

Dataset builder that takes care of strategy, batching and shuffling. Args: dataset_fn: function; A function that loads the dataset. batch_s

Dataset builder that takes care of strategy, batching and shuffling. Args: dataset_fn: function; A function that loads the dataset. batch_s

Dataset builder that takes care of strategy, batching and shuffling. Args: dataset_fn: function; A function that loads the dataset. batch_s

Dataset builder that takes care of strategy, batching and shuffling. Args: dataset_fn: function; A function that loads the dataset. batch_s

Dataset builder that takes care of strategy, batching and shuffling. Args: dataset_fn: function; A function that loads the dataset. batch_s

Dataset builder that takes care of strategy, batching and shuffling. Args: dataset_fn: function; A function that loads the dataset. batch_s

Dataset builder that takes care of strategy, batching and shuffling. Args: dataset_fn: function; A function that loads the dataset. batch_s

Randomly select a non-zero element from a 2D mask. Args: key: random key array: 2D array of non-zero elements Returns: point: [x,y]

Computes cost matrices for a batch of predictions. Relevant code: https://github.com/facebookresearch/detr/blob/647917626d5017e63c1217b99537deb2d

Compute inner product between videos and text embeddings.

Copy matched parameters from a restored one.

Creates a moment retrieval dataset iterator.

Creates a temporal localization dataset iterator.

Applies Transformer decoder-branch on encoded-input and target. Args: encoded: encoded input data from encoder. targets: target data.

Compute boundary map for each patch, given distance functions. The width of the boundary is determined by self.opts.delta. Args: dists

Embeds text queries. Args: queries: Tuple of query strings. Returns: Numpy arrays containing query embeddings.

Encode text. Args: encoder_input_tokens: (batch_size, max_cap_len) for ASR or (batch_size, num_caps_per_image, max_cap_len) for p

Gathers the beam slices indexed by beam_indices into new beam array. Args: nested: pytree of arrays or scalars (the latter ignored). beam_i

Gathers the beam slices indexed by beam_indices into new beam array. Args: nested: pytree of arrays or scalars (the latter ignored). beam_i

Gathers the beam slices indexed by beam_indices into new beam array. Args: nested: pytree of arrays or scalars (the latter ignored). beam_i

r"""Computes general features of kernel's linearization. Args: data: input data tensor of the shape [B..., L, H, D], where: B - batch dim

Returns the configuration for COCO detection using DeformableDETR.

Returns the RefCOCO/RefCOCOg/RefCOCO+/ eval source.

Returns the RefCOCO/RefCOCOg/RefCOCO+/ eval source.

Returns the RefCOCO/RefCOCOg/RefCOCO+/ train source.

Returns the RefCOCO/RefCOCOg/RefCOCO+/ train source.

Returns the RefCOCO/RefCOCOg/RefCOCO+/ train source.

Returns the RefCOCO/RefCOCOg/RefCOCO+/ train source.

Stochastic drop-layer probability. Args: layer: Layer index. num_layers: Total number of layers. stochastic_droplayer_rate: Proba

Generates a fake `batch`. Args: array_size: size of the label and output array. Returns: `batch`: Dictionary of None inputs and fake gro

Generates a fake `batch`. Returns: `batch`: Dictionary of None inputs and fake ground truth targets. outputs_noaux.pop('aux_outputs')

Calculates the loss for global supervision.

Return the height, width of image as float32 tf.Tensors.

Given a (dict of) Dataset object(s), returns iterators and metadata. Args: ds: A tf.data.Dataset instance or a dictionary of TFDS instances.

Given a (dict of) Dataset object(s), returns iterators and metadata. Different from the original flexio, this function uses a custom get_number_o

Given a (dict of) Dataset object(s), returns iterators and metadata. Args: ds: A tf.data.Dataset instance or a dictionary of TFDS instances.

Compute the objective of the model (see Equation 8 of the paper). Args: dists: Array of shape [N, 2, R, R, H', W'] with samples of the two

Calculates the loss function.

Returns a callable metric function for the model. Args: split: The split for which we calculate the metrics. It should be one of th

Forward image backbone and aggregate video features. Args: images: (total_batch_size, height, width, 3). Note for videos, the total

Initialize the positional embeddings.

Inspects whether the params are consistent with the expected keys. Based on https://github.com/google-research/big_vision/blob/main/big_vision/mo

Helper function to assert that tensor `t` does not have `nan`, `inf`.

Helper function to assert that tensor `t` does not have `nan`, `inf`.

Helper function to assert that tensor `t` does not have `nan`, `inf`.

Helper function to assert that tensor `t` does not have `nan`, `inf`.

Helper function to assert that `loss` is of shape [] and `is_valid`.

Helper function to assert that `loss` is of shape [] and `is_valid`.

Loads OVEN or entity-based pretraining dataset using TensorFlow Datasets. Args: batch_size: int; The batch size returned by the data pipeline.

Returns the loss. Args: logits: Output of model in shape [batch, length, num_classes]. batch: Batch of data that has 'label' and opti

Make 2D boxes out of 1D segments. We reuse tf.image.non_max_suppression_with_scores() for non-maximal suppression, which takes 2D boxes. Args:

Create boolean mask by resizing from given mask or set all True.

Returns the minimum of two integers, ignoring None values.

Plots a single bounding box into axes.

Aggregate text embedding for a UI element.

Proprocess images. Normalize pixels for non-padded pixels.

Feeds the inputs to the model and returns their representations. Args: train_state: TrainState, the state of training including the current

Creates a split from the ShapeNet dataset using TensorFlow Datasets. For the training set, we drop the last partial batch. This is fine to do bec

Gathers `x` using the indices in `idx`. `output[i] = x[i, idx[i]]` . This simple gather operation assumes that the first dimension is the batch d

Transposes a list of dicts to dict of lists. For example, given [{'a':1,'b':2}, {'a':3,'b':4}], the output is: {'a': ([1, 3]), 'b': ([2, 4]

Write to the given writer and training step. After writing, the state gets reset. Args: writer: The summary will be written with this

Writes predictions to file in JSON format. Args: path: Path to write the prediction annotation JSON file. is_groundtruth: bool; if th

Initializes detection model. Args: config: Hyper-parameter dictionary. dataset_meta_data: Dataset meta data specifies `target_is_oneh

Adds positional embedding.