Dataset function.

Dataset function.

Dataset function.

Dataset function.

Dataset function.

Make a random crop of `image_size`.

Decodes audio spectrogram. Args: spectrogram: input mel spectrogram inflate: if True, adds a channel dimension zero_centering: if True,

Computes the denominator.

Utility function for computing the shape resulting from a dilation. Args: shape: Shapes (input or kernel i.e. lhr or rhs) to which the dilation

Downloads the appropriate annotations file.

Downloads a file from a URL to a path.

Subsamples video tokens.

Drops degenerate (e.g. cropped out) boxes.

Subsample token x according to the keeped index in fg_inds. Args: x: batch x num_tokens x hidden_dim fg_inds: batch x num_fg_tokens

Embeds image and returns image features, class embeddings, and boxes.

Embed points. Args: points: (num_prompts, num_points, 2). In absolute coordinates. labels: (num_prompts, num_points) pad: bool

Erase an area.

Evaluates a trained linear probe. Args: representation_fn: A function that given a batch, returns representations that are passed t

A context manager for each state of eval.

Runs a single step of moment retrieval evaluation.

Performs a single evaluation pass over the dataset for the given step.

Extracts encoder output.

Extracts the image shape from the first example.

Computes (x,y) linear regression accuracy on (x_test, y_test). Args: x: An [n_examples, n_classes] matrix of feature representations. This will

Fills blank area for video.

Rotates an image four times, with 90 degrees between each rotation.

From grid.

Gathers the top instances based on IoU with ground-truth boxes.

Generate attention mask.

Generate attention mask.

Generate attention mask.

Returns fake inputs for visualization tests.

Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token dimension for the original embeddings. Args

Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token dimension for the original embeddings. Args

generate FCOS style regression targets. Args: grids: array in shape m x 2: all output pixel coordinates. boxes: array in shape n x 4:

pad arrays which will be used to generate heatmaps. Args: grids: array in shape m x 2: all output pixel coordinates. boxes: array in

Generates a noise that is a chunky version of the input noise.

Lazy-loads given dataset.

Lazy-loads given dataset.

Lazy-loads given dataset.

Lazy-loads given dataset.

Generates a gaussian noise.

Obtain a size x size meshgrid inside the given box. Args: box: XYXY-format boxes of shape (T, 4). size: Resolution of the grid. sparse:

Computes the 1d index & their weights to be used in interpolation.

Looks up the axis storing object instances for a given feature.

Creates datasets from builders, applies preprocessing, and merges them. Args: builders: List of TFDS dataset builders. splits: List contain

define positive/ negative pixels for regression and heatmap. Args: centers_discret: array in shape m x n x 2: snap each of the n centers to

Constructs the preprocess_fn that should be applied after mosaicing. Args: spec: Config string specifying the preprocessing. Returns: Pr

Constructs the preprocess_fn that should be applied before mosaicing. Args: builder: TFDS dataset builder. decoder_kwargs: Decoder kwargs.

generate heatmaps and normalized regression targets. Args: area: array in shape n: area if each ground truth box. centers_expanded: a

Creates a Grain-backed dataset from builder and applies preprocessing. Args: builder: TFDS dataset builder. start_index: Index dataset (Gra

Creates a Grain-backed dataset from builder and applies preprocessing. Args: builder: TFDS dataset builder. start_index: Index dataset (Gra

Creates dataset from builder and applies pre-mosaic preprocessing. Args: builder: TFDS dataset builder. split: Train/test/validation split.

Creates dataset from builder and applies preprocessing. Args: builder: TFDS dataset builder. split: Train/test/validation split. batch_

Creates dataset from builder and applies preprocessing. Args: builder: TFDS dataset builder. split: Train/test/validation split. batch_

Creates dataset using DMVR and applies preprocessing. Args: tfrecords: Path to tfrecords. context_features: Dictionary of context features

Creates dataset using DMVR and applies preprocessing. Args: tfrecords: Path to tfrecords. context_features: Dictionary of context features

Gets a unique set of queries and corresponding labels for one image.

Applies the global encoder. We support two strategies to merge encoded tokens from each view: In the first strategy, we extract the CLS toke

Greedily assigns workers to jobs based on an adjaceny matrix. Starting with an adjacency matrix representing the available connections in the bi-

Reshapes a batch of image tiles into a mosaic.

Converts an image to patches. Args: vid: A 5D tensor of shape [B, T, H, W, C]. patch_size: integer, dimension of a square patch.

Improves an assignment by backtracking the augmented path using state. Args: assignment: A bool [b, n, m] array, where each element of the inne

Inflate tensor with an extra mean channel.

Initialization function to be jitted.

Initialization function to be jitted.