Save masks to a directory as PNG files.

Get the context for the attention scaled dot-product kernel. We use Flash Attention by default, but fall back to all available kernels if Fla

Set up slurm job info

show bounding box on the image Parameters ---------- box : numpy.ndarray bounding box coordinates in the original image

show mask on the image Parameters ---------- mask : numpy.ndarray mask of the image ax : matplotlib.axes.Axes ax

Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002. Args: inputs: A float tensor of arbitrary shape.

Expects a tensor of shape Bx4. The coordinates can be in absolute image or normalized coordinates, if the coords are in absolute imag

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

Partition into non-overlapping windows with padding if needed. Args: x (tensor): input tokens with [B, H, W, C]. window_size

Window unpartition into original sequences and removing padding. Args: x (tensor): input tokens with [B * num_windows, window_size, w

Initialize the neck :param trunk: the backbone :param position_encoding: the positional encoding to use :param d_model: the d

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

Build SAM-style prompt encoder and mask decoder.

Build SAM-style prompt encoder and mask decoder.

Build SAM-style prompt encoder and mask decoder.

When outputting a single mask, if the stability score from the current single-mask output (based on output token 0) falls below a thr

When outputting a single mask, if the stability score from the current single-mask output (based on output token 0) falls below a thr

Embed bounding boxes.

Embeds box prompts.

Embeds box prompts.

Embeds mask inputs.

Embeds mask inputs.

Embeds point prompts.

Embeds point prompts.

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

Identical to the corresponding method in the parent (EfficientTAMVideoPredictor), but cloning the memories and their pos enc to enabl

Identical to the corresponding method in the parent (EfficientTAMVideoPredictor), but cloning the memories and their pos enc to enabl

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

Compute the losses related to the masks: the focal loss and the dice loss. and also the MAE or MSE loss between predicted IoUs and ac

Gets the batch size of the output given the batch size of the input prompts.

Gets the batch size of the output given the batch size of the input prompts.

Get a dummy object pointer based on an empty mask on the current frame.

Get a dummy object pointer based on an empty mask on the current frame.

Create a list of dataset indices (with repeats) to use for one epoch. Args: generator (torch.Generator): pseudo random n

Generates a placeholder for point coordinates and labels. Args: batch_size (int, optional): The size of the batch. If no

Compute stability scores of the mask logits based on the IoU between upper and lower thresholds.

Compute stability scores of the mask logits based on the IoU between upper and lower thresholds.

Initialize memory for the first frame.

load multiple png masks from the disk (path: f'{obj_id}/{frame_id:05d}.png') Args: frame_id: int, define the mask path

load single png from the disk (path: f'{self.obj_id}/{frame_id:05d}.png') Args: frame_id: int, define the mask path

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

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

Predict masks and IoUs for the first frame.

Predict masks and IoUs for subsequent frames using memory.

Compute the image backbone features on the fly for the given img_ids.

Compute the image backbone features on the fly for the given img_ids.

Prepare memory for the current frame.

Fuse the current frame's visual feature map with previous memory.

Fuse the current frame's visual feature map with previous memory.

Fuse the current frame's visual feature map with previous memory.

Reset all tracking inputs and results across the videos.

Reset all tracking inputs and results across the videos.

Run the memory encoder on `high_res_masks`. This is usually after applying non-overlapping constraints to object scores. Since their

Run the memory encoder on `high_res_masks`. This is usually after applying non-overlapping constraints to object scores. Since their

Run the memory encoder on `high_res_masks`. This is usually after applying non-overlapping constraints to object scores. Since their

Run the memory encoder on `high_res_masks`. This is usually after applying non-overlapping constraints to object scores. Since their

Run the forward / backward

Save a checkpoint while guarding against the job being killed in the middle of checkpoint saving (which corrupts the checkpoint file