Load the single mask for the given frame index and convert it to binary segments. Args: frame_idx (int): Index of the fr

Initializes the state variables for the video trainer. This method sets the initial state of various attributes used in the video

Resets the state of the video trainer. This method clears the internal state variables, setting them to their initial values:

Get the image feature on the input batch.

Propagate the input points across frames to track in the entire video.

Prepare and flatten visual features.

Add new mask to a frame.

Add new points to a frame.

Apply non-overlapping constraints to the object scores in pred_masks. Here we keep only the highest scoring object at each spatial lo

Simple wrapper for correctly getting worldsize in both distributed / non-distributed settings

Create the directory if it does not exist.

Apply non-overlapping constraints to the object scores in pred_masks. Here we keep only the highest scoring object at each spatial lo

Consolidate the per-object temporary outputs in `temp_output_dict_per_obj` on a frame into a single output for all objects, including

Consolidate the per-object temporary outputs in `temp_output_dict_per_obj` on a frame into a single output for all objects, including

Resize the object scores to the original video resolution (video_res_masks) and apply non-overlapping constraints for final output.

Resize the object scores to the original video resolution (video_res_masks) and apply non-overlapping constraints for final output.

Resize the object scores to the original video resolution (video_res_masks) and apply non-overlapping constraints for final output.

Resize the object scores to the original video resolution (video_res_masks) and apply non-overlapping constraints for final output.

Returns the positional encoding used to encode point prompts, applied to a dense set of points the shape of the image encoding.

Simple wrapper for correctly getting rank in both distributed / non-distributed settings

Go through the list of parameter names and select those that match any of the provided constraints

Consolidate the per-object temporary outputs in `temp_output_dict_per_obj` on a frame into a single output for all objects, including

Consolidate the per-object temporary outputs in `temp_output_dict_per_obj` on a frame into a single output for all objects, including

Forward SAM prompt encoders and mask heads. Inputs: - backbone_features: image features of [B, C, H, W] shape - poin

Compute the image features on a given frame.

Compute the image features on a given frame.

Get the total number of unique object ids received so far in this session.

Get the total number of unique object ids received so far in this session.

Get the total number of unique object ids received so far in this session.

Get the total number of unique object ids received so far in this session.

Predict masks for the given input prompts, using the currently set image. Input prompts are batched torch tensors and are expected to

Prepare and flatten visual features.

Wrapper over torch.distributed.all_reduce for performing reduction of tensor over all processes in both distributed / non-distributed sce

Wrapper over torch.distributed.broadcast for broadcasting a tensor from the source to all processes in both distributed / non-distributed sce

Add new points and labels to previous point inputs (add at the end).

For some backends, such as NCCL, communication only works if the tensor is on the GPU. This helper function converts to the correct devic

For some backends, such as NCCL, communication only works if the tensor is on the GPU. This converts the tensor back to original device.

Returns the positional encoding used to encode point prompts, applied to a dense set of points the shape of the image encoding.

Expects a torch tensor with length 2 in the last dimension. The coordinates can be in absolute image or normalized coordinates, If th

Split a multi-object output into per-object output slices and add them into `output_dict_per_obj`. The resulting slices share the sam

Split a multi-object output into per-object output slices and add them into `output_dict_per_obj`. The resulting slices share the sam

Encode the current image and its prediction into a memory feature.

Compute the image features on a given frame.

Compute the image features on a given frame.

`maskmem_pos_enc` is the same across frames and objects, so we cache it as a constant in the inference session to reduce session stor

`maskmem_pos_enc` is the same across frames and objects, so we cache it as a constant in the inference session to reduce session stor

Map client-side object id to model-side object index.

Map client-side object id to model-side object index.

Map client-side object id to model-side object index.

Map client-side object id to model-side object index.

Run tracking on a single frame based on current inputs and previous memory.

Run tracking on a single frame based on current inputs and previous memory.

Run tracking on a single frame based on current inputs and previous memory.

Run tracking on a single frame based on current inputs and previous memory.

Whether to use multimask output in the SAM head.

Add new points to a frame.

Wrapper over torch.distributed.barrier, returns without waiting if the distributed process group is not initialized instead of throwing error

Close writer, flushing pending logs to disk. Logs cannot be written after `close` is called.

Remove from the state dictionary the parameters matching the provided unix patterns Args: patterns: the list of unix patterns to exc

Get the best available device in the order: CUDA, MPS, CPU Returns: device string for torch.device

Get the connected components (8-connectivity) of binary masks of shape (N, 1, H, W). Inputs: - mask: A binary mask tensor of shape (N, 1