vision_agent.tools

vision_agent.tools

register_tool

register_tool(imports=None)

Source code in vision_agent/tools/__init__.py

def register_tool(imports: Optional[List] = None) -> Callable:
    def decorator(tool: Callable) -> Callable:
        import inspect

        from .tools import (  # noqa: F811
            get_tool_descriptions,
            get_tools_df,
            get_tools_info,
        )

        global TOOLS, TOOLS_DF, TOOL_DESCRIPTIONS, TOOL_DOCSTRING, TOOLS_INFO

        if tool not in TOOLS:
            TOOLS.append(tool)
            TOOLS_DF = get_tools_df(TOOLS)  # type: ignore
            TOOL_DESCRIPTIONS = get_tool_descriptions(TOOLS)  # type: ignore
            TOOL_DOCSTRING = get_tool_documentation(TOOLS)  # type: ignore
            TOOLS_INFO = get_tools_info(TOOLS)  # type: ignore

            globals()[tool.__name__] = tool
            if imports is not None:
                for import_ in imports:
                    __new_tools__.append(import_)
            __new_tools__.append(inspect.getsource(tool))
        return tool

    return decorator

vision_agent.tools.tools

COLORS module-attribute

COLORS = [
    (158, 218, 229),
    (219, 219, 141),
    (23, 190, 207),
    (188, 189, 34),
    (199, 199, 199),
    (247, 182, 210),
    (127, 127, 127),
    (227, 119, 194),
    (196, 156, 148),
    (197, 176, 213),
    (140, 86, 75),
    (148, 103, 189),
    (255, 152, 150),
    (152, 223, 138),
    (214, 39, 40),
    (44, 160, 44),
    (255, 187, 120),
    (174, 199, 232),
    (255, 127, 14),
    (31, 119, 180),
]

FUNCTION_TOOLS module-attribute

FUNCTION_TOOLS = [
    owlv2_object_detection,
    owlv2_sam2_instance_segmentation,
    owlv2_sam2_video_tracking,
    countgd_object_detection,
    countgd_sam2_instance_segmentation,
    countgd_sam2_video_tracking,
    florence2_ocr,
    florence2_sam2_instance_segmentation,
    florence2_sam2_video_tracking,
    florence2_object_detection,
    claude35_text_extraction,
    document_extraction,
    document_qa,
    ocr,
    qwen2_vl_images_vqa,
    qwen2_vl_video_vqa,
    depth_anything_v2,
    generate_pose_image,
    vit_nsfw_classification,
    video_temporal_localization,
    flux_image_inpainting,
    siglip_classification,
    minimum_distance,
]

UTIL_TOOLS module-attribute

UTIL_TOOLS = [
    extract_frames_and_timestamps,
    save_json,
    load_image,
    save_image,
    save_video,
    overlay_bounding_boxes,
    overlay_segmentation_masks,
    overlay_heat_map,
]

TOOLS module-attribute

TOOLS = FUNCTION_TOOLS + UTIL_TOOLS

TOOLS_DF module-attribute

TOOLS_DF = get_tools_df(TOOLS)

TOOL_DESCRIPTIONS module-attribute

TOOL_DESCRIPTIONS = get_tool_descriptions(TOOLS)

TOOL_DOCSTRING module-attribute

TOOL_DOCSTRING = get_tool_documentation(TOOLS)

TOOLS_INFO module-attribute

TOOLS_INFO = get_tools_info(FUNCTION_TOOLS)

UTILITIES_DOCSTRING module-attribute

UTILITIES_DOCSTRING = get_tool_documentation(UTIL_TOOLS)

sam2

sam2(image, detections)

'sam2' is a tool that can segment multiple objects given an input bounding box, label and score. It returns a set of masks along with the corresponding bounding boxes and labels.

