Visualize Predictions

The plot() method allows you to visualize predictions from a detection or segmentation model.

If you use a detection model to run inference (i.e. Grounding DINO), the plot() method will plot bounding boxes for each prediction.

If you use a segmentation model to run inference (i.e. Grounded SAM), the plot() method will plot segmentation masks for each prediction.

Here is an example of the method used to annotate predictions from a Grounding DINO model:

Bounding BoxSegmentation Mask

from autodistill_grounding_dino import GroundingDINO
from autodistill.detection import CaptionOntology
from autodistill.utils.plot import plot
import cv2

ontology = CaptionOntology(
    {
        "dog": "dog",
    }
)

model = GroundingDINO(ontology=ontology)

result = model.predict("./dog.jpeg")

plot(
    image=cv2.imread("./dog.jpeg"),
    classes=base_model.ontology.classes(),
    detections=result
)

from autodistill_grounded_sam import GroundedSAM
from autodistill.detection import CaptionOntology
from autodistill.utils import plot
import cv2

ontology = CaptionOntology(
    {
        "dog": "dog",
    }
)

model = GroundedSAM(ontology=ontology)

result = model.predict("./dog.jpeg")

plot(
    image=cv2.imread("./dog.jpeg"),
    classes=model.ontology.classes(),
    detections=result
)

Code Reference

Plot bounding boxes or segmentation masks on an image.

Parameters:

Name	Type	Description	Default
image	np.ndarray	The image to plot on	required
detections		The detections to plot	required
classes	List[str]	The classes to plot	required
raw		Whether to return the raw image or plot it interactively	False

Returns:

Type	Description
	The raw image (np.ndarray) if raw=True, otherwise None (image is plotted interactively

Source code in autodistill/utils.py

def plot(image: np.ndarray, detections, classes: List[str], raw=False):
    """
    Plot bounding boxes or segmentation masks on an image.

    Args:
        image: The image to plot on
        detections: The detections to plot
        classes: The classes to plot
        raw: Whether to return the raw image or plot it interactively

    Returns:
        The raw image (np.ndarray) if raw=True, otherwise None (image is plotted interactively
    """
    # TODO: When we have a classification annotator
    # in supervision, we can add it here
    if detections.mask is not None:
        annotator = sv.MaskAnnotator()
    else:
        annotator = sv.BoundingBoxAnnotator()

    label_annotator = sv.LabelAnnotator()

    labels = [
        f"{classes[class_id]} {confidence:0.2f}"
        for _, _, confidence, class_id, _, _ in detections
    ]

    annotated_frame = annotator.annotate(scene=image.copy(), detections=detections)
    annotated_frame = label_annotator.annotate(
        scene=annotated_frame, labels=labels, detections=detections
    )

    if raw:
        return annotated_frame

    sv.plot_image(annotated_frame, size=(8, 8))