vision-agent 0.2.32__py3-none-any.whl → 0.2.33__py3-none-any.whl

vision_agent/agent/vision_agent.py

@@ -9,6 +9,7 @@ from rich.console import Console
 from rich.syntax import Syntax
 from tabulate import tabulate
 
+import vision_agent.tools as T
 from vision_agent.agent import Agent
 from vision_agent.agent.vision_agent_prompts import (
     CODE,
@@ -22,7 +23,6 @@ from vision_agent.agent.vision_agent_prompts import (
 )
 from vision_agent.llm import LLM, OpenAILLM
 from vision_agent.lmm import LMM, OpenAILMM
-from vision_agent.tools import TOOL_DESCRIPTIONS, TOOLS_DF, UTILITIES_DOCSTRING
 from vision_agent.utils import Execute
 from vision_agent.utils.sim import Sim
 
@@ -31,6 +31,7 @@ _LOGGER = logging.getLogger(__name__)
 _MAX_TABULATE_COL_WIDTH = 80
 _EXECUTE = Execute(600)
 _CONSOLE = Console()
+_DEFAULT_IMPORT = "\n".join(T.__new_tools__)
 
 
 def format_memory(memory: List[Dict[str, str]]) -> str:
@@ -140,7 +141,7 @@ def write_and_test_code(
         )
     )
 
