Reference for ultralytics/models/yolo/yoloe/predict.py

def get_vpe(self, source):
    """
    Process the source to get the visual prompt embeddings (VPE).

    Args:
        source (str | Path | int | PIL.Image | np.ndarray | torch.Tensor | List | Tuple): The source
            of the image to make predictions on. Accepts various types including file paths, URLs, PIL
            images, numpy arrays, and torch tensors.

    Returns:
        (torch.Tensor): The visual prompt embeddings (VPE) from the model.
    """
    self.setup_source(source)
    assert len(self.dataset) == 1, "get_vpe only supports one image!"
    for _, im0s, _ in self.dataset:
        im = self.preprocess(im0s)
        return self.model(im, vpe=self.prompts, return_vpe=True)

def inference(self, im, *args, **kwargs):
    """
    Run inference with visual prompts.

    Args:
        im (torch.Tensor): Input image tensor.
        *args (Any): Variable length argument list.
        **kwargs (Any): Arbitrary keyword arguments.

    Returns:
        (torch.Tensor): Model prediction results.
    """
    return super().inference(im, vpe=self.prompts, *args, **kwargs)

def pre_transform(self, im):
    """
    Preprocess images and prompts before inference.

    This method applies letterboxing to the input image and transforms the visual prompts
    (bounding boxes or masks) accordingly.

    Args:
        im (list): List containing a single input image.

    Returns:
        (list): Preprocessed image ready for model inference.

    Raises:
        ValueError: If neither valid bounding boxes nor masks are provided in the prompts.
    """
    img = super().pre_transform(im)
    bboxes = self.prompts.pop("bboxes", None)
    masks = self.prompts.pop("masks", None)
    category = self.prompts["cls"]
    if len(img) == 1:
        visuals = self._process_single_image(img[0].shape[:2], im[0].shape[:2], category, bboxes, masks)
        self.prompts = visuals.unsqueeze(0).to(self.device)  # (1, N, H, W)
    else:
        # NOTE: only supports bboxes as prompts for now
        assert bboxes is not None, f"Expected bboxes, but got {bboxes}!"
        # NOTE: needs List[np.ndarray]
        assert isinstance(bboxes, list) and all(isinstance(b, np.ndarray) for b in bboxes), (
            f"Expected List[np.ndarray], but got {bboxes}!"
        )
        assert isinstance(category, list) and all(isinstance(b, np.ndarray) for b in category), (
            f"Expected List[np.ndarray], but got {category}!"
        )
        assert len(im) == len(category) == len(bboxes), (
            f"Expected same length for all inputs, but got {len(im)}vs{len(category)}vs{len(bboxes)}!"
        )
        visuals = [
            self._process_single_image(img[i].shape[:2], im[i].shape[:2], category[i], bboxes[i])
            for i in range(len(img))
        ]
        self.prompts = torch.nn.utils.rnn.pad_sequence(visuals, batch_first=True).to(self.device)

    return img

def set_prompts(self, prompts):
    """
    Set the visual prompts for the model.

    Args:
        prompts (dict): Dictionary containing class indices and bounding boxes or masks.
            Must include a 'cls' key with class indices.
    """
    self.prompts = prompts

def setup_model(self, model, verbose: bool = True):
    """
    Set up the model for prediction.

    Args:
        model (torch.nn.Module): Model to load or use.
        verbose (bool, optional): If True, provides detailed logging.
    """
    super().setup_model(model, verbose=verbose)
    self.done_warmup = True