PARAMETER	DESCRIPTION
image	The image that contains multiple instances of the object. TYPE: ndarray
detections	A list of dictionaries containing the score, label, and bounding box of the detected objects with normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box. TYPE: List[Dict[str, Any]]

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label, bounding box, and mask of the detected objects with normalized coordinates (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box. The mask is binary 2D numpy array where 1 indicates the object and 0 indicates the background.

Example

>>> sam2(image, [
        {'score': 0.49, 'label': 'flower', 'bbox': [0.1, 0.11, 0.35, 0.4]},
    ])
[
    {
        'score': 0.49,
        'label': 'flower',
        'bbox': [0.1, 0.11, 0.35, 0.4],
        'mask': array([[0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0],
            ...,
            [0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
    },
]

Source code in vision_agent/tools/tools.py

def sam2(
    image: np.ndarray,
    detections: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
    """'sam2' is a tool that can segment multiple objects given an input bounding box,
    label and score. It returns a set of masks along with the corresponding bounding
    boxes and labels.

    Parameters:
        image (np.ndarray): The image that contains multiple instances of the object.
        detections (List[Dict[str, Any]]): A list of dictionaries containing the score,
            label, and bounding box of the detected objects with normalized coordinates
            between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates
            of the top-left and xmax and ymax are the coordinates of the bottom-right of
            the bounding box.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label,
            bounding box, and mask of the detected objects with normalized coordinates
            (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left
            and xmax and ymax are the coordinates of the bottom-right of the bounding box.
            The mask is binary 2D numpy array where 1 indicates the object and 0 indicates
            the background.

    Example
    -------
        >>> sam2(image, [
                {'score': 0.49, 'label': 'flower', 'bbox': [0.1, 0.11, 0.35, 0.4]},
            ])
        [
            {
                'score': 0.49,
                'label': 'flower',
                'bbox': [0.1, 0.11, 0.35, 0.4],
                'mask': array([[0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0],
                    ...,
                    [0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
            },
        ]
    """
    image_size = image.shape[:2]
    ret = _sam2(image, detections, image_size)
    _display_tool_trace(
        sam2.__name__,
        {},
        ret["display_data"],
        ret["files"],
    )

    return ret["return_data"]  # type: ignore

od_sam2_video_tracking

od_sam2_video_tracking(
    od_model,
    prompt,
    frames,
    chunk_length=10,
    fine_tune_id=None,
)

Source code in vision_agent/tools/tools.py

def od_sam2_video_tracking(
    od_model: ODModels,
    prompt: str,
    frames: List[np.ndarray],
    chunk_length: Optional[int] = 10,
    fine_tune_id: Optional[str] = None,
) -> Dict[str, Any]:
    SEGMENT_SIZE = 50
    OVERLAP = 1  # Number of overlapping frames between segments

    image_size = frames[0].shape[:2]

    # Split frames into segments with overlap
    segments = split_frames_into_segments(frames, SEGMENT_SIZE, OVERLAP)

    def _apply_object_detection(  # inner method to avoid circular importing issues.
        od_model: ODModels,
        prompt: str,
        segment_index: int,
        frame_number: int,
        fine_tune_id: str,
        segment_frames: list,
    ) -> tuple:
        """
        Applies the specified object detection model to the given image.

        Args:
            od_model: The object detection model to use.
            prompt: The prompt for the object detection model.
            segment_index: The index of the current segment.
            frame_number: The number of the current frame.
            fine_tune_id: Optional fine-tune ID for the model.
            segment_frames: List of frames for the current segment.

        Returns:
            A tuple containing the object detection results and the name of the function used.
        """

        if od_model == ODModels.COUNTGD:
            segment_results = countgd_object_detection(
                prompt=prompt, image=segment_frames[frame_number]
            )
            function_name = "countgd_object_detection"

        elif od_model == ODModels.OWLV2:
            segment_results = owlv2_object_detection(
                prompt=prompt,
                image=segment_frames[frame_number],
                fine_tune_id=fine_tune_id,
            )
            function_name = "owlv2_object_detection"

        elif od_model == ODModels.FLORENCE2:
            segment_results = florence2_object_detection(
                prompt=prompt,
                image=segment_frames[frame_number],
                fine_tune_id=fine_tune_id,
            )
            function_name = "florence2_object_detection"

        elif od_model == ODModels.AGENTIC:
            segment_results = agentic_object_detection(
                prompt=prompt,
                image=segment_frames[frame_number],
                fine_tune_id=fine_tune_id,
            )
            function_name = "agentic_object_detection"

        elif od_model == ODModels.CUSTOM:
            segment_results = custom_object_detection(
                deployment_id=fine_tune_id,
                image=segment_frames[frame_number],
            )
            function_name = "custom_object_detection"

        else:
            raise NotImplementedError(
                f"Object detection model '{od_model}' is not implemented."
            )

        return segment_results, function_name

    # Process each segment and collect detections
    detections_per_segment: List[Any] = []
    for segment_index, segment in enumerate(segments):
        segment_detections = process_segment(
            segment_frames=segment,
            od_model=od_model,
            prompt=prompt,
            fine_tune_id=fine_tune_id,
            chunk_length=chunk_length,
            image_size=image_size,
            segment_index=segment_index,
            object_detection_tool=_apply_object_detection,
        )
        detections_per_segment.append(segment_detections)

    merged_detections = merge_segments(detections_per_segment)
    post_processed = post_process(merged_detections, image_size)

    buffer_bytes = frames_to_bytes(frames)
    files = [("video", buffer_bytes)]

    return {
        "files": files,
        "return_data": post_processed["return_data"],
        "display_data": post_processed["display_data"],
    }

owlv2_object_detection

owlv2_object_detection(
    prompt, image, box_threshold=0.1, fine_tune_id=None
)

'owlv2_object_detection' is a tool that can detect and count multiple objects given a text prompt such as category names or referring expressions on images. The categories in text prompt are separated by commas. It returns a list of bounding boxes with normalized coordinates, label names and associated probability scores.

PARAMETER	DESCRIPTION
prompt	The prompt to ground to the image. TYPE: str
image	The image to ground the prompt to. TYPE: ndarray
box_threshold	The threshold for the box detection. Defaults to 0.10. TYPE: float DEFAULT: 0.1
fine_tune_id	If you have a fine-tuned model, you can pass the fine-tuned model ID here to use it. TYPE: Optional[str] DEFAULT: None

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label, and bounding box of the detected objects with normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box.

Example

>>> owlv2_object_detection("car, dinosaur", image)
[
    {'score': 0.99, 'label': 'dinosaur', 'bbox': [0.1, 0.11, 0.35, 0.4]},
    {'score': 0.98, 'label': 'car', 'bbox': [0.2, 0.21, 0.45, 0.5},
]

Source code in vision_agent/tools/tools.py

def owlv2_object_detection(
    prompt: str,
    image: np.ndarray,
    box_threshold: float = 0.10,
    fine_tune_id: Optional[str] = None,
) -> List[Dict[str, Any]]:
    """'owlv2_object_detection' is a tool that can detect and count multiple objects
    given a text prompt such as category names or referring expressions on images. The
    categories in text prompt are separated by commas. It returns a list of bounding
    boxes with normalized coordinates, label names and associated probability scores.

    Parameters:
        prompt (str): The prompt to ground to the image.
        image (np.ndarray): The image to ground the prompt to.
        box_threshold (float, optional): The threshold for the box detection. Defaults
            to 0.10.
        fine_tune_id (Optional[str]): If you have a fine-tuned model, you can pass the
            fine-tuned model ID here to use it.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label, and
            bounding box of the detected objects with normalized coordinates between 0
            and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the
            top-left and xmax and ymax are the coordinates of the bottom-right of the
            bounding box.

    Example
    -------
        >>> owlv2_object_detection("car, dinosaur", image)
        [
            {'score': 0.99, 'label': 'dinosaur', 'bbox': [0.1, 0.11, 0.35, 0.4]},
            {'score': 0.98, 'label': 'car', 'bbox': [0.2, 0.21, 0.45, 0.5},
        ]
    """

    image_size = image.shape[:2]
    if image_size[0] < 1 or image_size[1] < 1:
        return []

    ret = _owlv2_object_detection(
        prompt, image, box_threshold, image_size, fine_tune_id=fine_tune_id
    )

    _display_tool_trace(
        owlv2_object_detection.__name__,
        {
            "prompts": prompt,
            "confidence": box_threshold,
        },
        ret["display_data"],
        ret["files"],
    )
    return ret["return_data"]  # type: ignore

owlv2_sam2_instance_segmentation

owlv2_sam2_instance_segmentation(
    prompt, image, box_threshold=0.1
)

'owlv2_sam2_instance_segmentation' is a tool that can detect and count multiple instances of objects given a text prompt such as category names or referring expressions on images. The categories in text prompt are separated by commas. It returns a list of bounding boxes with normalized coordinates, label names, masks and associated probability scores.

PARAMETER	DESCRIPTION
prompt	The object that needs to be counted. TYPE: str
image	The image that contains multiple instances of the object. TYPE: ndarray
box_threshold	The threshold for detection. Defaults to 0.10. TYPE: float DEFAULT: 0.1

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label, bounding box, and mask of the detected objects with normalized coordinates (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box. The mask is binary 2D numpy array where 1 indicates the object and 0 indicates the background.

Example

>>> owlv2_sam2_instance_segmentation("flower", image)
[
    {
        'score': 0.49,
        'label': 'flower',
        'bbox': [0.1, 0.11, 0.35, 0.4],
        'mask': array([[0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0],
            ...,
            [0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
    },
]

Source code in vision_agent/tools/tools.py

def owlv2_sam2_instance_segmentation(
    prompt: str,
    image: np.ndarray,
    box_threshold: float = 0.10,
) -> List[Dict[str, Any]]:
    """'owlv2_sam2_instance_segmentation' is a tool that can detect and count multiple
    instances of objects given a text prompt such as category names or referring
    expressions on images. The categories in text prompt are separated by commas. It
    returns a list of bounding boxes with normalized coordinates, label names, masks
    and associated probability scores.

    Parameters:
        prompt (str): The object that needs to be counted.
        image (np.ndarray): The image that contains multiple instances of the object.
        box_threshold (float, optional): The threshold for detection. Defaults
            to 0.10.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label,
            bounding box, and mask of the detected objects with normalized coordinates
            (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left
            and xmax and ymax are the coordinates of the bottom-right of the bounding box.
            The mask is binary 2D numpy array where 1 indicates the object and 0 indicates
            the background.

    Example
    -------
        >>> owlv2_sam2_instance_segmentation("flower", image)
        [
            {
                'score': 0.49,
                'label': 'flower',
                'bbox': [0.1, 0.11, 0.35, 0.4],
                'mask': array([[0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0],
                    ...,
                    [0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
            },
        ]
    """

    od_ret = _owlv2_object_detection(prompt, image, box_threshold, image.shape[:2])
    seg_ret = _sam2(
        image, od_ret["return_data"], image.shape[:2], image_bytes=od_ret["files"][0][1]
    )

    _display_tool_trace(
        owlv2_sam2_instance_segmentation.__name__,
        {
            "prompts": prompt,
            "confidence": box_threshold,
        },
        seg_ret["display_data"],
        seg_ret["files"],
    )

    return seg_ret["return_data"]  # type: ignore

owlv2_sam2_video_tracking

owlv2_sam2_video_tracking(
    prompt, frames, chunk_length=10, fine_tune_id=None
)

'owlv2_sam2_video_tracking' is a tool that can track and segment multiple objects in a video given a text prompt such as category names or referring expressions. The categories in the text prompt are separated by commas. It returns a list of bounding boxes, label names, masks and associated probability scores and is useful for tracking and counting without duplicating counts.

PARAMETER	DESCRIPTION
prompt	The prompt to ground to the image. TYPE: str
frames	The list of frames to ground the prompt to. TYPE: List[ndarray]
chunk_length	The number of frames to re-run owlv2 to find new objects. TYPE: Optional[int] DEFAULT: 10
fine_tune_id	If you have a fine-tuned model, you can pass the fine-tuned model ID here to use it. TYPE: Optional[str] DEFAULT: None

RETURNS

DESCRIPTION

List[List[Dict[str, Any]]]

List[List[Dict[str, Any]]]: A list of list of dictionaries containing the label, segmentation mask and bounding boxes. The outer list represents each frame and the inner list is the entities per frame. The detected objects have normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box. The mask is binary 2D numpy array where 1 indicates the object and 0 indicates the background. The label names are prefixed with their ID represent the total count.

Example

>>> owlv2_sam2_video_tracking("car, dinosaur", frames)
[
    [
        {
            'label': '0: dinosaur',
            'bbox': [0.1, 0.11, 0.35, 0.4],
            'mask': array([[0, 0, 0, ..., 0, 0, 0],
                [0, 0, 0, ..., 0, 0, 0],
                ...,
                [0, 0, 0, ..., 0, 0, 0],
                [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
        },
    ],
    ...
]

Source code in vision_agent/tools/tools.py

def owlv2_sam2_video_tracking(
    prompt: str,
    frames: List[np.ndarray],
    chunk_length: Optional[int] = 10,
    fine_tune_id: Optional[str] = None,
) -> List[List[Dict[str, Any]]]:
    """'owlv2_sam2_video_tracking' is a tool that can track and segment multiple
    objects in a video given a text prompt such as category names or referring
    expressions. The categories in the text prompt are separated by commas. It returns
    a list of bounding boxes, label names, masks and associated probability scores and
    is useful for tracking and counting without duplicating counts.

    Parameters:
        prompt (str): The prompt to ground to the image.
        frames (List[np.ndarray]): The list of frames to ground the prompt to.
        chunk_length (Optional[int]): The number of frames to re-run owlv2 to find
            new objects.
        fine_tune_id (Optional[str]): If you have a fine-tuned model, you can pass the
            fine-tuned model ID here to use it.

    Returns:
        List[List[Dict[str, Any]]]: A list of list of dictionaries containing the
            label, segmentation mask and bounding boxes. The outer list represents each
            frame and the inner list is the entities per frame. The detected objects
            have normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin
            and ymin are the coordinates of the top-left and xmax and ymax are the
            coordinates of the bottom-right of the bounding box. The mask is binary 2D
            numpy array where 1 indicates the object and 0 indicates the background.
            The label names are prefixed with their ID represent the total count.

    Example
    -------
        >>> owlv2_sam2_video_tracking("car, dinosaur", frames)
        [
            [
                {
                    'label': '0: dinosaur',
                    'bbox': [0.1, 0.11, 0.35, 0.4],
                    'mask': array([[0, 0, 0, ..., 0, 0, 0],
                        [0, 0, 0, ..., 0, 0, 0],
                        ...,
                        [0, 0, 0, ..., 0, 0, 0],
                        [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
                },
            ],
            ...
        ]
    """

    ret = od_sam2_video_tracking(
        ODModels.OWLV2,
        prompt=prompt,
        frames=frames,
        chunk_length=chunk_length,
        fine_tune_id=fine_tune_id,
    )
    _display_tool_trace(
        owlv2_sam2_video_tracking.__name__,
        {},
        ret["display_data"],
        ret["files"],
    )
    return ret["return_data"]  # type: ignore

florence2_object_detection

florence2_object_detection(
    prompt, image, fine_tune_id=None
)

'florence2_object_detection' is a tool that can detect multiple objects given a text prompt which can be object names or caption. You can optionally separate the object names in the text with commas. It returns a list of bounding boxes with normalized coordinates, label names and associated confidence scores of 1.0.

PARAMETER	DESCRIPTION
prompt	The prompt to ground to the image. Use exclusive categories that do not overlap such as 'person, car' and NOT 'person, athlete'. TYPE: str
image	The image to used to detect objects TYPE: ndarray
fine_tune_id	If you have a fine-tuned model, you can pass the fine-tuned model ID here to use it. TYPE: Optional[str] DEFAULT: None

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label, and bounding box of the detected objects with normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box. The scores are always 1.0 and cannot be thresholded

Example

>>> florence2_object_detection('person looking at a coyote', image)
[
    {'score': 1.0, 'label': 'person', 'bbox': [0.1, 0.11, 0.35, 0.4]},
    {'score': 1.0, 'label': 'coyote', 'bbox': [0.34, 0.21, 0.85, 0.5},
]

Source code in vision_agent/tools/tools.py

def florence2_object_detection(
    prompt: str, image: np.ndarray, fine_tune_id: Optional[str] = None
) -> List[Dict[str, Any]]:
    """'florence2_object_detection' is a tool that can detect multiple objects given a
    text prompt which can be object names or caption. You can optionally separate the
    object names in the text with commas. It returns a list of bounding boxes with
    normalized coordinates, label names and associated confidence scores of 1.0.

    Parameters:
        prompt (str): The prompt to ground to the image. Use exclusive categories that
            do not overlap such as 'person, car' and NOT 'person, athlete'.
        image (np.ndarray): The image to used to detect objects
        fine_tune_id (Optional[str]): If you have a fine-tuned model, you can pass the
            fine-tuned model ID here to use it.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label, and
            bounding box of the detected objects with normalized coordinates between 0
            and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the
            top-left and xmax and ymax are the coordinates of the bottom-right of the
            bounding box. The scores are always 1.0 and cannot be thresholded

    Example
    -------
        >>> florence2_object_detection('person looking at a coyote', image)
        [
            {'score': 1.0, 'label': 'person', 'bbox': [0.1, 0.11, 0.35, 0.4]},
            {'score': 1.0, 'label': 'coyote', 'bbox': [0.34, 0.21, 0.85, 0.5},
        ]
    """
    image_size = image.shape[:2]
    if image_size[0] < 1 or image_size[1] < 1:
        return []

    buffer_bytes = numpy_to_bytes(image)
    files = [("image", buffer_bytes)]
    payload = {
        "prompts": [s.strip() for s in prompt.split(",")],
        "model": "florence2",
    }
    metadata = {"function_name": "florence2_object_detection"}

    if fine_tune_id is not None:
        landing_api = LandingPublicAPI()
        status = landing_api.check_fine_tuning_job(UUID(fine_tune_id))
        if status is not JobStatus.SUCCEEDED:
            raise FineTuneModelIsNotReady(
                f"Fine-tuned model {fine_tune_id} is not ready yet"
            )

        payload["jobId"] = fine_tune_id

    detections = send_task_inference_request(
        payload,
        "text-to-object-detection",
        files=files,
        metadata=metadata,
    )

    # get the first frame
    bboxes = detections[0]
    bboxes_formatted = [
        {
            "label": bbox["label"],
            "bbox": normalize_bbox(bbox["bounding_box"], image_size),
            "score": round(bbox["score"], 2),
        }
        for bbox in bboxes
    ]

    _display_tool_trace(
        florence2_object_detection.__name__,
        payload,
        detections[0],
        files,
    )
    return [bbox for bbox in bboxes_formatted]

florence2_sam2_instance_segmentation

florence2_sam2_instance_segmentation(
    prompt, image, fine_tune_id=None
)

'florence2_sam2_instance_segmentation' is a tool that can segment multiple objects given a text prompt such as category names or referring expressions. The categories in the text prompt are separated by commas. It returns a list of bounding boxes, label names, mask file names and associated probability scores of 1.0.

PARAMETER	DESCRIPTION
prompt	The prompt to ground to the image. Use exclusive categories that do not overlap such as 'person, car' and NOT 'person, athlete'. TYPE: str
image	The image to ground the prompt to. TYPE: ndarray
fine_tune_id	If you have a fine-tuned model, you can pass the fine-tuned model ID here to use it. TYPE: Optional[str] DEFAULT: None

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label, bounding box, and mask of the detected objects with normalized coordinates (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box. The mask is binary 2D numpy array where 1 indicates the object and 0 indicates the background.

Example

>>> florence2_sam2_instance_segmentation("car, dinosaur", image)
[
    {
        'score': 1.0,
        'label': 'dinosaur',
        'bbox': [0.1, 0.11, 0.35, 0.4],
        'mask': array([[0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0],
            ...,
            [0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
    },
]

Source code in vision_agent/tools/tools.py

def florence2_sam2_instance_segmentation(
    prompt: str, image: np.ndarray, fine_tune_id: Optional[str] = None
) -> List[Dict[str, Any]]:
    """'florence2_sam2_instance_segmentation' is a tool that can segment multiple
    objects given a text prompt such as category names or referring expressions. The
    categories in the text prompt are separated by commas. It returns a list of
    bounding boxes, label names, mask file names and associated probability scores of
    1.0.

    Parameters:
        prompt (str): The prompt to ground to the image. Use exclusive categories that
            do not overlap such as 'person, car' and NOT 'person, athlete'.
        image (np.ndarray): The image to ground the prompt to.
        fine_tune_id (Optional[str]): If you have a fine-tuned model, you can pass the
            fine-tuned model ID here to use it.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label,
            bounding box, and mask of the detected objects with normalized coordinates
            (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left
            and xmax and ymax are the coordinates of the bottom-right of the bounding box.
            The mask is binary 2D numpy array where 1 indicates the object and 0 indicates
            the background.

    Example
    -------
        >>> florence2_sam2_instance_segmentation("car, dinosaur", image)
        [
            {
                'score': 1.0,
                'label': 'dinosaur',
                'bbox': [0.1, 0.11, 0.35, 0.4],
                'mask': array([[0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0],
                    ...,
                    [0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
            },
        ]
    """
    if image.shape[0] < 1 or image.shape[1] < 1:
        return []

    buffer_bytes = numpy_to_bytes(image)
    files = [("image", buffer_bytes)]
    payload = {
        "prompt": prompt,
        "model": "florence2sam2",
    }
    metadata = {"function_name": "florence2_sam2_instance_segmentation"}

    if fine_tune_id is not None:
        landing_api = LandingPublicAPI()
        status = landing_api.check_fine_tuning_job(UUID(fine_tune_id))
        if status is not JobStatus.SUCCEEDED:
            raise FineTuneModelIsNotReady(
                f"Fine-tuned model {fine_tune_id} is not ready yet"
            )

        payload["jobId"] = fine_tune_id

    detections = send_task_inference_request(
        payload,
        "text-to-instance-segmentation",
        files=files,
        metadata=metadata,
    )

    # get the first frame
    frame = detections[0]
    return_data = []
    for detection in frame:
        mask = rle_decode_array(detection["mask"])
        label = detection["label"]
        bbox = normalize_bbox(detection["bounding_box"], detection["mask"]["size"])
        return_data.append({"label": label, "bbox": bbox, "mask": mask, "score": 1.0})

    _display_tool_trace(
        florence2_sam2_instance_segmentation.__name__,
        payload,
        detections[0],
        files,
    )
    return return_data

florence2_sam2_video_tracking

florence2_sam2_video_tracking(
    prompt, frames, chunk_length=10, fine_tune_id=None
)

'florence2_sam2_video_tracking' is a tool that can track and segment multiple objects in a video given a text prompt such as category names or referring expressions. The categories in the text prompt are separated by commas. It returns a list of bounding boxes, label names, masks and associated probability scores and is useful for tracking and counting without duplicating counts.

PARAMETER	DESCRIPTION
prompt	The prompt to ground to the image. Use exclusive categories that do not overlap such as 'person, car' and NOT 'person, athlete'. TYPE: str
frames	The list of frames to ground the prompt to. TYPE: List[ndarray]
chunk_length	The number of frames to re-run florence2 to find new objects. TYPE: Optional[int] DEFAULT: 10
fine_tune_id	If you have a fine-tuned model, you can pass the fine-tuned model ID here to use it. TYPE: Optional[str] DEFAULT: None

RETURNS

DESCRIPTION

List[List[Dict[str, Any]]]

List[List[Dict[str, Any]]]: A list of list of dictionaries containing the label, segmentation mask and bounding boxes. The outer list represents each frame and the inner list is the entities per frame. The detected objects have normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box. The mask is binary 2D numpy array where 1 indicates the object and 0 indicates the background. The label names are prefixed with their ID represent the total count.

Example

>>> florence2_sam2_video_tracking("car, dinosaur", frames)
[
    [
        {
            'label': '0: dinosaur',
            'bbox': [0.1, 0.11, 0.35, 0.4],
            'mask': array([[0, 0, 0, ..., 0, 0, 0],
                [0, 0, 0, ..., 0, 0, 0],
                ...,
                [0, 0, 0, ..., 0, 0, 0],
                [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
        },
    ],
    ...
]

Source code in vision_agent/tools/tools.py

def florence2_sam2_video_tracking(
    prompt: str,
    frames: List[np.ndarray],
    chunk_length: Optional[int] = 10,
    fine_tune_id: Optional[str] = None,
) -> List[List[Dict[str, Any]]]:
    """'florence2_sam2_video_tracking' is a tool that can track and segment multiple
    objects in a video given a text prompt such as category names or referring
    expressions. The categories in the text prompt are separated by commas. It returns
    a list of bounding boxes, label names, masks and associated probability scores and
    is useful for tracking and counting without duplicating counts.

    Parameters:
        prompt (str): The prompt to ground to the image. Use exclusive categories that
            do not overlap such as 'person, car' and NOT 'person, athlete'.
        frames (List[np.ndarray]): The list of frames to ground the prompt to.
        chunk_length (Optional[int]): The number of frames to re-run florence2 to find
            new objects.
        fine_tune_id (Optional[str]): If you have a fine-tuned model, you can pass the
            fine-tuned model ID here to use it.

    Returns:
        List[List[Dict[str, Any]]]: A list of list of dictionaries containing the
            label, segmentation mask and bounding boxes. The outer list represents each
            frame and the inner list is the entities per frame. The detected objects
            have normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin
            and ymin are the coordinates of the top-left and xmax and ymax are the
            coordinates of the bottom-right of the bounding box. The mask is binary 2D
            numpy array where 1 indicates the object and 0 indicates the background.
            The label names are prefixed with their ID represent the total count.

    Example
    -------
        >>> florence2_sam2_video_tracking("car, dinosaur", frames)
        [
            [
                {
                    'label': '0: dinosaur',
                    'bbox': [0.1, 0.11, 0.35, 0.4],
                    'mask': array([[0, 0, 0, ..., 0, 0, 0],
                        [0, 0, 0, ..., 0, 0, 0],
                        ...,
                        [0, 0, 0, ..., 0, 0, 0],
                        [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
                },
            ],
            ...
        ]
    """
    if len(frames) == 0 or not isinstance(frames, List):
        raise ValueError("Must provide a list of numpy arrays for frames")

    buffer_bytes = frames_to_bytes(frames)
    files = [("video", buffer_bytes)]
    payload = {
        "prompt": prompt,
        "model": "florence2sam2",
    }
    metadata = {"function_name": "florence2_sam2_video_tracking"}

    if chunk_length is not None:
        payload["chunk_length_frames"] = chunk_length  # type: ignore

    if fine_tune_id is not None:
        landing_api = LandingPublicAPI()
        status = landing_api.check_fine_tuning_job(UUID(fine_tune_id))
        if status is not JobStatus.SUCCEEDED:
            raise FineTuneModelIsNotReady(
                f"Fine-tuned model {fine_tune_id} is not ready yet"
            )

        payload["jobId"] = fine_tune_id

    detections = send_task_inference_request(
        payload,
        "text-to-instance-segmentation",
        files=files,
        metadata=metadata,
    )

    return_data = []
    for frame in detections:
        return_frame_data = []
        for detection in frame:
            mask = rle_decode_array(detection["mask"])
            label = str(detection["id"]) + ": " + detection["label"]
            return_frame_data.append(
                {"label": label, "mask": mask, "score": 1.0, "rle": detection["mask"]}
            )
        return_data.append(return_frame_data)
    return_data = add_bboxes_from_masks(return_data)
    return_data = nms(return_data, iou_threshold=0.95)

    _display_tool_trace(
        florence2_sam2_video_tracking.__name__,
        payload,
        [
            [
                {
                    "label": e["label"],
                    "score": e["score"],
                    "bbox": denormalize_bbox(e["bbox"], frames[0].shape[:2]),
                    "mask": e["rle"],
                }
                for e in lst
            ]
            for lst in return_data
        ],
        files,
    )
    # We save the RLE for display purposes, re-calculting RLE can get very expensive.
    # Deleted here because we are returning the numpy masks instead
    for frame in return_data:
        for obj in frame:
            del obj["rle"]
    return return_data

florence2_ocr

florence2_ocr(image)

'florence2_ocr' is a tool that can detect text and text regions in an image. Each text region contains one line of text. It returns a list of detected text, the text region as a bounding box with normalized coordinates, and confidence scores. The results are sorted from top-left to bottom right.

PARAMETER	DESCRIPTION
image	The image to extract text from. TYPE: ndarray

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the detected text, bbox with normalized coordinates, and confidence score.

Example

>>> florence2_ocr(image)
[
    {'label': 'hello world', 'bbox': [0.1, 0.11, 0.35, 0.4], 'score': 0.99},
]

Source code in vision_agent/tools/tools.py

def florence2_ocr(image: np.ndarray) -> List[Dict[str, Any]]:
    """'florence2_ocr' is a tool that can detect text and text regions in an image.
    Each text region contains one line of text. It returns a list of detected text,
    the text region as a bounding box with normalized coordinates, and confidence
    scores. The results are sorted from top-left to bottom right.

    Parameters:
        image (np.ndarray): The image to extract text from.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the detected text, bbox
            with normalized coordinates, and confidence score.

    Example
    -------
        >>> florence2_ocr(image)
        [
            {'label': 'hello world', 'bbox': [0.1, 0.11, 0.35, 0.4], 'score': 0.99},
        ]
    """

    image_size = image.shape[:2]
    if image_size[0] < 1 or image_size[1] < 1:
        return []
    image_b64 = convert_to_b64(image)
    data = {
        "image": image_b64,
        "task": "<OCR_WITH_REGION>",
        "function_name": "florence2_ocr",
    }

    detections = send_inference_request(data, "florence2", v2=True)
    detections = detections["<OCR_WITH_REGION>"]
    return_data = []
    for i in range(len(detections["quad_boxes"])):
        return_data.append(
            {
                "label": detections["labels"][i],
                "bbox": normalize_bbox(
                    convert_quad_box_to_bbox(detections["quad_boxes"][i]), image_size
                ),
                "score": 1.0,
            }
        )
    _display_tool_trace(
        florence2_ocr.__name__,
        {},
        detections,
        image_b64,
    )
    return return_data

countgd_object_detection

countgd_object_detection(prompt, image, box_threshold=0.23)

'countgd_object_detection' is a tool that can detect multiple instances of an object given a text prompt. It is particularly useful when trying to detect and count a large number of objects. You can optionally separate object names in the prompt with commas. It returns a list of bounding boxes with normalized coordinates, label names and associated confidence scores.

PARAMETER	DESCRIPTION
prompt	The object that needs to be counted. TYPE: str
image	The image that contains multiple instances of the object. TYPE: ndarray
box_threshold	The threshold for detection. Defaults to 0.23. TYPE: float DEFAULT: 0.23

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label, and bounding box of the detected objects with normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box.

Example

>>> countgd_object_detection("flower", image)
[
    {'score': 0.49, 'label': 'flower', 'bbox': [0.1, 0.11, 0.35, 0.4]},
    {'score': 0.68, 'label': 'flower', 'bbox': [0.2, 0.21, 0.45, 0.5},
    {'score': 0.78, 'label': 'flower', 'bbox': [0.3, 0.35, 0.48, 0.52},
    {'score': 0.98, 'label': 'flower', 'bbox': [0.44, 0.24, 0.49, 0.58},
]

Source code in vision_agent/tools/tools.py

def countgd_object_detection(
    prompt: str,
    image: np.ndarray,
    box_threshold: float = 0.23,
) -> List[Dict[str, Any]]:
    """'countgd_object_detection' is a tool that can detect multiple instances of an
    object given a text prompt. It is particularly useful when trying to detect and
    count a large number of objects. You can optionally separate object names in the
    prompt with commas. It returns a list of bounding boxes with normalized
    coordinates, label names and associated confidence scores.

    Parameters:
        prompt (str): The object that needs to be counted.
        image (np.ndarray): The image that contains multiple instances of the object.
        box_threshold (float, optional): The threshold for detection. Defaults
            to 0.23.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label, and
            bounding box of the detected objects with normalized coordinates between 0
            and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the
            top-left and xmax and ymax are the coordinates of the bottom-right of the
            bounding box.

    Example
    -------
        >>> countgd_object_detection("flower", image)
        [
            {'score': 0.49, 'label': 'flower', 'bbox': [0.1, 0.11, 0.35, 0.4]},
            {'score': 0.68, 'label': 'flower', 'bbox': [0.2, 0.21, 0.45, 0.5},
            {'score': 0.78, 'label': 'flower', 'bbox': [0.3, 0.35, 0.48, 0.52},
            {'score': 0.98, 'label': 'flower', 'bbox': [0.44, 0.24, 0.49, 0.58},
        ]
    """
    image_size = image.shape[:2]
    if image_size[0] < 1 or image_size[1] < 1:
        return []

    ret = _countgd_object_detection(prompt, image, box_threshold, image_size)
    _display_tool_trace(
        countgd_object_detection.__name__,
        {
            "prompts": prompt,
            "confidence": box_threshold,
        },
        ret["display_data"],
        ret["files"],
    )
    return ret["return_data"]  # type: ignore

countgd_sam2_instance_segmentation

countgd_sam2_instance_segmentation(
    prompt, image, box_threshold=0.23
)

'countgd_sam2_instance_segmentation' is a tool that can detect multiple instances of an object given a text prompt. It is particularly useful when trying to detect and count a large number of objects. You can optionally separate object names in the prompt with commas. It returns a list of bounding boxes with normalized coordinates, label names, masks associated confidence scores.

PARAMETER	DESCRIPTION
prompt	The object that needs to be counted. TYPE: str
image	The image that contains multiple instances of the object. TYPE: ndarray
box_threshold	The threshold for detection. Defaults to 0.23. TYPE: float DEFAULT: 0.23

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label, bounding box, and mask of the detected objects with normalized coordinates (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box. The mask is binary 2D numpy array where 1 indicates the object and 0 indicates the background.

Example

>>> countgd_sam2_instance_segmentation("flower", image)
[
    {
        'score': 0.49,
        'label': 'flower',
        'bbox': [0.1, 0.11, 0.35, 0.4],
        'mask': array([[0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0],
            ...,
            [0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
    },
]

Source code in vision_agent/tools/tools.py

def countgd_sam2_instance_segmentation(
    prompt: str,
    image: np.ndarray,
    box_threshold: float = 0.23,
) -> List[Dict[str, Any]]:
    """'countgd_sam2_instance_segmentation' is a tool that can detect multiple
    instances of an object given a text prompt. It is particularly useful when trying
    to detect and count a large number of objects. You can optionally separate object
    names in the prompt with commas. It returns a list of bounding boxes with
    normalized coordinates, label names, masks associated confidence scores.

    Parameters:
        prompt (str): The object that needs to be counted.
        image (np.ndarray): The image that contains multiple instances of the object.
        box_threshold (float, optional): The threshold for detection. Defaults
            to 0.23.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label,
            bounding box, and mask of the detected objects with normalized coordinates
            (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left
            and xmax and ymax are the coordinates of the bottom-right of the bounding box.
            The mask is binary 2D numpy array where 1 indicates the object and 0 indicates
            the background.

    Example
    -------
        >>> countgd_sam2_instance_segmentation("flower", image)
        [
            {
                'score': 0.49,
                'label': 'flower',
                'bbox': [0.1, 0.11, 0.35, 0.4],
                'mask': array([[0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0],
                    ...,
                    [0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
            },
        ]
    """

    od_ret = _countgd_object_detection(prompt, image, box_threshold, image.shape[:2])
    seg_ret = _sam2(
        image, od_ret["return_data"], image.shape[:2], image_bytes=od_ret["files"][0][1]
    )

    _display_tool_trace(
        countgd_sam2_instance_segmentation.__name__,
        {
            "prompts": prompt,
            "confidence": box_threshold,
        },
        seg_ret["display_data"],
        seg_ret["files"],
    )

    return seg_ret["return_data"]  # type: ignore

countgd_sam2_video_tracking

countgd_sam2_video_tracking(
    prompt, frames, chunk_length=10
)

'countgd_sam2_video_tracking' is a tool that can track and segment multiple objects in a video given a text prompt such as category names or referring expressions. The categories in the text prompt are separated by commas. It returns a list of bounding boxes, label names, masks and associated probability scores and is useful for tracking and counting without duplicating counts.

PARAMETER	DESCRIPTION
prompt	The prompt to ground to the image. TYPE: str
frames	The list of frames to ground the prompt to. TYPE: List[ndarray]
chunk_length	The number of frames to re-run countgd to find new objects. TYPE: Optional[int] DEFAULT: 10

RETURNS

DESCRIPTION

List[List[Dict[str, Any]]]

List[List[Dict[str, Any]]]: A list of list of dictionaries containing the label, segmentation mask and bounding boxes. The outer list represents each frame and the inner list is the entities per frame. The detected objects have normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box. The mask is binary 2D numpy array where 1 indicates the object and 0 indicates the background. The label names are prefixed with their ID represent the total count.

Example

>>> countgd_sam2_video_tracking("car, dinosaur", frames)
[
    [
        {
            'label': '0: dinosaur',
            'bbox': [0.1, 0.11, 0.35, 0.4],
            'mask': array([[0, 0, 0, ..., 0, 0, 0],
                [0, 0, 0, ..., 0, 0, 0],
                ...,
                [0, 0, 0, ..., 0, 0, 0],
                [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
        },
    ],
    ...
]

Source code in vision_agent/tools/tools.py

def countgd_sam2_video_tracking(
    prompt: str,
    frames: List[np.ndarray],
    chunk_length: Optional[int] = 10,
) -> List[List[Dict[str, Any]]]:
    """'countgd_sam2_video_tracking' is a tool that can track and segment multiple
    objects in a video given a text prompt such as category names or referring
    expressions. The categories in the text prompt are separated by commas. It returns
    a list of bounding boxes, label names, masks and associated probability scores and
    is useful for tracking and counting without duplicating counts.

    Parameters:
        prompt (str): The prompt to ground to the image.
        frames (List[np.ndarray]): The list of frames to ground the prompt to.
        chunk_length (Optional[int]): The number of frames to re-run countgd to find
            new objects.

    Returns:
        List[List[Dict[str, Any]]]: A list of list of dictionaries containing the
            label, segmentation mask and bounding boxes. The outer list represents each
            frame and the inner list is the entities per frame. The detected objects
            have normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin
            and ymin are the coordinates of the top-left and xmax and ymax are the
            coordinates of the bottom-right of the bounding box. The mask is binary 2D
            numpy array where 1 indicates the object and 0 indicates the background.
            The label names are prefixed with their ID represent the total count.

    Example
    -------
        >>> countgd_sam2_video_tracking("car, dinosaur", frames)
        [
            [
                {
                    'label': '0: dinosaur',
                    'bbox': [0.1, 0.11, 0.35, 0.4],
                    'mask': array([[0, 0, 0, ..., 0, 0, 0],
                        [0, 0, 0, ..., 0, 0, 0],
                        ...,
                        [0, 0, 0, ..., 0, 0, 0],
                        [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
                },
            ],
            ...
        ]
    """

    ret = od_sam2_video_tracking(
        ODModels.COUNTGD, prompt=prompt, frames=frames, chunk_length=chunk_length
    )
    _display_tool_trace(
        countgd_sam2_video_tracking.__name__,
        {},
        ret["display_data"],
        ret["files"],
    )
    return ret["return_data"]  # type: ignore

countgd_visual_prompt_object_detection

countgd_visual_prompt_object_detection(
    visual_prompts, image, box_threshold=0.23
)

'countgd_visual_prompt_object_detection' is a tool that can precisely count multiple instances of an object given few visual example prompts. It returns a list of bounding boxes with normalized coordinates, label names and associated confidence scores.

PARAMETER	DESCRIPTION
visual_prompts	Bounding boxes of the object in format [xmin, ymin, xmax, ymax]. Upto 3 bounding boxes can be provided. image (np.ndarray): The image that contains multiple instances of the object. box_threshold (float, optional): The threshold for detection. Defaults to 0.23. TYPE: List[List[float]]

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label, and bounding box of the detected objects with normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box.

Example

>>> countgd_visual_object_detection(
    visual_prompts=[[0.1, 0.1, 0.4, 0.42], [0.2, 0.3, 0.25, 0.35]],
    image=image
)
[
    {'score': 0.49, 'label': 'object', 'bounding_box': [0.1, 0.11, 0.35, 0.4]},
    {'score': 0.68, 'label': 'object', 'bounding_box': [0.2, 0.21, 0.45, 0.5},
    {'score': 0.78, 'label': 'object', 'bounding_box': [0.3, 0.35, 0.48, 0.52},
    {'score': 0.98, 'label': 'object', 'bounding_box': [0.44, 0.24, 0.49, 0.58},
]

Source code in vision_agent/tools/tools.py

def countgd_visual_prompt_object_detection(
    visual_prompts: List[List[float]],
    image: np.ndarray,
    box_threshold: float = 0.23,
) -> List[Dict[str, Any]]:
    """'countgd_visual_prompt_object_detection' is a tool that can precisely count
    multiple instances of an object given few visual example prompts. It returns a list
    of bounding boxes with normalized coordinates, label names and associated
    confidence scores.

    Parameters:
        visual_prompts (List[List[float]]): Bounding boxes of the object in format
            [xmin, ymin, xmax, ymax]. Upto 3 bounding boxes can be provided. image
            (np.ndarray): The image that contains multiple instances of the object.
            box_threshold (float, optional): The threshold for detection. Defaults to
            0.23.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label, and
            bounding box of the detected objects with normalized coordinates between 0
            and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the
            top-left and xmax and ymax are the coordinates of the bottom-right of the
            bounding box.

    Example
    -------
        >>> countgd_visual_object_detection(
            visual_prompts=[[0.1, 0.1, 0.4, 0.42], [0.2, 0.3, 0.25, 0.35]],
            image=image
        )
        [
            {'score': 0.49, 'label': 'object', 'bounding_box': [0.1, 0.11, 0.35, 0.4]},
            {'score': 0.68, 'label': 'object', 'bounding_box': [0.2, 0.21, 0.45, 0.5},
            {'score': 0.78, 'label': 'object', 'bounding_box': [0.3, 0.35, 0.48, 0.52},
            {'score': 0.98, 'label': 'object', 'bounding_box': [0.44, 0.24, 0.49, 0.58},
        ]
    """
    image_size = image.shape[:2]
    if image_size[0] < 1 or image_size[1] < 1:
        return []

    buffer_bytes = numpy_to_bytes(image)
    files = [("image", buffer_bytes)]
    visual_prompts = [
        denormalize_bbox(bbox, image.shape[:2]) for bbox in visual_prompts
    ]
    payload = {"visual_prompts": json.dumps(visual_prompts), "model": "countgd"}
    metadata = {"function_name": "countgd_visual_prompt_object_detection"}

    detections = send_task_inference_request(
        payload, "visual-prompts-to-object-detection", files=files, metadata=metadata
    )

    # get the first frame
    bboxes_per_frame = detections[0]
    bboxes_formatted = [
        {
            "label": bbox["label"],
            "bbox": normalize_bbox(bbox["bounding_box"], image_size),
            "score": round(bbox["score"], 2),
        }
        for bbox in bboxes_per_frame
    ]
    _display_tool_trace(
        countgd_visual_prompt_object_detection.__name__,
        payload,
        [
            {
                "label": e["label"],
                "score": e["score"],
                "bbox": denormalize_bbox(e["bbox"], image_size),
            }
            for e in bboxes_formatted
        ],
        files,
    )

    return bboxes_formatted

custom_object_detection

custom_object_detection(
    deployment_id, image, box_threshold=0.1
)

'custom_object_detection' is a tool that can detect instances of an object given a deployment_id of a previously finetuned object detection model. It is particularly useful when trying to detect objects that are not well detected by generalist models. It returns a list of bounding boxes with normalized coordinates, label names and associated confidence scores.

PARAMETER	DESCRIPTION
deployment_id	The id of the finetuned model. TYPE: str
image	The image that contains instances of the object. TYPE: ndarray
box_threshold	The threshold for detection. Defaults to 0.1. TYPE: float DEFAULT: 0.1

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label, and bounding box of the detected objects with normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box.

Example

>>> custom_object_detection("abcd1234-5678efg", image)
[
    {'score': 0.49, 'label': 'flower', 'bbox': [0.1, 0.11, 0.35, 0.4]},
    {'score': 0.68, 'label': 'flower', 'bbox': [0.2, 0.21, 0.45, 0.5]},
    {'score': 0.78, 'label': 'flower', 'bbox': [0.3, 0.35, 0.48, 0.52]},
    {'score': 0.98, 'label': 'flower', 'bbox': [0.44, 0.24, 0.49, 0.58]},
]

Source code in vision_agent/tools/tools.py

def custom_object_detection(
    deployment_id: str,
    image: np.ndarray,
    box_threshold: float = 0.1,
) -> List[Dict[str, Any]]:
    """'custom_object_detection' is a tool that can detect instances of an
    object given a deployment_id of a previously finetuned object detection model.
    It is particularly useful when trying to detect objects that are not well detected by generalist models.
    It returns a list of bounding boxes with normalized
    coordinates, label names and associated confidence scores.

    Parameters:
        deployment_id (str): The id of the finetuned model.
        image (np.ndarray): The image that contains instances of the object.
        box_threshold (float, optional): The threshold for detection. Defaults
            to 0.1.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label, and
            bounding box of the detected objects with normalized coordinates between 0
            and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the
            top-left and xmax and ymax are the coordinates of the bottom-right of the
            bounding box.

    Example
    -------
        >>> custom_object_detection("abcd1234-5678efg", image)
        [
            {'score': 0.49, 'label': 'flower', 'bbox': [0.1, 0.11, 0.35, 0.4]},
            {'score': 0.68, 'label': 'flower', 'bbox': [0.2, 0.21, 0.45, 0.5]},
            {'score': 0.78, 'label': 'flower', 'bbox': [0.3, 0.35, 0.48, 0.52]},
            {'score': 0.98, 'label': 'flower', 'bbox': [0.44, 0.24, 0.49, 0.58]},
        ]
    """
    image_size = image.shape[:2]
    if image_size[0] < 1 or image_size[1] < 1:
        return []

    files = [("image", numpy_to_bytes(image))]
    payload = {
        "deployment_id": deployment_id,
        "confidence": box_threshold,
    }
    detections: List[List[Dict[str, Any]]] = send_inference_request(
        payload, "custom-object-detection", files=files, v2=True
    )

    bboxes = detections[0]
    bboxes_formatted = [
        {
            "label": bbox["label"],
            "bbox": normalize_bbox(bbox["bounding_box"], image_size),
            "score": bbox["score"],
        }
        for bbox in bboxes
    ]
    display_data = [
        {
            "label": bbox["label"],
            "bbox": bbox["bounding_box"],
            "score": bbox["score"],
        }
        for bbox in bboxes
    ]

    _display_tool_trace(
        custom_object_detection.__name__,
        payload,
        display_data,
        files,
    )
    return bboxes_formatted

custom_od_sam2_video_tracking

custom_od_sam2_video_tracking(
    deployment_id, frames, chunk_length=10
)

'custom_od_sam2_video_tracking' is a tool that can segment multiple objects given a custom model with predefined category names. It returns a list of bounding boxes, label names, mask file names and associated probability scores.

PARAMETER	DESCRIPTION
deployment_id	The id of the deployed custom model. TYPE: str
image	The image to ground the prompt to. TYPE: ndarray
chunk_length	The number of frames to re-run florence2 to find new objects. TYPE: Optional[int] DEFAULT: 10

RETURNS	DESCRIPTION
List[List[Dict[str, Any]]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label, bounding box, and mask of the detected objects with normalized coordinates (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box. The mask is binary 2D numpy array where 1 indicates the object and 0 indicates the background.

Example

>>> custom_od_sam2_video_tracking("abcd1234-5678efg", frames)
[
    [
        {
            'label': '0: dinosaur',
            'bbox': [0.1, 0.11, 0.35, 0.4],
            'mask': array([[0, 0, 0, ..., 0, 0, 0],
                [0, 0, 0, ..., 0, 0, 0],
                ...,
                [0, 0, 0, ..., 0, 0, 0],
                [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
        },
    ],
    ...
]

Source code in vision_agent/tools/tools.py

def custom_od_sam2_video_tracking(
    deployment_id: str,
    frames: List[np.ndarray],
    chunk_length: Optional[int] = 10,
) -> List[List[Dict[str, Any]]]:
    """'custom_od_sam2_video_tracking' is a tool that can segment multiple objects given a
    custom model with predefined category names.
    It returns a list of bounding boxes, label names,
    mask file names and associated probability scores.

    Parameters:
        deployment_id (str): The id of the deployed custom model.
        image (np.ndarray): The image to ground the prompt to.
        chunk_length (Optional[int]): The number of frames to re-run florence2 to find
            new objects.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label,
            bounding box, and mask of the detected objects with normalized coordinates
            (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left
            and xmax and ymax are the coordinates of the bottom-right of the bounding box.
            The mask is binary 2D numpy array where 1 indicates the object and 0 indicates
            the background.

    Example
    -------
        >>> custom_od_sam2_video_tracking("abcd1234-5678efg", frames)
        [
            [
                {
                    'label': '0: dinosaur',
                    'bbox': [0.1, 0.11, 0.35, 0.4],
                    'mask': array([[0, 0, 0, ..., 0, 0, 0],
                        [0, 0, 0, ..., 0, 0, 0],
                        ...,
                        [0, 0, 0, ..., 0, 0, 0],
                        [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
                },
            ],
            ...
        ]
    """

    ret = od_sam2_video_tracking(
        ODModels.CUSTOM,
        prompt="",
        frames=frames,
        chunk_length=chunk_length,
        fine_tune_id=deployment_id,
    )
    _display_tool_trace(
        custom_od_sam2_video_tracking.__name__,
        {},
        ret["display_data"],
        ret["files"],
    )
    return ret["return_data"]  # type: ignore

qwen2_vl_images_vqa

qwen2_vl_images_vqa(prompt, images)

'qwen2_vl_images_vqa' is a tool that can answer any questions about arbitrary images including regular images or images of documents or presentations. It can be very useful for document QA or OCR text extraction. It returns text as an answer to the question.

PARAMETER	DESCRIPTION
prompt	The question about the document image TYPE: str
images	The reference images used for the question TYPE: List[ndarray]

RETURNS	DESCRIPTION
str	A string which is the answer to the given prompt. TYPE: str

Example

>>> qwen2_vl_images_vqa('Give a summary of the document', images)
'The document talks about the history of the United States of America and its...'

Source code in vision_agent/tools/tools.py

def qwen2_vl_images_vqa(prompt: str, images: List[np.ndarray]) -> str:
    """'qwen2_vl_images_vqa' is a tool that can answer any questions about arbitrary
    images including regular images or images of documents or presentations. It can be
    very useful for document QA or OCR text extraction. It returns text as an answer to
    the question.

    Parameters:
        prompt (str): The question about the document image
        images (List[np.ndarray]): The reference images used for the question

    Returns:
        str: A string which is the answer to the given prompt.

    Example
    -------
        >>> qwen2_vl_images_vqa('Give a summary of the document', images)
        'The document talks about the history of the United States of America and its...'
    """
    if isinstance(images, np.ndarray):
        images = [images]

    for image in images:
        if image.shape[0] < 1 or image.shape[1] < 1:
            raise ValueError(f"Image is empty, image shape: {image.shape}")

    files = [("images", numpy_to_bytes(image)) for image in images]
    payload = {
        "prompt": prompt,
        "model": "qwen2vl",
        "function_name": "qwen2_vl_images_vqa",
    }
    data: Dict[str, Any] = send_inference_request(
        payload, "image-to-text", files=files, v2=True
    )
    _display_tool_trace(
        qwen2_vl_images_vqa.__name__,
        payload,
        cast(str, data),
        files,
    )
    return cast(str, data)

qwen2_vl_video_vqa

qwen2_vl_video_vqa(prompt, frames)

'qwen2_vl_video_vqa' is a tool that can answer any questions about arbitrary videos including regular videos or videos of documents or presentations. It returns text as an answer to the question.

PARAMETER	DESCRIPTION
prompt	The question about the video TYPE: str
frames	The reference frames used for the question TYPE: List[ndarray]

RETURNS	DESCRIPTION
str	A string which is the answer to the given prompt. TYPE: str

Example

>>> qwen2_vl_video_vqa('Which football player made the goal?', frames)
'Lionel Messi'

Source code in vision_agent/tools/tools.py

def qwen2_vl_video_vqa(prompt: str, frames: List[np.ndarray]) -> str:
    """'qwen2_vl_video_vqa' is a tool that can answer any questions about arbitrary videos
    including regular videos or videos of documents or presentations. It returns text
    as an answer to the question.

    Parameters:
        prompt (str): The question about the video
        frames (List[np.ndarray]): The reference frames used for the question

    Returns:
        str: A string which is the answer to the given prompt.

    Example
    -------
        >>> qwen2_vl_video_vqa('Which football player made the goal?', frames)
        'Lionel Messi'
    """

    if len(frames) == 0 or not isinstance(frames, List):
        raise ValueError("Must provide a list of numpy arrays for frames")

    buffer_bytes = frames_to_bytes(frames)
    files = [("video", buffer_bytes)]
    payload = {
        "prompt": prompt,
        "model": "qwen2vl",
        "function_name": "qwen2_vl_video_vqa",
    }
    data: Dict[str, Any] = send_inference_request(
        payload, "image-to-text", files=files, v2=True
    )
    _display_tool_trace(
        qwen2_vl_video_vqa.__name__,
        payload,
        cast(str, data),
        files,
    )
    return cast(str, data)

ocr

ocr(image)

'ocr' extracts text from an image. It returns a list of detected text, bounding boxes with normalized coordinates, and confidence scores. The results are sorted from top-left to bottom right.

PARAMETER	DESCRIPTION
image	The image to extract text from. TYPE: ndarray

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the detected text, bbox with normalized coordinates, and confidence score.

Example

>>> ocr(image)
[
    {'label': 'hello world', 'bbox': [0.1, 0.11, 0.35, 0.4], 'score': 0.99},
]

Source code in vision_agent/tools/tools.py

def ocr(image: np.ndarray) -> List[Dict[str, Any]]:
    """'ocr' extracts text from an image. It returns a list of detected text, bounding
    boxes with normalized coordinates, and confidence scores. The results are sorted
    from top-left to bottom right.

    Parameters:
        image (np.ndarray): The image to extract text from.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the detected text, bbox
            with normalized coordinates, and confidence score.

    Example
    -------
        >>> ocr(image)
        [
            {'label': 'hello world', 'bbox': [0.1, 0.11, 0.35, 0.4], 'score': 0.99},
        ]
    """

    pil_image = Image.fromarray(image).convert("RGB")
    image_size = pil_image.size[::-1]
    if image_size[0] < 1 or image_size[1] < 1:
        return []
    image_buffer = io.BytesIO()
    pil_image.save(image_buffer, format="PNG")
    buffer_bytes = image_buffer.getvalue()
    image_buffer.close()

    res = requests.post(
        _OCR_URL,
        files={"images": buffer_bytes},
        data={"language": "en"},
        headers={"contentType": "multipart/form-data", "apikey": _API_KEY},
    )

    if res.status_code != 200:
        raise ValueError(f"OCR request failed with status code {res.status_code}")

    data = res.json()
    output = []
    for det in data[0]:
        label = det["text"]
        box = [
            det["location"][0]["x"],
            det["location"][0]["y"],
            det["location"][2]["x"],
            det["location"][2]["y"],
        ]
        box = normalize_bbox(box, image_size)
        output.append({"label": label, "bbox": box, "score": round(det["score"], 2)})

    _display_tool_trace(
        ocr.__name__,
        {},
        data,
        cast(List[Tuple[str, bytes]], [("image", buffer_bytes)]),
    )
    return sorted(output, key=lambda x: (x["bbox"][1], x["bbox"][0]))

claude35_text_extraction

claude35_text_extraction(image)

'claude35_text_extraction' is a tool that can extract text from an image. It returns the extracted text as a string and can be used as an alternative to OCR if you do not need to know the exact bounding box of the text.

PARAMETER	DESCRIPTION
image	The image to extract text from. TYPE: ndarray

RETURNS	DESCRIPTION
str	The extracted text from the image. TYPE: str

Source code in vision_agent/tools/tools.py

def claude35_text_extraction(image: np.ndarray) -> str:
    """'claude35_text_extraction' is a tool that can extract text from an image. It
    returns the extracted text as a string and can be used as an alternative to OCR if
    you do not need to know the exact bounding box of the text.

    Parameters:
        image (np.ndarray): The image to extract text from.

    Returns:
        str: The extracted text from the image.
    """

    lmm = AnthropicLMM()
    buffer = io.BytesIO()
    Image.fromarray(image).save(buffer, format="PNG")
    image_bytes = buffer.getvalue()
    image_b64 = "data:image/png;base64," + encode_image_bytes(image_bytes)
    text = lmm.generate(
        "Extract and return any text you see in this image and nothing else. If you do not read any text respond with an empty string.",
        [image_b64],
    )
    return cast(str, text)

document_extraction

document_extraction(image)

'document_extraction' is a tool that can extract structured information out of documents with different layouts. It returns the extracted data in a structured hierarchical format containing text, tables, pictures, charts, and other information.

PARAMETER	DESCRIPTION
image	The document image to analyze TYPE: ndarray

RETURNS	DESCRIPTION
Dict[str, Any]	Dict[str, Any]: A dictionary containing the extracted information.

Example

>>> document_analysis(image)
{'pages':
    [{'bbox': [0, 0, 1.0, 1.0],
            'chunks': [{'bbox': [0.8, 0.1, 1.0, 0.2],
                        'label': 'page_header',
                        'order': 75
                        'caption': 'Annual Report 2024',
                        'summary': 'This annual report summarizes ...' },
                       {'bbox': [0.2, 0.9, 0.9, 1.0],
                        'label': 'table',
                        'order': 1119,
                        'caption': [{'Column 1': 'Value 1', 'Column 2': 'Value 2'},
                        'summary': 'This table illustrates a trend of ...'},
            ],

Source code in vision_agent/tools/tools.py

def document_extraction(image: np.ndarray) -> Dict[str, Any]:
    """'document_extraction' is a tool that can extract structured information out of
    documents with different layouts. It returns the extracted data in a structured
    hierarchical format containing text, tables, pictures, charts, and other
    information.

    Parameters:
        image (np.ndarray): The document image to analyze

    Returns:
        Dict[str, Any]: A dictionary containing the extracted information.

    Example
    -------
        >>> document_analysis(image)
        {'pages':
            [{'bbox': [0, 0, 1.0, 1.0],
                    'chunks': [{'bbox': [0.8, 0.1, 1.0, 0.2],
                                'label': 'page_header',
                                'order': 75
                                'caption': 'Annual Report 2024',
                                'summary': 'This annual report summarizes ...' },
                               {'bbox': [0.2, 0.9, 0.9, 1.0],
                                'label': 'table',
                                'order': 1119,
                                'caption': [{'Column 1': 'Value 1', 'Column 2': 'Value 2'},
                                'summary': 'This table illustrates a trend of ...'},
                    ],
    """

    image_file = numpy_to_bytes(image)

    files = [("image", image_file)]

    payload = {
        "model": "document-analysis",
    }

    data: Dict[str, Any] = send_inference_request(
        payload=payload,
        endpoint_name="document-analysis",
        files=files,
        v2=True,
        metadata_payload={"function_name": "document_analysis"},
    )

    # don't display normalized bboxes
    _display_tool_trace(
        document_extraction.__name__,
        payload,
        data,
        files,
    )

    def normalize(data: Any) -> Dict[str, Any]:
        if isinstance(data, Dict):
            if "bbox" in data:
                data["bbox"] = normalize_bbox(data["bbox"], image.shape[:2])
            for key in data:
                data[key] = normalize(data[key])
        elif isinstance(data, List):
            for i in range(len(data)):
                data[i] = normalize(data[i])
        return data  # type: ignore

    data = normalize(data)

    return data

document_qa

document_qa(prompt, image)

'document_qa' is a tool that can answer any questions about arbitrary documents, presentations, or tables. It's very useful for document QA tasks, you can ask it a specific question or ask it to return a JSON object answering multiple questions about the document.

PARAMETER	DESCRIPTION
prompt	The question to be answered about the document image. TYPE: str
image	The document image to analyze. TYPE: ndarray

RETURNS	DESCRIPTION
str	The answer to the question based on the document's context. TYPE: str

Example

>>> document_qa(image, question)
'The answer to the question ...'

Source code in vision_agent/tools/tools.py

def document_qa(
    prompt: str,
    image: np.ndarray,
) -> str:
    """'document_qa' is a tool that can answer any questions about arbitrary documents,
    presentations, or tables. It's very useful for document QA tasks, you can ask it a
    specific question or ask it to return a JSON object answering multiple questions
    about the document.

    Parameters:
        prompt (str): The question to be answered about the document image.
        image (np.ndarray): The document image to analyze.

    Returns:
        str: The answer to the question based on the document's context.

    Example
    -------
        >>> document_qa(image, question)
        'The answer to the question ...'
    """

    image_file = numpy_to_bytes(image)

    files = [("image", image_file)]

    payload = {
        "model": "document-analysis",
    }

    data: Dict[str, Any] = send_inference_request(
        payload=payload,
        endpoint_name="document-analysis",
        files=files,
        v2=True,
        metadata_payload={"function_name": "document_qa"},
    )

    def normalize(data: Any) -> Dict[str, Any]:
        if isinstance(data, Dict):
            if "bbox" in data:
                data["bbox"] = normalize_bbox(data["bbox"], image.shape[:2])
            for key in data:
                data[key] = normalize(data[key])
        elif isinstance(data, List):
            for i in range(len(data)):
                data[i] = normalize(data[i])
        return data  # type: ignore

    data = normalize(data)

    prompt = f"""
Document Context:
{data}\n
Question: {prompt}\n
Answer the question directly using only the information from the document, do not answer with any additional text besides the answer. If the answer is not definitively contained in the document, say "I cannot find the answer in the provided document."
    """

    lmm = AnthropicLMM()
    llm_output = lmm.generate(prompt=prompt)
    llm_output = cast(str, llm_output)

    _display_tool_trace(
        document_qa.__name__,
        payload,
        llm_output,
        files,
    )

    return llm_output

video_temporal_localization

video_temporal_localization(
    prompt, frames, model="qwen2vl", chunk_length_frames=2
)

'video_temporal_localization' will run qwen2vl on each chunk_length_frames value selected for the video. It can detect multiple objects independently per chunk_length_frames given a text prompt such as a referring expression but does not track objects across frames. It returns a list of floats with a value of 1.0 if the objects are found in a given chunk_length_frames of the video.

PARAMETER	DESCRIPTION
prompt	The question about the video TYPE: str
frames	The reference frames used for the question TYPE: List[ndarray]
model	The model to use for the inference. Valid values are 'qwen2vl', 'gpt4o'. TYPE: str DEFAULT: 'qwen2vl'
chunk_length_frames	length of each chunk in frames TYPE: Optional[int] DEFAULT: 2

RETURNS	DESCRIPTION
List[float]	List[float]: A list of floats with a value of 1.0 if the objects to be found are present in the chunk_length_frames of the video.

Example

>>> video_temporal_localization('Did a goal happened?', frames)
[0.0, 0.0, 0.0, 1.0, 1.0, 0.0]

Source code in vision_agent/tools/tools.py

def video_temporal_localization(
    prompt: str,
    frames: List[np.ndarray],
    model: str = "qwen2vl",
    chunk_length_frames: Optional[int] = 2,
) -> List[float]:
    """'video_temporal_localization' will run qwen2vl on each chunk_length_frames
    value selected for the video. It can detect multiple objects independently per
    chunk_length_frames given a text prompt such as a referring expression
    but does not track objects across frames.
    It returns a list of floats with a value of 1.0 if the objects are found in a given
    chunk_length_frames of the video.

    Parameters:
        prompt (str): The question about the video
        frames (List[np.ndarray]): The reference frames used for the question
        model (str): The model to use for the inference. Valid values are
            'qwen2vl', 'gpt4o'.
        chunk_length_frames (Optional[int]): length of each chunk in frames

    Returns:
        List[float]: A list of floats with a value of 1.0 if the objects to be found
            are present in the chunk_length_frames of the video.

    Example
    -------
        >>> video_temporal_localization('Did a goal happened?', frames)
        [0.0, 0.0, 0.0, 1.0, 1.0, 0.0]
    """

    buffer_bytes = frames_to_bytes(frames)
    files = [("video", buffer_bytes)]
    payload: Dict[str, Any] = {
        "prompt": prompt,
        "model": model,
        "function_name": "video_temporal_localization",
    }
    if chunk_length_frames is not None:
        payload["chunk_length_frames"] = chunk_length_frames

    data = send_inference_request(
        payload, "video-temporal-localization", files=files, v2=True
    )
    _display_tool_trace(
        video_temporal_localization.__name__,
        payload,
        data,
        files,
    )
    return [cast(float, value) for value in data]

vit_image_classification

vit_image_classification(image)

'vit_image_classification' is a tool that can classify an image. It returns a list of classes and their probability scores based on image content.

PARAMETER	DESCRIPTION
image	The image to classify or tag TYPE: ndarray

RETURNS	DESCRIPTION
Dict[str, Any]	Dict[str, Any]: A dictionary containing the labels and scores. One dictionary contains a list of labels and other a list of scores.

Example

>>> vit_image_classification(image)
{"labels": ["leopard", "lemur, otter", "bird"], "scores": [0.68, 0.30, 0.02]},

Source code in vision_agent/tools/tools.py

def vit_image_classification(image: np.ndarray) -> Dict[str, Any]:
    """'vit_image_classification' is a tool that can classify an image. It returns a
    list of classes and their probability scores based on image content.

    Parameters:
        image (np.ndarray): The image to classify or tag

    Returns:
        Dict[str, Any]: A dictionary containing the labels and scores. One dictionary
            contains a list of labels and other a list of scores.

    Example
    -------
        >>> vit_image_classification(image)
        {"labels": ["leopard", "lemur, otter", "bird"], "scores": [0.68, 0.30, 0.02]},
    """
    if image.shape[0] < 1 or image.shape[1] < 1:
        return {"labels": [], "scores": []}

    image_b64 = convert_to_b64(image)
    data = {
        "image": image_b64,
        "tool": "image_classification",
        "function_name": "vit_image_classification",
    }
    resp_data: dict[str, Any] = send_inference_request(data, "tools")
    resp_data["scores"] = [round(prob, 4) for prob in resp_data["scores"]]
    _display_tool_trace(
        vit_image_classification.__name__,
        data,
        resp_data,
        image_b64,
    )
    return resp_data

vit_nsfw_classification

vit_nsfw_classification(image)

'vit_nsfw_classification' is a tool that can classify an image as 'nsfw' or 'normal'. It returns the predicted label and their probability scores based on image content.

PARAMETER	DESCRIPTION
image	The image to classify or tag TYPE: ndarray

RETURNS	DESCRIPTION
Dict[str, Any]	Dict[str, Any]: A dictionary containing the labels and scores. One dictionary contains a list of labels and other a list of scores.

Example

>>> vit_nsfw_classification(image)
{"label": "normal", "scores": 0.68},

Source code in vision_agent/tools/tools.py

def vit_nsfw_classification(image: np.ndarray) -> Dict[str, Any]:
    """'vit_nsfw_classification' is a tool that can classify an image as 'nsfw' or 'normal'.
    It returns the predicted label and their probability scores based on image content.

    Parameters:
        image (np.ndarray): The image to classify or tag

    Returns:
        Dict[str, Any]: A dictionary containing the labels and scores. One dictionary
            contains a list of labels and other a list of scores.

    Example
    -------
        >>> vit_nsfw_classification(image)
        {"label": "normal", "scores": 0.68},
    """
    if image.shape[0] < 1 or image.shape[1] < 1:
        raise ValueError(f"Image is empty, image shape: {image.shape}")

    image_b64 = convert_to_b64(image)
    data = {
        "image": image_b64,
        "function_name": "vit_nsfw_classification",
    }
    resp_data: dict[str, Any] = send_inference_request(
        data, "nsfw-classification", v2=True
    )
    resp_data["score"] = round(resp_data["score"], 4)
    _display_tool_trace(
        vit_nsfw_classification.__name__,
        data,
        resp_data,
        image_b64,
    )
    return resp_data

detr_segmentation

detr_segmentation(image)

'detr_segmentation' is a tool that can segment common objects in an image without any text prompt. It returns a list of detected objects as labels, their regions as masks and their scores.

PARAMETER	DESCRIPTION
image	The image used to segment things and objects TYPE: ndarray

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label and mask of the detected objects. The mask is binary 2D numpy array where 1 indicates the object and 0 indicates the background.

Example

>>> detr_segmentation(image)
[
    {
        'score': 0.45,
        'label': 'window',
        'mask': array([[0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0],
            ...,
            [0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
    },
    {
        'score': 0.70,
        'label': 'bird',
        'mask': array([[0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0],
            ...,
            [0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
    },
]

Source code in vision_agent/tools/tools.py

def detr_segmentation(image: np.ndarray) -> List[Dict[str, Any]]:
    """'detr_segmentation' is a tool that can segment common objects in an
    image without any text prompt. It returns a list of detected objects
    as labels, their regions as masks and their scores.

    Parameters:
        image (np.ndarray): The image used to segment things and objects

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label
            and mask of the detected objects. The mask is binary 2D numpy array where 1
            indicates the object and 0 indicates the background.

    Example
    -------
        >>> detr_segmentation(image)
        [
            {
                'score': 0.45,
                'label': 'window',
                'mask': array([[0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0],
                    ...,
                    [0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
            },
            {
                'score': 0.70,
                'label': 'bird',
                'mask': array([[0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0],
                    ...,
                    [0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
            },
        ]
    """
    if image.shape[0] < 1 or image.shape[1] < 1:
        return []
    image_b64 = convert_to_b64(image)
    data = {
        "image": image_b64,
        "tool": "panoptic_segmentation",
        "function_name": "detr_segmentation",
    }

    answer = send_inference_request(data, "tools")
    return_data = []

    for i in range(len(answer["scores"])):
        return_data.append(
            {
                "score": round(answer["scores"][i], 2),
                "label": answer["labels"][i],
                "mask": rle_decode(
                    mask_rle=answer["masks"][i], shape=answer["mask_shape"][0]
                ),
            }
        )
    _display_tool_trace(
        detr_segmentation.__name__,
        {},
        return_data,
        image_b64,
    )
    return return_data

depth_anything_v2

depth_anything_v2(image)

'depth_anything_v2' is a tool that runs depth_anythingv2 model to generate a depth image from a given RGB image. The returned depth image is monochrome and represents depth values as pixel intesities with pixel values ranging from 0 to 255.

PARAMETER	DESCRIPTION
image	The image to used to generate depth image TYPE: ndarray

RETURNS	DESCRIPTION
ndarray	np.ndarray: A grayscale depth image with pixel values ranging from 0 to 255.

Example

>>> depth_anything_v2(image)
array([[0, 0, 0, ..., 0, 0, 0],
        [0, 20, 24, ..., 0, 100, 103],
        ...,
        [10, 11, 15, ..., 202, 202, 205],
        [10, 10, 10, ..., 200, 200, 200]], dtype=uint8),

Source code in vision_agent/tools/tools.py

def depth_anything_v2(image: np.ndarray) -> np.ndarray:
    """'depth_anything_v2' is a tool that runs depth_anythingv2 model to generate a
    depth image from a given RGB image. The returned depth image is monochrome and
    represents depth values as pixel intesities with pixel values ranging from 0 to 255.

    Parameters:
        image (np.ndarray): The image to used to generate depth image

    Returns:
        np.ndarray: A grayscale depth image with pixel values ranging from 0 to 255.

    Example
    -------
        >>> depth_anything_v2(image)
        array([[0, 0, 0, ..., 0, 0, 0],
                [0, 20, 24, ..., 0, 100, 103],
                ...,
                [10, 11, 15, ..., 202, 202, 205],
                [10, 10, 10, ..., 200, 200, 200]], dtype=uint8),
    """
    if image.shape[0] < 1 or image.shape[1] < 1:
        raise ValueError(f"Image is empty, image shape: {image.shape}")

    image_b64 = convert_to_b64(image)
    data = {
        "image": image_b64,
        "function_name": "depth_anything_v2",
    }

    depth_map = send_inference_request(data, "depth-anything-v2", v2=True)
    depth_map_np = np.array(depth_map["map"])
    depth_map_np = (depth_map_np - depth_map_np.min()) / (
        depth_map_np.max() - depth_map_np.min()
    )
    depth_map_np = (255 * depth_map_np).astype(np.uint8)
    _display_tool_trace(
        depth_anything_v2.__name__,
        {},
        depth_map,
        image_b64,
    )
    return depth_map_np

generate_pose_image

generate_pose_image(image)

'generate_pose_image' is a tool that generates a open pose bone/stick image from a given RGB image. The returned bone image is RGB with the pose amd keypoints colored and background as black.

PARAMETER	DESCRIPTION
image	The image to used to generate pose image TYPE: ndarray

RETURNS	DESCRIPTION
ndarray	np.ndarray: A bone or pose image indicating the pose and keypoints

Example

>>> generate_pose_image(image)
array([[0, 0, 0, ..., 0, 0, 0],
        [0, 20, 24, ..., 0, 100, 103],
        ...,
        [10, 11, 15, ..., 202, 202, 205],
        [10, 10, 10, ..., 200, 200, 200]], dtype=uint8),

Source code in vision_agent/tools/tools.py

def generate_pose_image(image: np.ndarray) -> np.ndarray:
    """'generate_pose_image' is a tool that generates a open pose bone/stick image from
    a given RGB image. The returned bone image is RGB with the pose amd keypoints colored
    and background as black.

    Parameters:
        image (np.ndarray): The image to used to generate pose image

    Returns:
        np.ndarray: A bone or pose image indicating the pose and keypoints

    Example
    -------
        >>> generate_pose_image(image)
        array([[0, 0, 0, ..., 0, 0, 0],
                [0, 20, 24, ..., 0, 100, 103],
                ...,
                [10, 11, 15, ..., 202, 202, 205],
                [10, 10, 10, ..., 200, 200, 200]], dtype=uint8),
    """
    image_b64 = convert_to_b64(image)
    data = {
        "image": image_b64,
        "function_name": "generate_pose_image",
    }

    pos_img = send_inference_request(data, "pose-detector", v2=True)
    return_data = np.array(b64_to_pil(pos_img["data"]).convert("RGB"))
    _display_tool_trace(
        generate_pose_image.__name__,
        {},
        pos_img,
        image_b64,
    )
    return return_data

template_match

template_match(image, template_image)

'template_match' is a tool that can detect all instances of a template in a given image. It returns the locations of the detected template, a corresponding similarity score of the same

PARAMETER	DESCRIPTION
image	The image used for searching the template TYPE: ndarray
template_image	The template image or crop to search in the image TYPE: ndarray

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score and bounding box of the detected template with normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box.

Example

>>> template_match(image, template)
[
    {'score': 0.79, 'bbox': [0.1, 0.11, 0.35, 0.4]},
    {'score': 0.38, 'bbox': [0.2, 0.21, 0.45, 0.5},
]

Source code in vision_agent/tools/tools.py

def template_match(
    image: np.ndarray, template_image: np.ndarray
) -> List[Dict[str, Any]]:
    """'template_match' is a tool that can detect all instances of a template in
    a given image. It returns the locations of the detected template, a corresponding
    similarity score of the same

    Parameters:
        image (np.ndarray): The image used for searching the template
        template_image (np.ndarray): The template image or crop to search in the image

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score and
            bounding box of the detected template with normalized coordinates between 0
            and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the
            top-left and xmax and ymax are the coordinates of the bottom-right of the
            bounding box.

    Example
    -------
        >>> template_match(image, template)
        [
            {'score': 0.79, 'bbox': [0.1, 0.11, 0.35, 0.4]},
            {'score': 0.38, 'bbox': [0.2, 0.21, 0.45, 0.5},
        ]
    """
    image_size = image.shape[:2]
    image_b64 = convert_to_b64(image)
    template_image_b64 = convert_to_b64(template_image)
    data = {
        "image": image_b64,
        "template": template_image_b64,
        "tool": "template_match",
        "function_name": "template_match",
    }

    answer = send_inference_request(data, "tools")
    return_data = []
    for i in range(len(answer["bboxes"])):
        return_data.append(
            {
                "label": "match",
                "score": round(answer["scores"][i], 2),
                "bbox": normalize_bbox(answer["bboxes"][i], image_size),
            }
        )
    _display_tool_trace(
        template_match.__name__,
        {"template_image": template_image_b64},
        return_data,
        image_b64,
    )
    return return_data

flux_image_inpainting

flux_image_inpainting(prompt, image, mask)

'flux_image_inpainting' performs image inpainting to fill the masked regions, given by mask, in the image, given image based on the text prompt and surrounding image context. It can be used to edit regions of an image according to the prompt given.

PARAMETER	DESCRIPTION
prompt	A detailed text description guiding what should be generated in the masked area. More detailed and specific prompts typically yield better results. TYPE: str
image	The source image to be inpainted. The image will serve as the base context for the inpainting process. TYPE: ndarray
mask	A binary mask image with 0's and 1's, where 1 indicates areas to be inpainted and 0 indicates areas to be preserved. TYPE: ndarray

RETURNS	DESCRIPTION
ndarray	np.ndarray: The generated image(s) as a numpy array in RGB format with values ranging from 0 to 255.

---

Example: >>> # Generate inpainting >>> result = flux_image_inpainting( ... prompt="a modern black leather sofa with white pillows", ... image=image, ... mask=mask, ... ) >>> save_image(result, "inpainted_room.png")

Source code in vision_agent/tools/tools.py

def flux_image_inpainting(
    prompt: str,
    image: np.ndarray,
    mask: np.ndarray,
) -> np.ndarray:
    """'flux_image_inpainting' performs image inpainting to fill the masked regions,
    given by mask, in the image, given image based on the text prompt and surrounding image context.
    It can be used to edit regions of an image according to the prompt given.

    Parameters:
        prompt (str): A detailed text description guiding what should be generated
            in the masked area. More detailed and specific prompts typically yield better results.
        image (np.ndarray): The source image to be inpainted.
            The image will serve as the base context for the inpainting process.
        mask (np.ndarray): A binary mask image with 0's and 1's,
            where 1 indicates areas to be inpainted and 0 indicates areas to be preserved.

    Returns:
        np.ndarray: The generated image(s) as a numpy array in RGB format with values
            ranging from 0 to 255.

    -------
    Example:
        >>> # Generate inpainting
        >>> result = flux_image_inpainting(
        ...     prompt="a modern black leather sofa with white pillows",
        ...     image=image,
        ...     mask=mask,
        ... )
        >>> save_image(result, "inpainted_room.png")
    """

    min_dim = 8

    if any(dim < min_dim for dim in image.shape[:2] + mask.shape[:2]):
        raise ValueError(f"Image and mask must be at least {min_dim}x{min_dim} pixels")

    max_size = (512, 512)

    if image.shape[0] > max_size[0] or image.shape[1] > max_size[1]:
        scaling_factor = min(max_size[0] / image.shape[0], max_size[1] / image.shape[1])
        new_size = (
            int(image.shape[1] * scaling_factor),
            int(image.shape[0] * scaling_factor),
        )
        new_size = ((new_size[0] // 8) * 8, (new_size[1] // 8) * 8)
        image = cv2.resize(image, new_size, interpolation=cv2.INTER_AREA)
        mask = cv2.resize(mask, new_size, interpolation=cv2.INTER_NEAREST)

    elif image.shape[0] % 8 != 0 or image.shape[1] % 8 != 0:
        new_size = ((image.shape[1] // 8) * 8, (image.shape[0] // 8) * 8)
        image = cv2.resize(image, new_size, interpolation=cv2.INTER_AREA)
        mask = cv2.resize(mask, new_size, interpolation=cv2.INTER_NEAREST)

    if np.array_equal(mask, mask.astype(bool).astype(int)):
        mask = np.where(mask > 0, 255, 0).astype(np.uint8)
    else:
        raise ValueError("Mask should contain only binary values (0 or 1)")

    image_file = numpy_to_bytes(image)
    mask_file = numpy_to_bytes(mask)

    files = [
        ("image", image_file),
        ("mask_image", mask_file),
    ]

    payload = {
        "prompt": prompt,
        "task": "inpainting",
        "height": image.shape[0],
        "width": image.shape[1],
        "strength": 0.99,
        "guidance_scale": 18,
        "num_inference_steps": 20,
        "seed": None,
    }

    response = send_inference_request(
        payload=payload,
        endpoint_name="flux1",
        files=files,
        v2=True,
        metadata_payload={"function_name": "flux_image_inpainting"},
    )

    output_image = np.array(b64_to_pil(response[0]).convert("RGB"))
    _display_tool_trace(
        flux_image_inpainting.__name__,
        payload,
        output_image,
        files,
    )
    return output_image

siglip_classification

siglip_classification(image, labels)

'siglip_classification' is a tool that can classify an image or a cropped detection given a list of input labels or tags. It returns the same list of the input labels along with their probability scores based on image content.

PARAMETER	DESCRIPTION
image	The image to classify or tag TYPE: ndarray
labels	The list of labels or tags that is associated with the image TYPE: List[str]

RETURNS	DESCRIPTION
Dict[str, Any]	Dict[str, Any]: A dictionary containing the labels and scores. One dictionary contains a list of given labels and other a list of scores.

Example

>>> siglip_classification(image, ['dog', 'cat', 'bird'])
{"labels": ["dog", "cat", "bird"], "scores": [0.68, 0.30, 0.02]},

Source code in vision_agent/tools/tools.py

def siglip_classification(image: np.ndarray, labels: List[str]) -> Dict[str, Any]:
    """'siglip_classification' is a tool that can classify an image or a cropped detection given a list
    of input labels or tags. It returns the same list of the input labels along with
    their probability scores based on image content.

    Parameters:
        image (np.ndarray): The image to classify or tag
        labels (List[str]): The list of labels or tags that is associated with the image

    Returns:
        Dict[str, Any]: A dictionary containing the labels and scores. One dictionary
            contains a list of given labels and other a list of scores.

    Example
    -------
        >>> siglip_classification(image, ['dog', 'cat', 'bird'])
        {"labels": ["dog", "cat", "bird"], "scores": [0.68, 0.30, 0.02]},
    """

    if image.shape[0] < 1 or image.shape[1] < 1:
        return {"labels": [], "scores": []}

    image_file = numpy_to_bytes(image)

    files = [("image", image_file)]

    payload = {
        "model": "siglip",
        "labels": labels,
    }

    response: dict[str, Any] = send_inference_request(
        payload=payload,
        endpoint_name="classification",
        files=files,
        v2=True,
        metadata_payload={"function_name": "siglip_classification"},
    )

    _display_tool_trace(
        siglip_classification.__name__,
        payload,
        response,
        files,
    )
    return response

agentic_object_detection

agentic_object_detection(prompt, image, fine_tune_id=None)

'agentic_object_detection' is a tool that can detect and count multiple objects given a text prompt such as category names or referring expressions on images. The categories in text prompt are separated by commas. It returns a list of bounding boxes with normalized coordinates, label names and associated probability scores.

PARAMETER	DESCRIPTION
prompt	The prompt to ground to the image. TYPE: str
image	The image to ground the prompt to. TYPE: ndarray
fine_tune_id	If you have a fine-tuned model, you can pass the fine-tuned model ID here to use it. TYPE: Optional[str] DEFAULT: None

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label, and bounding box of the detected objects with normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box.

Example

>>> agentic_object_detection("car", image)
[
    {'score': 0.99, 'label': 'car', 'bbox': [0.1, 0.11, 0.35, 0.4]},
    {'score': 0.98, 'label': 'car', 'bbox': [0.2, 0.21, 0.45, 0.5},
]

Source code in vision_agent/tools/tools.py

def agentic_object_detection(
    prompt: str,
    image: np.ndarray,
    fine_tune_id: Optional[str] = None,
) -> List[Dict[str, Any]]:
    """'agentic_object_detection' is a tool that can detect and count multiple objects
    given a text prompt such as category names or referring expressions on images. The
    categories in text prompt are separated by commas. It returns a list of bounding
    boxes with normalized coordinates, label names and associated probability scores.

    Parameters:
        prompt (str): The prompt to ground to the image.
        image (np.ndarray): The image to ground the prompt to.
        fine_tune_id (Optional[str]): If you have a fine-tuned model, you can pass the
            fine-tuned model ID here to use it.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label, and
            bounding box of the detected objects with normalized coordinates between 0
            and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the
            top-left and xmax and ymax are the coordinates of the bottom-right of the
            bounding box.

    Example
    -------
        >>> agentic_object_detection("car", image)
        [
            {'score': 0.99, 'label': 'car', 'bbox': [0.1, 0.11, 0.35, 0.4]},
            {'score': 0.98, 'label': 'car', 'bbox': [0.2, 0.21, 0.45, 0.5},
        ]
    """

    image_size = image.shape[:2]
    if image_size[0] < 1 or image_size[1] < 1:
        return []

    ret = _agentic_object_detection(
        prompt, image, image_size, fine_tune_id=fine_tune_id
    )

    _display_tool_trace(
        agentic_object_detection.__name__,
        {"prompts": prompt},
        ret["display_data"],
        ret["files"],
    )
    return ret["return_data"]  # type: ignore

agentic_sam2_instance_segmentation

agentic_sam2_instance_segmentation(prompt, image)

'agentic_sam2_instance_segmentation' is a tool that can detect and count multiple instances of objects given a text prompt such as category names or referring expressions on images. The categories in text prompt are separated by commas. It returns a list of bounding boxes with normalized coordinates, label names, masks and associated probability scores.

PARAMETER	DESCRIPTION
prompt	The object that needs to be counted. TYPE: str
image	The image that contains multiple instances of the object. TYPE: ndarray

RETURNS	DESCRIPTION
List[Dict[str, Any]]	List[Dict[str, Any]]: A list of dictionaries containing the score, label, bounding box, and mask of the detected objects with normalized coordinates (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box. The mask is binary 2D numpy array where 1 indicates the object and 0 indicates the background.

Example

>>> agentic_sam2_instance_segmentation("flower", image)
[
    {
        'score': 0.49,
        'label': 'flower',
        'bbox': [0.1, 0.11, 0.35, 0.4],
        'mask': array([[0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0],
            ...,
            [0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
    },
]

Source code in vision_agent/tools/tools.py

def agentic_sam2_instance_segmentation(
    prompt: str, image: np.ndarray
) -> List[Dict[str, Any]]:
    """'agentic_sam2_instance_segmentation' is a tool that can detect and count multiple
    instances of objects given a text prompt such as category names or referring
    expressions on images. The categories in text prompt are separated by commas. It
    returns a list of bounding boxes with normalized coordinates, label names, masks
    and associated probability scores.

    Parameters:
        prompt (str): The object that needs to be counted.
        image (np.ndarray): The image that contains multiple instances of the object.

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing the score, label,
            bounding box, and mask of the detected objects with normalized coordinates
            (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left
            and xmax and ymax are the coordinates of the bottom-right of the bounding box.
            The mask is binary 2D numpy array where 1 indicates the object and 0 indicates
            the background.

    Example
    -------
        >>> agentic_sam2_instance_segmentation("flower", image)
        [
            {
                'score': 0.49,
                'label': 'flower',
                'bbox': [0.1, 0.11, 0.35, 0.4],
                'mask': array([[0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0],
                    ...,
                    [0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
            },
        ]
    """

    od_ret = _agentic_object_detection(prompt, image, image.shape[:2])
    seg_ret = _sam2(
        image, od_ret["return_data"], image.shape[:2], image_bytes=od_ret["files"][0][1]
    )

    _display_tool_trace(
        agentic_sam2_instance_segmentation.__name__,
        {
            "prompts": prompt,
        },
        seg_ret["display_data"],
        seg_ret["files"],
    )

    return seg_ret["return_data"]  # type: ignore

agentic_sam2_video_tracking

agentic_sam2_video_tracking(
    prompt, frames, chunk_length=10, fine_tune_id=None
)

'agentic_sam2_video_tracking' is a tool that can track and segment multiple objects in a video given a text prompt such as category names or referring expressions. The categories in the text prompt are separated by commas. It returns a list of bounding boxes, label names, masks and associated probability scores and is useful for tracking and counting without duplicating counts.

PARAMETER	DESCRIPTION
prompt	The prompt to ground to the image. TYPE: str
frames	The list of frames to ground the prompt to. TYPE: List[ndarray]
chunk_length	The number of frames to re-run agentic object detection to to find new objects. TYPE: Optional[int] DEFAULT: 10
fine_tune_id	If you have a fine-tuned model, you can pass the fine-tuned model ID here to use it. TYPE: Optional[str] DEFAULT: None

RETURNS

DESCRIPTION

List[List[Dict[str, Any]]]

List[List[Dict[str, Any]]]: A list of list of dictionaries containing the label, segmentation mask and bounding boxes. The outer list represents each frame and the inner list is the entities per frame. The detected objects have normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin and ymin are the coordinates of the top-left and xmax and ymax are the coordinates of the bottom-right of the bounding box. The mask is binary 2D numpy array where 1 indicates the object and 0 indicates the background. The label names are prefixed with their ID represent the total count.

Example

>>> agentic_sam2_video_tracking("dinosaur", frames)
[
    [
        {
            'label': '0: dinosaur',
            'bbox': [0.1, 0.11, 0.35, 0.4],
            'mask': array([[0, 0, 0, ..., 0, 0, 0],
                [0, 0, 0, ..., 0, 0, 0],
                ...,
                [0, 0, 0, ..., 0, 0, 0],
                [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
        },
    ],
    ...
]

Source code in vision_agent/tools/tools.py

def agentic_sam2_video_tracking(
    prompt: str,
    frames: List[np.ndarray],
    chunk_length: Optional[int] = 10,
    fine_tune_id: Optional[str] = None,
) -> List[List[Dict[str, Any]]]:
    """'agentic_sam2_video_tracking' is a tool that can track and segment multiple
    objects in a video given a text prompt such as category names or referring
    expressions. The categories in the text prompt are separated by commas. It returns
    a list of bounding boxes, label names, masks and associated probability scores and
    is useful for tracking and counting without duplicating counts.

    Parameters:
        prompt (str): The prompt to ground to the image.
        frames (List[np.ndarray]): The list of frames to ground the prompt to.
        chunk_length (Optional[int]): The number of frames to re-run agentic object detection to
            to find new objects.
        fine_tune_id (Optional[str]): If you have a fine-tuned model, you can pass the
            fine-tuned model ID here to use it.

    Returns:
        List[List[Dict[str, Any]]]: A list of list of dictionaries containing the
            label, segmentation mask and bounding boxes. The outer list represents each
            frame and the inner list is the entities per frame. The detected objects
            have normalized coordinates between 0 and 1 (xmin, ymin, xmax, ymax). xmin
            and ymin are the coordinates of the top-left and xmax and ymax are the
            coordinates of the bottom-right of the bounding box. The mask is binary 2D
            numpy array where 1 indicates the object and 0 indicates the background.
            The label names are prefixed with their ID represent the total count.

    Example
    -------
        >>> agentic_sam2_video_tracking("dinosaur", frames)
        [
            [
                {
                    'label': '0: dinosaur',
                    'bbox': [0.1, 0.11, 0.35, 0.4],
                    'mask': array([[0, 0, 0, ..., 0, 0, 0],
                        [0, 0, 0, ..., 0, 0, 0],
                        ...,
                        [0, 0, 0, ..., 0, 0, 0],
                        [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
                },
            ],
            ...
        ]
    """

    ret = od_sam2_video_tracking(
        ODModels.AGENTIC,
        prompt=prompt,
        frames=frames,
        chunk_length=chunk_length,
        fine_tune_id=fine_tune_id,
    )
    _display_tool_trace(
        agentic_sam2_video_tracking.__name__,
        {},
        ret["display_data"],
        ret["files"],
    )
    return ret["return_data"]  # type: ignore

minimum_distance

minimum_distance(det1, det2, image_size)

'minimum_distance' calculates the minimum distance between two detections which can include bounding boxes and or masks. This will return the closest distance between the objects, not the distance between the centers of the objects.

PARAMETER	DESCRIPTION
det1	The first detection of boxes or masks. TYPE: Dict[str, Any]
det2	The second detection of boxes or masks. TYPE: Dict[str, Any]
image_size	The size of the image given as (height, width). TYPE: Tuple[int, int]

RETURNS	DESCRIPTION
float	The closest distance between the two detections. TYPE: float

Example

>>> closest_distance(det1, det2, image_size)
141.42

Source code in vision_agent/tools/tools.py

def minimum_distance(
    det1: Dict[str, Any], det2: Dict[str, Any], image_size: Tuple[int, int]
) -> float:
    """'minimum_distance' calculates the minimum distance between two detections which
    can include bounding boxes and or masks. This will return the closest distance
    between the objects, not the distance between the centers of the objects.

    Parameters:
        det1 (Dict[str, Any]): The first detection of boxes or masks.
        det2 (Dict[str, Any]): The second detection of boxes or masks.
        image_size (Tuple[int, int]): The size of the image given as (height, width).

    Returns:
        float: The closest distance between the two detections.

    Example
    -------
        >>> closest_distance(det1, det2, image_size)
        141.42
    """

    if "mask" in det1 and "mask" in det2:
        return closest_mask_distance(det1["mask"], det2["mask"])
    elif "bbox" in det1 and "bbox" in det2:
        return closest_box_distance(det1["bbox"], det2["bbox"], image_size)
    else:
        raise ValueError("Both detections must have either bbox or mask")

closest_mask_distance

closest_mask_distance(mask1, mask2)

'closest_mask_distance' calculates the closest distance between two masks.

PARAMETER	DESCRIPTION
mask1	The first mask. TYPE: ndarray
mask2	The second mask. TYPE: ndarray

RETURNS	DESCRIPTION
float	The closest distance between the two masks. TYPE: float

Example

>>> closest_mask_distance(mask1, mask2)
0.5

Source code in vision_agent/tools/tools.py

def closest_mask_distance(mask1: np.ndarray, mask2: np.ndarray) -> float:
    """'closest_mask_distance' calculates the closest distance between two masks.

    Parameters:
        mask1 (np.ndarray): The first mask.
        mask2 (np.ndarray): The second mask.

    Returns:
        float: The closest distance between the two masks.

    Example
    -------
        >>> closest_mask_distance(mask1, mask2)
        0.5
    """

    mask1 = np.clip(mask1, 0, 1)
    mask2 = np.clip(mask2, 0, 1)
    contours1, _ = cv2.findContours(mask1, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    contours2, _ = cv2.findContours(mask2, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    largest_contour1 = max(contours1, key=cv2.contourArea)
    largest_contour2 = max(contours2, key=cv2.contourArea)
    polygon1 = cv2.approxPolyDP(largest_contour1, 1.0, True)
    polygon2 = cv2.approxPolyDP(largest_contour2, 1.0, True)
    min_distance = np.inf

    small_polygon, larger_contour = (
        (polygon1, largest_contour2)
        if len(largest_contour1) < len(largest_contour2)
        else (polygon2, largest_contour1)
    )

    # For each point in the first polygon
    for point in small_polygon:
        # Calculate the distance to the second polygon, -1 is to invert result as point inside the polygon is positive

        distance = (
            cv2.pointPolygonTest(
                larger_contour, (point[0, 0].item(), point[0, 1].item()), True
            )
            * -1
        )

        # If the distance is negative, the point is inside the polygon, so the distance is 0
        if distance < 0:
            continue
        else:
            # Update the minimum distance if the point is outside the polygon
            min_distance = min(min_distance, distance)

    return min_distance if min_distance != np.inf else 0.0

closest_box_distance

closest_box_distance(box1, box2, image_size)

'closest_box_distance' calculates the closest distance between two bounding boxes.

PARAMETER	DESCRIPTION
box1	The first bounding box. TYPE: List[float]
box2	The second bounding box. TYPE: List[float]
image_size	The size of the image given as (height, width). TYPE: Tuple[int, int]

RETURNS	DESCRIPTION
float	The closest distance between the two bounding boxes. TYPE: float

Example

>>> closest_box_distance([100, 100, 200, 200], [300, 300, 400, 400])
141.42

Source code in vision_agent/tools/tools.py

def closest_box_distance(
    box1: List[float], box2: List[float], image_size: Tuple[int, int]
) -> float:
    """'closest_box_distance' calculates the closest distance between two bounding boxes.

    Parameters:
        box1 (List[float]): The first bounding box.
        box2 (List[float]): The second bounding box.
        image_size (Tuple[int, int]): The size of the image given as (height, width).

    Returns:
        float: The closest distance between the two bounding boxes.

    Example
    -------
        >>> closest_box_distance([100, 100, 200, 200], [300, 300, 400, 400])
        141.42
    """

    x11, y11, x12, y12 = denormalize_bbox(box1, image_size)
    x21, y21, x22, y22 = denormalize_bbox(box2, image_size)

    horizontal_distance = np.max([0, x21 - x12, x11 - x22])
    vertical_distance = np.max([0, y21 - y12, y11 - y22])
    return cast(float, np.sqrt(horizontal_distance**2 + vertical_distance**2))

extract_frames_and_timestamps

extract_frames_and_timestamps(video_uri, fps=5)

'extract_frames_and_timestamps' extracts frames and timestamps from a video which can be a file path, url or youtube link, returns a list of dictionaries with keys "frame" and "timestamp" where "frame" is a numpy array and "timestamp" is the relative time in seconds where the frame was captured. The frame is a numpy array.

PARAMETER	DESCRIPTION
video_uri	The path to the video file, url or youtube link TYPE: Union[str, Path]
fps	The frame rate per second to extract the frames. Defaults to 5. TYPE: float DEFAULT: 5

RETURNS	DESCRIPTION
List[Dict[str, Union[ndarray, float]]]	List[Dict[str, Union[np.ndarray, float]]]: A list of dictionaries containing the extracted frame as a numpy array and the timestamp in seconds.

Example

>>> extract_frames("path/to/video.mp4")
[{"frame": np.ndarray, "timestamp": 0.0}, ...]

Source code in vision_agent/tools/tools.py

def extract_frames_and_timestamps(
    video_uri: Union[str, Path], fps: float = 5
) -> List[Dict[str, Union[np.ndarray, float]]]:
    """'extract_frames_and_timestamps' extracts frames and timestamps from a video
    which can be a file path, url or youtube link, returns a list of dictionaries
    with keys "frame" and "timestamp" where "frame" is a numpy array and "timestamp" is
    the relative time in seconds where the frame was captured. The frame is a numpy
    array.

    Parameters:
        video_uri (Union[str, Path]): The path to the video file, url or youtube link
        fps (float, optional): The frame rate per second to extract the frames. Defaults
            to 5.

    Returns:
        List[Dict[str, Union[np.ndarray, float]]]: A list of dictionaries containing the
            extracted frame as a numpy array and the timestamp in seconds.

    Example
    -------
        >>> extract_frames("path/to/video.mp4")
        [{"frame": np.ndarray, "timestamp": 0.0}, ...]
    """
    if isinstance(fps, str):
        # fps could be a string when it's passed in from a web endpoint deployment
        fps = float(fps)

    def reformat(
        frames_and_timestamps: List[Tuple[np.ndarray, float]],
    ) -> List[Dict[str, Union[np.ndarray, float]]]:
        return [
            {"frame": frame, "timestamp": timestamp}
            for frame, timestamp in frames_and_timestamps
        ]

    if str(video_uri).startswith(
        (
            "http://www.youtube.com/",
            "https://www.youtube.com/",
            "http://youtu.be/",
            "https://youtu.be/",
        )
    ):
        with tempfile.TemporaryDirectory() as temp_dir:
            yt = YouTube(str(video_uri))
            # Download the highest resolution video
            video = (
                yt.streams.filter(progressive=True, file_extension="mp4")
                .order_by("resolution")
                .desc()
                .first()
            )
            if not video:
                raise Exception("No suitable video stream found")
            video_file_path = video.download(output_path=temp_dir)

            return reformat(extract_frames_from_video(video_file_path, fps))
    elif str(video_uri).startswith(("http", "https")):
        _, image_suffix = os.path.splitext(video_uri)
        with tempfile.NamedTemporaryFile(delete=False, suffix=image_suffix) as tmp_file:
            # Download the video and save it to the temporary file
            with urllib.request.urlopen(str(video_uri)) as response:
                tmp_file.write(response.read())
            return reformat(extract_frames_from_video(tmp_file.name, fps))

    return reformat(extract_frames_from_video(str(video_uri), fps))

save_json

save_json(data, file_path)

'save_json' is a utility function that saves data as a JSON file. It is helpful for saving data that contains NumPy arrays which are not JSON serializable.

PARAMETER	DESCRIPTION
data	The data to save. TYPE: Any
file_path	The path to save the JSON file. TYPE: str

Example

>>> save_json(data, "path/to/file.json")

Source code in vision_agent/tools/tools.py

def save_json(data: Any, file_path: str) -> None:
    """'save_json' is a utility function that saves data as a JSON file. It is helpful
    for saving data that contains NumPy arrays which are not JSON serializable.

    Parameters:
        data (Any): The data to save.
        file_path (str): The path to save the JSON file.

    Example
    -------
        >>> save_json(data, "path/to/file.json")
    """

    class NumpyEncoder(json.JSONEncoder):
        def default(self, obj: Any):  # type: ignore
            if isinstance(obj, np.ndarray):
                return obj.tolist()
            elif isinstance(obj, np.bool_):
                return bool(obj)
            return json.JSONEncoder.default(self, obj)

    Path(file_path).parent.mkdir(parents=True, exist_ok=True)
    with open(file_path, "w") as f:
        json.dump(data, f, cls=NumpyEncoder)

load_image

load_image(image_path)

'load_image' is a utility function that loads an image from the given file path string or an URL.

PARAMETER	DESCRIPTION
image_path	The path or URL to the image. TYPE: str

RETURNS	DESCRIPTION
ndarray	np.ndarray: The image as a NumPy array.

Example

>>> load_image("path/to/image.jpg")

Source code in vision_agent/tools/tools.py

def load_image(image_path: str) -> np.ndarray:
    """'load_image' is a utility function that loads an image from the given file path string or an URL.

    Parameters:
        image_path (str): The path or URL to the image.

    Returns:
        np.ndarray: The image as a NumPy array.

    Example
    -------
        >>> load_image("path/to/image.jpg")
    """
    # NOTE: sometimes the generated code pass in a NumPy array
    if isinstance(image_path, np.ndarray):
        return image_path
    if image_path.startswith(("http", "https")):
        _, image_suffix = os.path.splitext(image_path)
        with tempfile.NamedTemporaryFile(delete=False, suffix=image_suffix) as tmp_file:
            # Download the image and save it to the temporary file
            with urllib.request.urlopen(image_path) as response:
                tmp_file.write(response.read())
            image_path = tmp_file.name
    image = Image.open(image_path).convert("RGB")
    return np.array(image)

save_image

save_image(image, file_path)

'save_image' is a utility function that saves an image to a file path.

PARAMETER	DESCRIPTION
image	The image to save. TYPE: ndarray
file_path	The path to save the image file. TYPE: str

Example

>>> save_image(image)

Source code in vision_agent/tools/tools.py

def save_image(image: np.ndarray, file_path: str) -> None:
    """'save_image' is a utility function that saves an image to a file path.

    Parameters:
        image (np.ndarray): The image to save.
        file_path (str): The path to save the image file.

    Example
    -------
        >>> save_image(image)
    """
    Path(file_path).parent.mkdir(parents=True, exist_ok=True)
    if not isinstance(image, np.ndarray) or (
        image.shape[0] == 0 and image.shape[1] == 0
    ):
        raise ValueError("The image is not a valid NumPy array with shape (H, W, C)")

    pil_image = Image.fromarray(image.astype(np.uint8)).convert("RGB")
    if should_report_tool_traces():
        from IPython.display import display

        display(pil_image)

    pil_image.save(file_path)

save_video

save_video(frames, output_video_path=None, fps=1)

'save_video' is a utility function that saves a list of frames as a mp4 video file on disk.

PARAMETER	DESCRIPTION
frames	A list of frames to save. TYPE: list[ndarray]
output_video_path	The path to save the video file. If not provided, a temporary file will be created. TYPE: str DEFAULT: None
fps	The number of frames composes a second in the video. TYPE: float DEFAULT: 1

RETURNS	DESCRIPTION
str	The path to the saved video file. TYPE: str

Example

>>> save_video(frames)
"/tmp/tmpvideo123.mp4"

Source code in vision_agent/tools/tools.py

def save_video(
    frames: List[np.ndarray], output_video_path: Optional[str] = None, fps: float = 1
) -> str:
    """'save_video' is a utility function that saves a list of frames as a mp4 video file on disk.

    Parameters:
        frames (list[np.ndarray]): A list of frames to save.
        output_video_path (str): The path to save the video file. If not provided, a temporary file will be created.
        fps (float): The number of frames composes a second in the video.

    Returns:
        str: The path to the saved video file.

    Example
    -------
        >>> save_video(frames)
        "/tmp/tmpvideo123.mp4"
    """
    if isinstance(fps, str):
        # fps could be a string when it's passed in from a web endpoint deployment
        fps = float(fps)
    if fps <= 0:
        raise ValueError(f"fps must be greater than 0 got {fps}")

    if not isinstance(frames, list) or len(frames) == 0:
        raise ValueError("Frames must be a list of NumPy arrays")

    for frame in frames:
        if not isinstance(frame, np.ndarray) or (
            frame.shape[0] == 0 and frame.shape[1] == 0
        ):
            raise ValueError("A frame is not a valid NumPy array with shape (H, W, C)")

    if output_video_path is None:
        output_video_path = tempfile.NamedTemporaryFile(
            delete=False, suffix=".mp4"
        ).name
    else:
        Path(output_video_path).parent.mkdir(parents=True, exist_ok=True)

    output_video_path = video_writer(frames, fps, output_video_path)
    _save_video_to_result(output_video_path)
    return output_video_path

overlay_bounding_boxes

overlay_bounding_boxes(medias, bboxes)

'overlay_bounding_boxes' is a utility function that displays bounding boxes on an image. It will draw a box around the detected object with the label and score.

PARAMETER	DESCRIPTION
medias	The image or frames to display the bounding boxes on. TYPE: Union[ndarray, List[ndarra]]
bboxes	A list of dictionaries or a list of list of dictionaries containing the bounding boxes. TYPE: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]

RETURNS	DESCRIPTION
Union[ndarray, List[ndarray]]	np.ndarray: The image with the bounding boxes, labels and scores displayed.

Example

>>> image_with_bboxes = overlay_bounding_boxes(
    image, [{'score': 0.99, 'label': 'dinosaur', 'bbox': [0.1, 0.11, 0.35, 0.4]}],
)

Source code in vision_agent/tools/tools.py

def overlay_bounding_boxes(
    medias: Union[np.ndarray, List[np.ndarray]],
    bboxes: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]],
) -> Union[np.ndarray, List[np.ndarray]]:
    """'overlay_bounding_boxes' is a utility function that displays bounding boxes on
    an image. It will draw a box around the detected object with the label and score.

    Parameters:
        medias (Union[np.ndarray, List[np.ndarra]]): The image or frames to display the
            bounding boxes on.
        bboxes (Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]): A list of
            dictionaries or a list of list of dictionaries containing the bounding
            boxes.

    Returns:
        np.ndarray: The image with the bounding boxes, labels and scores displayed.

    Example
    -------
        >>> image_with_bboxes = overlay_bounding_boxes(
            image, [{'score': 0.99, 'label': 'dinosaur', 'bbox': [0.1, 0.11, 0.35, 0.4]}],
        )
    """

    medias_int: List[np.ndarray] = (
        [medias] if isinstance(medias, np.ndarray) else medias
    )
    if len(bboxes) == 0:
        bbox_int: List[List[Dict[str, Any]]] = [[] for _ in medias_int]
    else:
        if isinstance(bboxes[0], dict):
            bbox_int = [cast(List[Dict[str, Any]], bboxes)]
        else:
            bbox_int = cast(List[List[Dict[str, Any]]], bboxes)

    labels = set([bb["label"] for b in bbox_int for bb in b])

    if len(labels) > len(COLORS):
        _LOGGER.warning(
            "Number of unique labels exceeds the number of available colors. Some labels may have the same color."
        )

    color = {label: COLORS[i % len(COLORS)] for i, label in enumerate(labels)}

    frame_out = []
    for i, frame in enumerate(medias_int):
        pil_image = Image.fromarray(frame.astype(np.uint8)).convert("RGB")

        bboxes = bbox_int[i]
        bboxes = sorted(bboxes, key=lambda x: x["label"], reverse=True)

        # if more than 50 boxes use small boxes to indicate objects else use regular boxes
        if len(bboxes) > 50:
            pil_image = _plot_counting(pil_image, bboxes, color)
        else:
            width, height = pil_image.size
            fontsize = max(12, int(min(width, height) / 40))
            draw = ImageDraw.Draw(pil_image)
            font = ImageFont.truetype(
                str(
                    resources.files("vision_agent.fonts").joinpath(
                        "default_font_ch_en.ttf"
                    )
                ),
                fontsize,
            )

            for elt in bboxes:
                label = elt["label"]
                box = elt["bbox"]
                scores = elt["score"]

                # denormalize the box if it is normalized
                box = denormalize_bbox(box, (height, width))
                draw.rectangle(box, outline=color[label], width=4)
                text = f"{label}: {scores:.2f}"
                text_box = draw.textbbox((box[0], box[1]), text=text, font=font)
                draw.rectangle(
                    (box[0], box[1], text_box[2], text_box[3]), fill=color[label]
                )
                draw.text((box[0], box[1]), text, fill="black", font=font)

        frame_out.append(np.array(pil_image))
    return_frame = frame_out[0] if len(frame_out) == 1 else frame_out

    return return_frame  # type: ignore

overlay_segmentation_masks

overlay_segmentation_masks(
    medias,
    masks,
    draw_label=True,
    secondary_label_key="tracking_label",
)

'overlay_segmentation_masks' is a utility function that displays segmentation masks. It will overlay a colored mask on the detected object with the label.

PARAMETER	DESCRIPTION
medias	The image or frames to display the masks on. TYPE: Union[ndarray, List[ndarray]]
masks	A list of dictionaries or a list of list of dictionaries containing the masks, labels and scores. TYPE: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]
draw_label	If True, the labels will be displayed on the image. TYPE: bool DEFAULT: True
secondary_label_key	The key to use for the secondary tracking label which is needed in videos to display tracking information. TYPE: str DEFAULT: 'tracking_label'

RETURNS	DESCRIPTION
Union[ndarray, List[ndarray]]	np.ndarray: The image with the masks displayed.

Example

>>> image_with_masks = overlay_segmentation_masks(
    image,
    [{
        'score': 0.99,
        'label': 'dinosaur',
        'mask': array([[0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0],
            ...,
            [0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
    }],
)

Source code in vision_agent/tools/tools.py

def overlay_segmentation_masks(
    medias: Union[np.ndarray, List[np.ndarray]],
    masks: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]],
    draw_label: bool = True,
    secondary_label_key: str = "tracking_label",
) -> Union[np.ndarray, List[np.ndarray]]:
    """'overlay_segmentation_masks' is a utility function that displays segmentation
    masks. It will overlay a colored mask on the detected object with the label.

    Parameters:
        medias (Union[np.ndarray, List[np.ndarray]]): The image or frames to display
            the masks on.
        masks (Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]): A list of
            dictionaries or a list of list of dictionaries containing the masks, labels
            and scores.
        draw_label (bool, optional): If True, the labels will be displayed on the image.
        secondary_label_key (str, optional): The key to use for the secondary
            tracking label which is needed in videos to display tracking information.

    Returns:
        np.ndarray: The image with the masks displayed.

    Example
    -------
        >>> image_with_masks = overlay_segmentation_masks(
            image,
            [{
                'score': 0.99,
                'label': 'dinosaur',
                'mask': array([[0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0],
                    ...,
                    [0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0]], dtype=uint8),
            }],
        )
    """
    if not masks:
        return medias

    medias_int: List[np.ndarray] = (
        [medias] if isinstance(medias, np.ndarray) else medias
    )
    masks_int = [masks] if isinstance(masks[0], dict) else masks
    masks_int = cast(List[List[Dict[str, Any]]], masks_int)

    labels = set()
    for mask_i in masks_int:
        for mask_j in mask_i:
            labels.add(mask_j["label"])
    color = {label: COLORS[i % len(COLORS)] for i, label in enumerate(labels)}

    width, height = Image.fromarray(medias_int[0]).size
    fontsize = max(12, int(min(width, height) / 40))
    font = ImageFont.truetype(
        str(resources.files("vision_agent.fonts").joinpath("default_font_ch_en.ttf")),
        fontsize,
    )

    frame_out = []
    for i, frame in enumerate(medias_int):
        pil_image = Image.fromarray(frame.astype(np.uint8)).convert("RGBA")
        for elt in masks_int[i]:
            mask = elt["mask"]
            label = elt["label"]
            tracking_lbl = elt.get(secondary_label_key, None)

            # Create semi-transparent mask overlay
            np_mask = np.zeros((pil_image.size[1], pil_image.size[0], 4))
            np_mask[mask > 0, :] = color[label] + (255 * 0.7,)
            mask_img = Image.fromarray(np_mask.astype(np.uint8))
            pil_image = Image.alpha_composite(pil_image, mask_img)

            # Draw contour border
            mask_uint8 = mask.astype(np.uint8) * 255
            contours, _ = cv2.findContours(
                mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
            )
            border_mask = np.zeros(
                (pil_image.size[1], pil_image.size[0], 4), dtype=np.uint8
            )
            cv2.drawContours(border_mask, contours, -1, color[label] + (255,), 8)
            border_img = Image.fromarray(border_mask)
            pil_image = Image.alpha_composite(pil_image, border_img)

            if draw_label:
                draw = ImageDraw.Draw(pil_image)
                text = tracking_lbl if tracking_lbl else label
                text_box = draw.textbbox((0, 0), text=text, font=font)
                x, y = _get_text_coords_from_mask(
                    mask,
                    v_gap=(text_box[3] - text_box[1]) + 10,
                    h_gap=(text_box[2] - text_box[0]) // 2,
                )
                if x != 0 and y != 0:
                    text_box = draw.textbbox((x, y), text=text, font=font)
                    draw.rectangle((x, y, text_box[2], text_box[3]), fill=color[label])
                    draw.text((x, y), text, fill="black", font=font)
        frame_out.append(np.array(pil_image))
    return_frame = frame_out[0] if len(frame_out) == 1 else frame_out

    return return_frame  # type: ignore

overlay_heat_map

overlay_heat_map(image, heat_map, alpha=0.8)

'overlay_heat_map' is a utility function that displays a heat map on an image.

PARAMETER	DESCRIPTION
image	The image to display the heat map on. TYPE: ndarray
heat_map	A dictionary containing the heat map under the key 'heat_map'. TYPE: Dict[str, Any]
alpha	The transparency of the overlay. Defaults to 0.8. TYPE: float DEFAULT: 0.8

RETURNS	DESCRIPTION
ndarray	np.ndarray: The image with the heat map displayed.

Example

>>> image_with_heat_map = overlay_heat_map(
    image,
    {
        'heat_map': array([[0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 0, 0, 0],
            ...,
            [0, 0, 0, ..., 0, 0, 0],
            [0, 0, 0, ..., 125, 125, 125]], dtype=uint8),
    },
)

Source code in vision_agent/tools/tools.py

def overlay_heat_map(
    image: np.ndarray, heat_map: Dict[str, Any], alpha: float = 0.8
) -> np.ndarray:
    """'overlay_heat_map' is a utility function that displays a heat map on an image.

    Parameters:
        image (np.ndarray): The image to display the heat map on.
        heat_map (Dict[str, Any]): A dictionary containing the heat map under the key
            'heat_map'.
        alpha (float, optional): The transparency of the overlay. Defaults to 0.8.

    Returns:
        np.ndarray: The image with the heat map displayed.

    Example
    -------
        >>> image_with_heat_map = overlay_heat_map(
            image,
            {
                'heat_map': array([[0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 0, 0, 0],
                    ...,
                    [0, 0, 0, ..., 0, 0, 0],
                    [0, 0, 0, ..., 125, 125, 125]], dtype=uint8),
            },
        )
    """
    pil_image = Image.fromarray(image.astype(np.uint8)).convert("RGB")

    if "heat_map" not in heat_map or len(heat_map["heat_map"]) == 0:
        return image

    pil_image = pil_image.convert("L")
    mask = Image.fromarray(heat_map["heat_map"])
    mask = mask.resize(pil_image.size)

    overlay = Image.new("RGBA", mask.size)
    odraw = ImageDraw.Draw(overlay)
    odraw.bitmap((0, 0), mask, fill=(255, 0, 0, round(alpha * 255)))
    combined = Image.alpha_composite(
        pil_image.convert("RGBA"), overlay.resize(pil_image.size)
    )
    return np.array(combined)

vision_agent.tools.tool_utils

ToolCallTrace

Bases: BaseModel

endpoint_url instance-attribute

endpoint_url

type instance-attribute

type

request instance-attribute

request

response instance-attribute

response

error instance-attribute

error

files instance-attribute

files

should_report_tool_traces

should_report_tool_traces()

Source code in vision_agent/tools/tool_utils.py

def should_report_tool_traces() -> bool:
    return bool(os.environ.get("REPORT_TOOL_TRACES", False))

send_inference_request

send_inference_request(
    payload,
    endpoint_name,
    files=None,
    v2=False,
    metadata_payload=None,
    is_form=False,
)

Source code in vision_agent/tools/tool_utils.py

def send_inference_request(
    payload: Dict[str, Any],
    endpoint_name: str,
    files: Optional[List[Tuple[Any, ...]]] = None,
    v2: bool = False,
    metadata_payload: Optional[Dict[str, Any]] = None,
    is_form: bool = False,
) -> Any:
    url = f"{_LND_API_URL_v2 if v2 else _LND_API_URL}/{endpoint_name}"
    if "TOOL_ENDPOINT_URL" in os.environ:
        url = os.environ["TOOL_ENDPOINT_URL"]

    headers = {"apikey": _LND_API_KEY}
    if "TOOL_ENDPOINT_AUTH" in os.environ:
        headers["Authorization"] = os.environ["TOOL_ENDPOINT_AUTH"]
        headers.pop("apikey")

    if runtime_tag := os.environ.get("RUNTIME_TAG", ""):
        headers["runtime_tag"] = runtime_tag

    session = _create_requests_session(
        url=url,
        num_retry=3,
        headers=headers,
    )

    function_name = "unknown"
    if "function_name" in payload:
        function_name = payload["function_name"]
    elif metadata_payload is not None and "function_name" in metadata_payload:
        function_name = metadata_payload["function_name"]

    response = _call_post(url, payload, session, files, function_name, is_form)

    return response["data"]

send_task_inference_request

send_task_inference_request(
    payload,
    task_name,
    files=None,
    metadata=None,
    is_form=False,
)

Source code in vision_agent/tools/tool_utils.py

def send_task_inference_request(
    payload: Dict[str, Any],
    task_name: str,
    files: Optional[List[Tuple[Any, ...]]] = None,
    metadata: Optional[Dict[str, Any]] = None,
    is_form: bool = False,
) -> Any:
    url = f"{_LND_API_URL_v2}/{task_name}"
    headers = {"apikey": _LND_API_KEY}
    session = _create_requests_session(
        url=url,
        num_retry=3,
        headers=headers,
    )

    function_name = "unknown"
    if metadata is not None and "function_name" in metadata:
        function_name = metadata["function_name"]
    response = _call_post(url, payload, session, files, function_name, is_form)
    return response["data"]

get_tool_documentation

get_tool_documentation(funcs)

Source code in vision_agent/tools/tool_utils.py

def get_tool_documentation(funcs: List[Callable[..., Any]]) -> str:
    docstrings = ""
    for func in funcs:
        docstrings += f"{func.__name__}{inspect.signature(func)}:\n{func.__doc__}\n\n"

    return docstrings

get_tool_descriptions

get_tool_descriptions(funcs)

Source code in vision_agent/tools/tool_utils.py

def get_tool_descriptions(funcs: List[Callable[..., Any]]) -> str:
    descriptions = ""
    for func in funcs:
        description = func.__doc__
        if description is None:
            description = ""

        if "Parameters:" in description:
            description = (
                description[: description.find("Parameters:")]
                .replace("\n", " ")
                .strip()
            )

        description = " ".join(description.split())
        descriptions += f"- {func.__name__}{inspect.signature(func)}: {description}\n"
    return descriptions

get_tool_descriptions_by_names

get_tool_descriptions_by_names(
    tool_name, funcs, util_funcs
)

Source code in vision_agent/tools/tool_utils.py

def get_tool_descriptions_by_names(
    tool_name: Optional[List[str]],
    funcs: List[Callable[..., Any]],
    util_funcs: List[
        Callable[..., Any]
    ],  # util_funcs will always be added to the list of functions
) -> str:
    if tool_name is None:
        return get_tool_descriptions(funcs + util_funcs)

    invalid_names = [
        name for name in tool_name if name not in {func.__name__ for func in funcs}
    ]

    if invalid_names:
        raise ValueError(f"Invalid customized tool names: {', '.join(invalid_names)}")

    filtered_funcs = (
        funcs
        if not tool_name
        else [func for func in funcs if func.__name__ in tool_name]
    )
    return get_tool_descriptions(filtered_funcs + util_funcs)

get_tools_df

get_tools_df(funcs)

Source code in vision_agent/tools/tool_utils.py

def get_tools_df(funcs: List[Callable[..., Any]]) -> pd.DataFrame:
    data: Dict[str, List[str]] = {"desc": [], "doc": [], "name": []}

    for func in funcs:
        desc = func.__doc__
        if desc is None:
            desc = ""
        desc = desc[: desc.find("Parameters:")].replace("\n", " ").strip()
        desc = " ".join(desc.split())

        doc = f"{func.__name__}{inspect.signature(func)}:\n{func.__doc__}"
        data["desc"].append(desc)
        data["doc"].append(doc)
        data["name"].append(func.__name__)

    return pd.DataFrame(data)  # type: ignore

get_tools_info

get_tools_info(funcs)

Source code in vision_agent/tools/tool_utils.py

def get_tools_info(funcs: List[Callable[..., Any]]) -> Dict[str, str]:
    data: Dict[str, str] = {}

    for func in funcs:
        desc = func.__doc__
        if desc is None:
            desc = ""

        data[func.__name__] = f"{func.__name__}{inspect.signature(func)}:\n{desc}"

    return data

filter_bboxes_by_threshold

filter_bboxes_by_threshold(bboxes, threshold)

Source code in vision_agent/tools/tool_utils.py

def filter_bboxes_by_threshold(
    bboxes: BoundingBoxes, threshold: float
) -> BoundingBoxes:
    return list(filter(lambda bbox: bbox.score >= threshold, bboxes))

add_bboxes_from_masks

add_bboxes_from_masks(all_preds)

Source code in vision_agent/tools/tool_utils.py

def add_bboxes_from_masks(
    all_preds: List[List[Dict[str, Any]]],
) -> List[List[Dict[str, Any]]]:
    for frame_preds in all_preds:
        for preds in frame_preds:
            if np.sum(preds["mask"]) == 0:
                preds["bbox"] = []
            else:
                rows, cols = np.where(preds["mask"])
                bbox = [
                    float(np.min(cols)),
                    float(np.min(rows)),
                    float(np.max(cols)),
                    float(np.max(rows)),
                ]
                bbox = normalize_bbox(bbox, preds["mask"].shape)
                preds["bbox"] = bbox

    return all_preds

calculate_iou

calculate_iou(bbox1, bbox2)

Source code in vision_agent/tools/tool_utils.py

def calculate_iou(bbox1: List[float], bbox2: List[float]) -> float:
    x1, y1, x2, y2 = bbox1
    x3, y3, x4, y4 = bbox2

    x_overlap = max(0, min(x2, x4) - max(x1, x3))
    y_overlap = max(0, min(y2, y4) - max(y1, y3))
    intersection = x_overlap * y_overlap

    area1 = (x2 - x1) * (y2 - y1)
    area2 = (x4 - x3) * (y4 - y3)
    union = area1 + area2 - intersection

    return intersection / union if union > 0 else 0

single_nms

single_nms(preds, iou_threshold)

Source code in vision_agent/tools/tool_utils.py

def single_nms(
    preds: List[Dict[str, Any]], iou_threshold: float
) -> List[Dict[str, Any]]:
    for i in range(len(preds)):
        for j in range(i + 1, len(preds)):
            if calculate_iou(preds[i]["bbox"], preds[j]["bbox"]) > iou_threshold:
                if preds[i]["score"] > preds[j]["score"]:
                    preds[j]["score"] = 0
                else:
                    preds[i]["score"] = 0

    return [pred for pred in preds if pred["score"] > 0]

nms

nms(all_preds, iou_threshold)

Source code in vision_agent/tools/tool_utils.py

def nms(
    all_preds: List[List[Dict[str, Any]]], iou_threshold: float
) -> List[List[Dict[str, Any]]]:
    return_preds = []
    for frame_preds in all_preds:
        frame_preds = single_nms(frame_preds, iou_threshold)
        return_preds.append(frame_preds)

    return return_preds