-    success, result = _EXECUTE.run_isolation(f"{code}\n{test}")
+    success, result = _EXECUTE.run_isolation(f"{_DEFAULT_IMPORT}\n{code}\n{test}")
     if verbosity == 2:
         _LOGGER.info("Initial code and tests:")
         log_progress(
@@ -184,7 +185,7 @@ def write_and_test_code(
             {"code": f"{code}\n{test}", "feedback": fixed_code_and_test["reflections"]}
         )
 
-        success, result = _EXECUTE.run_isolation(f"{code}\n{test}")
+        success, result = _EXECUTE.run_isolation(f"{_DEFAULT_IMPORT}\n{code}\n{test}")
         if verbosity == 2:
             log_progress(
                 {
@@ -300,7 +301,7 @@ class VisionAgent(Agent):
         )
 
         self.tool_recommender = (
-            Sim(TOOLS_DF, sim_key="desc")
+            Sim(T.TOOLS_DF, sim_key="desc")
             if tool_recommender is None
             else tool_recommender
         )
@@ -358,6 +359,10 @@ class VisionAgent(Agent):
                 if chat_i["role"] == "user":
                     chat_i["content"] += f" Image name {media}"
 
+        # re-grab custom tools
+        global _DEFAULT_IMPORT
+        _DEFAULT_IMPORT = "\n".join(T.__new_tools__)
+
         code = ""
         test = ""
         working_memory: List[Dict[str, str]] = []
@@ -369,7 +374,7 @@ class VisionAgent(Agent):
         while not success and retries < self.max_retries:
             plan_i = write_plan(
                 chat,
-                TOOL_DESCRIPTIONS,
+                T.TOOL_DESCRIPTIONS,
                 format_memory(working_memory),
                 self.planner,
                 media=[media] if media else None,
@@ -397,7 +402,7 @@ class VisionAgent(Agent):
             results = write_and_test_code(
                 FULL_TASK.format(user_request=chat[0]["content"], subtasks=plan_i_str),
                 tool_info,
-                UTILITIES_DOCSTRING,
+                T.UTILITIES_DOCSTRING,
                 format_memory(working_memory),
                 self.coder,
                 self.tester,

vision_agent/tools/__init__.py

@@ -1,3 +1,5 @@
+from typing import Callable, List, Optional
+
 from .prompts import CHOOSE_PARAMS, SYSTEM_PROMPT
 from .tools import (
     TOOL_DESCRIPTIONS,
@@ -16,9 +18,39 @@ from .tools import (
     load_image,
     ocr,
     overlay_bounding_boxes,
+    overlay_heat_map,
     overlay_segmentation_masks,
     save_image,
     save_json,
     visual_prompt_counting,
     zero_shot_counting,
 )
+
+__new_tools__ = [
+    "import vision_agent as va",
+    "from vision_agent.tools import register_tool",
+]
+
+
+def register_tool(imports: Optional[List] = None) -> Callable:
+    def decorator(tool: Callable) -> Callable:
+        import inspect
+
+        from .tools import get_tool_descriptions, get_tool_documentation, get_tools_df
+
+        global TOOLS, TOOLS_DF, TOOL_DESCRIPTIONS, TOOL_DOCSTRING
+
+        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
+
+            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.py

@@ -203,7 +203,7 @@ def ocr(image: np.ndarray) -> List[Dict[str, Any]]:
 
     Returns:
         List[Dict[str, Any]]: A list of dictionaries containing the detected text, bbox,
-        and confidence score.
+            and confidence score.
 
     Example
     -------
@@ -247,14 +247,16 @@ def ocr(image: np.ndarray) -> List[Dict[str, Any]]:
 
 
 def zero_shot_counting(image: np.ndarray) -> Dict[str, Any]:
-    """'zero_shot_counting' is a tool that counts the dominant foreground object given an image and no other information about the content.
-    It returns only the count of the objects in the image.
+    """'zero_shot_counting' is a tool that counts the dominant foreground object given
+    an image and no other information about the content. It returns only the count of
+    the objects in the image.
 
     Parameters:
         image (np.ndarray): The image that contains lot of instances of a single object
 
     Returns:
-        Dict[str, Any]: A dictionary containing the key 'count' and the count as a value. E.g. {count: 12}.
+        Dict[str, Any]: A dictionary containing the key 'count' and the count as a
+            value. E.g. {count: 12}.
 
     Example
     -------
@@ -276,14 +278,16 @@ def zero_shot_counting(image: np.ndarray) -> Dict[str, Any]:
 def visual_prompt_counting(
     image: np.ndarray, visual_prompt: Dict[str, List[float]]
 ) -> Dict[str, Any]:
-    """'visual_prompt_counting' is a tool that counts the dominant foreground object given an image and a visual prompt which is a bounding box describing the object.
+    """'visual_prompt_counting' is a tool that counts the dominant foreground object
+    given an image and a visual prompt which is a bounding box describing the object.
     It returns only the count of the objects in the image.
 
     Parameters:
         image (np.ndarray): The image that contains lot of instances of a single object
 
     Returns:
-        Dict[str, Any]: A dictionary containing the key 'count' and the count as a value. E.g. {count: 12}.
+        Dict[str, Any]: A dictionary containing the key 'count' and the count as a
+            value. E.g. {count: 12}.
 
     Example
     -------
@@ -308,15 +312,17 @@ def visual_prompt_counting(
 
 
 def image_question_answering(image: np.ndarray, prompt: str) -> str:
-    """'image_question_answering_' is a tool that can answer questions about the visual contents of an image given a question and an image.
-    It returns an answer to the question
+    """'image_question_answering_' is a tool that can answer questions about the visual
+    contents of an image given a question and an image. It returns an answer to the
+    question
 
     Parameters:
         image (np.ndarray): The reference image used for the question
         prompt (str): The question about the image
 
     Returns:
-        str: A string which is the answer to the given prompt. E.g. {'text': 'This image contains a cat sitting on a table with a bowl of milk.'}.
+        str: A string which is the answer to the given prompt. E.g. {'text': 'This
+            image contains a cat sitting on a table with a bowl of milk.'}.
 
     Example
     -------
@@ -338,14 +344,16 @@ def image_question_answering(image: np.ndarray, prompt: str) -> str:
 
 def clip(image: np.ndarray, classes: List[str]) -> Dict[str, Any]:
     """'clip' is a tool that can classify an image given a list of input classes or tags.
-    It returns the same list of the input classes along with their probability scores based on image content.
+    It returns the same list of the input classes along with their probability scores
+    based on image content.
 
     Parameters:
         image (np.ndarray): The image to classify or tag
         classes (List[str]): The list of classes 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.
+        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
     -------
@@ -366,8 +374,8 @@ def clip(image: np.ndarray, classes: List[str]) -> Dict[str, Any]:
 
 
 def image_caption(image: np.ndarray) -> str:
-    """'image_caption' is a tool that can caption an image based on its contents.
-    It returns a text describing the image.
+    """'image_caption' is a tool that can caption an image based on its contents. It
+    returns a text describing the image.
 
     Parameters:
         image (np.ndarray): The image to caption
@@ -619,6 +627,51 @@ def overlay_segmentation_masks(
     return np.array(pil_image.convert("RGB"))
 
 
+def overlay_heat_map(
+    image: np.ndarray, heat_map: Dict[str, Any], alpha: float = 0.8
+) -> np.ndarray:
+    """'display_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 = display_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.convert("RGB"))
+
+
 def get_tool_documentation(funcs: List[Callable[..., Any]]) -> str:
     docstrings = ""
     for func in funcs:
@@ -634,9 +687,13 @@ def get_tool_descriptions(funcs: List[Callable[..., Any]]) -> str:
         if description is None:
             description = ""
 
-        description = (
-            description[: description.find("Parameters:")].replace("\n", " ").strip()
-        )
+        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
@@ -676,10 +733,18 @@ TOOLS = [
     save_image,
     overlay_bounding_boxes,
     overlay_segmentation_masks,
+    overlay_heat_map,
 ]
 TOOLS_DF = get_tools_df(TOOLS)  # type: ignore
 TOOL_DESCRIPTIONS = get_tool_descriptions(TOOLS)  # type: ignore
 TOOL_DOCSTRING = get_tool_documentation(TOOLS)  # type: ignore
 UTILITIES_DOCSTRING = get_tool_documentation(
-    [save_json, load_image, save_image, overlay_bounding_boxes]
+    [
+        save_json,
+        load_image,
+        save_image,
+        overlay_bounding_boxes,
+        overlay_segmentation_masks,
+        overlay_heat_map,
+    ]
 )

{vision_agent-0.2.32.dist-info → vision_agent-0.2.33.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vision-agent
-Version: 0.2.32
+Version: 0.2.33
 Summary: Toolset for Vision Agent
 Author: Landing AI
 Author-email: dev@landing.ai
@@ -153,6 +153,34 @@ you. For example:
 }]
 ```
 
+You can also add custom tools to the agent:
+
+```python
+import vision_agent as va
+
+@va.tools.register_tool(imports=["import numpy as np"])
+def custom_tool(image_path: str) -> str:
+    """My custom tool documentation.
+
+    Parameters:
+        image_path (str): The path to the image.
+
+    Returns:
+        str: The result of the tool.
+
+    Example
+    -------
+    >>> custom_tool("image.jpg")
+    """
+
+    import numpy as np
+    return np.zeros((10, 10))
+```
+
+You need to ensure you call `@va.tools.register_tool` with any imports it might use and
+ensure the documentation is in the same format above with description, `Parameters:`,
+`Returns:`, and `Example\n-------`. You can find an example use case [here](examples/custom_tools/).
+
 ### Azure Setup
 If you want to use Azure OpenAI models, you can set the environment variable:
 
@@ -166,9 +194,10 @@ You can then run Vision Agent using the Azure OpenAI models:
 ```python
 >>> import vision_agent as va
 >>> agent = va.agent.VisionAgent(
->>>     task_model=va.llm.AzureOpenAILLM(),
->>>     answer_model=va.lmm.AzureOpenAILMM(),
->>>     reflection_model=va.lmm.AzureOpenAILMM(),
+>>>     planner=va.llm.AzureOpenAILLM(),
+>>>     coder=va.lmm.AzureOpenAILMM(),
+>>>     tester=va.lmm.AzureOpenAILMM(),
+>>>     debugger=va.lmm.AzureOpenAILMM(),
 >>> )
 ```
 

{vision_agent-0.2.32.dist-info → vision_agent-0.2.33.dist-info}/RECORD

@@ -11,7 +11,7 @@ vision_agent/agent/easytool_v2.py,sha256=CjY-sSj3abxnSq3ZHZMt-7YvRWDXEZsC6RN8FFI
 vision_agent/agent/easytool_v2_prompts.py,sha256=MZSIwovYgB-f-kdJ6btaNDVXptJn47bfOL3-Zn6NiC0,8573
 vision_agent/agent/reflexion.py,sha256=AlM5AvBJvCslXlYQdZiadq4oVHsNBm3IF_03DglTxRo,10506
 vision_agent/agent/reflexion_prompts.py,sha256=G7UAeNz_g2qCb2yN6OaIC7bQVUkda4m3z42EG8wAyfE,9342
-vision_agent/agent/vision_agent.py,sha256=gvGNPSyZ1WpIzBcxz03xAe96cSdVBI0zCYIfrPs9jw0,15326
+vision_agent/agent/vision_agent.py,sha256=0_3hRcN4b0O52kDk8eq8TMIwBO2FA2XCC4JxAWK0ZEU,15482
 vision_agent/agent/vision_agent_prompts.py,sha256=0YbiS59IEWbiE43gCvOqfWrpudIAhTn8FHzXW0Y-Gaw,8201
 vision_agent/fonts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vision_agent/fonts/default_font_ch_en.ttf,sha256=1YM0Z3XqLDjSNbF7ihQFSAIUdjF9m1rtHiNC_6QosTE,1594400
@@ -19,18 +19,18 @@ vision_agent/llm/__init__.py,sha256=BoUm_zSAKnLlE8s-gKTSQugXDqVZKPqYlWwlTLdhcz4,
 vision_agent/llm/llm.py,sha256=UZ73GqQHE-NKOJWsrOTWfmdHYsbCBkJ5rZ7dhcSCHHw,5951
 vision_agent/lmm/__init__.py,sha256=nnNeKD1k7q_4vLb1x51O_EUTYaBgGfeiCx5F433gr3M,67
 vision_agent/lmm/lmm.py,sha256=NwcZYLTzi95LSMAk0sTtw7G_zBLa9lU-DHM5GUUCiK4,10622
-vision_agent/tools/__init__.py,sha256=1kyJy4euA8t73_ALhKZIUOjVb2A1IyYztu-MJJJ0TYI,505
+vision_agent/tools/__init__.py,sha256=oZa_sslb1UqEgpdWROChDcz5JHdB475ejJX78FMLYvE,1512
 vision_agent/tools/easytool_tools.py,sha256=pZc5dQlYINlV4nYbbzsDi3-wauA-fCeD2iGmJUMoUfE,47373
 vision_agent/tools/prompts.py,sha256=V1z4YJLXZuUl_iZ5rY0M5hHc_2tmMEUKr0WocXKGt4E,1430
 vision_agent/tools/tool_utils.py,sha256=wzRacbUpqk9hhfX_Y08rL8qP0XCN2w-8IZoYLi3Upn4,869
-vision_agent/tools/tools.py,sha256=mio0A1l5QcyRC5IgaD4Trfqg7hFTZ8rOjx1dYivwb4Q,21585
+vision_agent/tools/tools.py,sha256=nXjefpW9L-Xuos73ObDqpmJfOyUAJVrzoiHsxEE7O10,23346
 vision_agent/utils/__init__.py,sha256=xsHFyJSDbLdonB9Dh74cwZnVTiT__2OQF3Brd3Nmglc,116
 vision_agent/utils/execute.py,sha256=8_SfK-IkHH4lXF0JVyV7sDFszZn9HKsh1bFITKGCJ1g,3881
 vision_agent/utils/image_utils.py,sha256=_cdiS5YrLzqkq_ZgFUO897m5M4_SCIThwUy4lOklfB8,7700
 vision_agent/utils/sim.py,sha256=oUZ-6eu8Io-UNt9GXJ0XRKtP-Wc0sPWVzYGVpB2yDFk,3001
 vision_agent/utils/type_defs.py,sha256=BlI8ywWHAplC7kYWLvt4AOdnKpEW3qWEFm-GEOSkrFQ,1792
 vision_agent/utils/video.py,sha256=xTElFSFp1Jw4ulOMnk81Vxsh-9dTxcWUO6P9fzEi3AM,7653
-vision_agent-0.2.32.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-vision_agent-0.2.32.dist-info/METADATA,sha256=okG2Wf1LavPw4agq2iH7tY_rLlrabKgOOCL43pzfEMY,5942
-vision_agent-0.2.32.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
-vision_agent-0.2.32.dist-info/RECORD,,
+vision_agent-0.2.33.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+vision_agent-0.2.33.dist-info/METADATA,sha256=IAp4r_1YMih7p1kgZQ9rmjrRqD1vM0ManpXMYOMoLW4,6698
+vision_agent-0.2.33.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
+vision_agent-0.2.33.dist-info/RECORD,,