vision-agent 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

vision_agent/agent/vision_agent.py

@@ -314,6 +314,7 @@ def _handle_extract_frames(
                 image_to_data[image] = {
                     "bboxes": [],
                     "masks": [],
+                    "heat_map": [],
                     "labels": [],
                     "scores": [],
                 }
@@ -340,9 +341,12 @@ def _handle_viz_tools(
             return image_to_data
 
     for param, call_result in zip(parameters, tool_result["call_results"]):
-        # calls can fail, so we need to check if the call was successful
+        # Calls can fail, so we need to check if the call was successful. It can either:
+        # 1. return a str or some error that's not a dictionary
+        # 2. return a dictionary but not have the necessary keys
+
         if not isinstance(call_result, dict) or (
-            "bboxes" not in call_result and "masks" not in call_result
+            "bboxes" not in call_result and "heat_map" not in call_result
         ):
             return image_to_data
 
@@ -352,6 +356,7 @@ def _handle_viz_tools(
             image_to_data[image] = {
                 "bboxes": [],
                 "masks": [],
+                "heat_map": [],
                 "labels": [],
                 "scores": [],
             }
@@ -360,12 +365,28 @@ def _handle_viz_tools(
         image_to_data[image]["labels"].extend(call_result.get("labels", []))
         image_to_data[image]["scores"].extend(call_result.get("scores", []))
         image_to_data[image]["masks"].extend(call_result.get("masks", []))
+        # only single heatmap is returned
+        image_to_data[image]["heat_map"].append(call_result.get("heat_map", []))
         if "mask_shape" in call_result:
             image_to_data[image]["mask_shape"] = call_result["mask_shape"]
 
     return image_to_data
 
 
+def sample_n_evenly_spaced(lst: Sequence, n: int) -> Sequence:
+    if n <= 0:
+        return []
+    elif len(lst) == 0:
+        return []
+    elif n == 1:
+        return [lst[0]]
+    elif n >= len(lst):
+        return lst
+
+    spacing = (len(lst) - 1) / (n - 1)
+    return [lst[round(spacing * i)] for i in range(n)]
+
+
 def visualize_result(all_tool_results: List[Dict]) -> Sequence[Union[str, Path]]:
     image_to_data: Dict[str, Dict] = {}
     for tool_result in all_tool_results:
@@ -466,9 +487,14 @@ class VisionAgent(Agent):
         """Invoke the vision agent.
 
         Parameters:
-            input: a prompt that describe the task or a conversation in the format of
+            chat: A conversation in the format of
                 [{"role": "user", "content": "describe your task here..."}].
-            image: the input image referenced in the prompt parameter.
+            image: The input image referenced in the chat parameter.
+            reference_data: A dictionary containing the reference image, mask or bounding
+                box in the format of:
+                {"image": "image.jpg", "mask": "mask.jpg", "bbox": [0.1, 0.2, 0.1, 0.2]}
+                where the bounding box coordinates are normalized.
+            visualize_output: Whether to visualize the output.
 
         Returns:
             The result of the vision agent in text.
@@ -508,12 +534,14 @@ class VisionAgent(Agent):
         """Chat with the vision agent and return the final answer and all tool results.
 
         Parameters:
-            chat: a conversation in the format of
+            chat: A conversation in the format of
                 [{"role": "user", "content": "describe your task here..."}].
-            image: the input image referenced in the chat parameter.
-            reference_data: a dictionary containing the reference image and mask. in the
-                format of {"image": "image.jpg", "mask": "mask.jpg}
-            visualize_output: whether to visualize the output.
+            image: The input image referenced in the chat parameter.
+            reference_data: A dictionary containing the reference image, mask or bounding
+                box in the format of:
+                {"image": "image.jpg", "mask": "mask.jpg", "bbox": [0.1, 0.2, 0.1, 0.2]}
+                where the bounding box coordinates are normalized.
+            visualize_output: Whether to visualize the output.
 
         Returns:
             A tuple where the first item is the final answer and the second item is a
@@ -584,7 +612,7 @@ class VisionAgent(Agent):
             visualized_output = visualize_result(all_tool_results)
             all_tool_results.append({"visualized_output": visualized_output})
             if len(visualized_output) > 0:
-                reflection_images = visualized_output
+                reflection_images = sample_n_evenly_spaced(visualized_output, 3)
             elif image is not None:
                 reflection_images = [image]
             else:

vision_agent/image_utils.py

@@ -211,7 +211,7 @@ def overlay_masks(
     }
 
     for label, mask in zip(masks["labels"], masks["masks"]):
-        if isinstance(mask, str):
+        if isinstance(mask, str) or isinstance(mask, Path):
             mask = np.array(Image.open(mask))
         np_mask = np.zeros((image.size[1], image.size[0], 4))
         np_mask[mask > 0, :] = color[label] + (255 * alpha,)
@@ -221,7 +221,7 @@ def overlay_masks(
 
 
 def overlay_heat_map(
-    image: Union[str, Path, np.ndarray, ImageType], masks: Dict, alpha: float = 0.8
+    image: Union[str, Path, np.ndarray, ImageType], heat_map: Dict, alpha: float = 0.8
 ) -> ImageType:
     r"""Plots heat map on to an image.
 
@@ -238,14 +238,12 @@ def overlay_heat_map(
     elif isinstance(image, np.ndarray):
         image = Image.fromarray(image)
 
-    if "masks" not in masks:
+    if "heat_map" not in heat_map:
         return image.convert("RGB")
 
-    # Only one heat map per image, so no need to loop through masks
     image = image.convert("L")
-
-    if isinstance(masks["masks"][0], str):
-        mask = b64_to_pil(masks["masks"][0])
+    # Only one heat map per image, so no need to loop through masks
+    mask = Image.fromarray(heat_map["heat_map"][0])
 
     overlay = Image.new("RGBA", mask.size)
     odraw = ImageDraw.Draw(overlay)

vision_agent/lmm/lmm.py

@@ -9,10 +9,7 @@ from typing import Any, Callable, Dict, List, Optional, Union, cast
 import requests
 from openai import AzureOpenAI, OpenAI
 
-from vision_agent.tools import (
-    CHOOSE_PARAMS,
-    SYSTEM_PROMPT,
-)
+from vision_agent.tools import CHOOSE_PARAMS, SYSTEM_PROMPT
 
 _LOGGER = logging.getLogger(__name__)
 

vision_agent/tools/__init__.py

@@ -12,12 +12,12 @@ from .tools import (  # Counter,
     GroundingDINO,
     GroundingSAM,
     ImageCaption,
-    ZeroShotCounting,
-    VisualPromptCounting,
-    VisualQuestionAnswering,
     ImageQuestionAnswering,
     SegArea,
     SegIoU,
     Tool,
+    VisualPromptCounting,
+    VisualQuestionAnswering,
+    ZeroShotCounting,
     register_tool,
 )

vision_agent/tools/tools.py

@@ -11,15 +11,16 @@ from PIL import Image
 from PIL.Image import Image as ImageType
 
 from vision_agent.image_utils import (
+    b64_to_pil,
     convert_to_b64,
     denormalize_bbox,
     get_image_size,
     normalize_bbox,
     rle_decode,
 )
+from vision_agent.lmm import OpenAILMM
 from vision_agent.tools.video import extract_frames_from_video
 from vision_agent.type_defs import LandingaiAPIKey
-from vision_agent.lmm import OpenAILMM
 
 _LOGGER = logging.getLogger(__name__)
 _LND_API_KEY = LandingaiAPIKey().api_key
@@ -516,7 +517,9 @@ class ZeroShotCounting(Tool):
             "image": image_b64,
             "tool": "zero_shot_counting",
         }
-        return _send_inference_request(data, "tools")
+        resp_data = _send_inference_request(data, "tools")
+        resp_data["heat_map"] = np.array(b64_to_pil(resp_data["heat_map"][0]))
+        return resp_data
 
 
 class VisualPromptCounting(Tool):
@@ -585,7 +588,9 @@ class VisualPromptCounting(Tool):
             "prompt": prompt,
             "tool": "few_shot_counting",
         }
-        return _send_inference_request(data, "tools")
+        resp_data = _send_inference_request(data, "tools")
+        resp_data["heat_map"] = np.array(b64_to_pil(resp_data["heat_map"][0]))
+        return resp_data
 
 
 class VisualQuestionAnswering(Tool):

{vision_agent-0.2.4.dist-info → vision_agent-0.2.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vision-agent
-Version: 0.2.4
+Version: 0.2.6
 Summary: Toolset for Vision Agent
 Author: Landing AI
 Author-email: dev@landing.ai

{vision_agent-0.2.4.dist-info → vision_agent-0.2.6.dist-info}/RECORD

@@ -5,21 +5,21 @@ vision_agent/agent/easytool.py,sha256=oMHnBg7YBtIPgqQUNcZgq7uMgpPThs99_UnO7ERkMV
 vision_agent/agent/easytool_prompts.py,sha256=zdQQw6WpXOmvwOMtlBlNKY5a3WNlr65dbUvMIGiqdeo,4526
 vision_agent/agent/reflexion.py,sha256=4gz30BuFMeGxSsTzoDV4p91yE0R8LISXp28IaOI6wdM,10506
 vision_agent/agent/reflexion_prompts.py,sha256=G7UAeNz_g2qCb2yN6OaIC7bQVUkda4m3z42EG8wAyfE,9342
-vision_agent/agent/vision_agent.py,sha256=Ehb97lyPs7lYM9ipx07yxm6c2kUqz2OnjGQsv-nMwKA,24849
+vision_agent/agent/vision_agent.py,sha256=xepqtPqxwIEj0V9OyPtlBr4hsE67BazmqcXHBjUO8a4,25971
 vision_agent/agent/vision_agent_prompts.py,sha256=W3Z72FpUt71UIJSkjAcgtQqxeMqkYuATqHAN5fYY26c,7342
 vision_agent/fonts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vision_agent/fonts/default_font_ch_en.ttf,sha256=1YM0Z3XqLDjSNbF7ihQFSAIUdjF9m1rtHiNC_6QosTE,1594400
-vision_agent/image_utils.py,sha256=YvP5KE9NrWdgJKuHW2NR1glzfObkxtcXBknpmj3Gsbs,7554
+vision_agent/image_utils.py,sha256=BfUueXZKOvjdI_J7vi2vE57FfpaHOkxrTY0d3aQ1zgI,7552
 vision_agent/llm/__init__.py,sha256=BoUm_zSAKnLlE8s-gKTSQugXDqVZKPqYlWwlTLdhcz4,48
 vision_agent/llm/llm.py,sha256=1BkrSVBWEClyqLc0Rmyw4heLhi_ZVm6JO7-i1wd1ziw,5383
 vision_agent/lmm/__init__.py,sha256=nnNeKD1k7q_4vLb1x51O_EUTYaBgGfeiCx5F433gr3M,67
-vision_agent/lmm/lmm.py,sha256=sECjGMaGrv1QHq7OiFr-9LoBM5uRLjAqd0Ypp-zyFlw,10552
-vision_agent/tools/__init__.py,sha256=X6yJhWa8iKkQm4Mgf1KcV0_o39-Nrg3E56QAB5gWCO0,413
+vision_agent/lmm/lmm.py,sha256=gK90vMxh0OcGSuIZQikBkDXm4pfkdFk1R2y7rtWDl84,10539
+vision_agent/tools/__init__.py,sha256=HfUr0JQUwk0Kyieen93df9lMbbdpVf9Q6CcVFmKv_q4,413
 vision_agent/tools/prompts.py,sha256=V1z4YJLXZuUl_iZ5rY0M5hHc_2tmMEUKr0WocXKGt4E,1430
-vision_agent/tools/tools.py,sha256=hYgRTHMCBwjT0kkT2SY5MN0FK89vuuecu-x1VqRlGbU,42779
+vision_agent/tools/tools.py,sha256=DZX-w17XWttR4j8bQKY90QxLOs3-ZD5qOHA_53LL7Dk,43013
 vision_agent/tools/video.py,sha256=xTElFSFp1Jw4ulOMnk81Vxsh-9dTxcWUO6P9fzEi3AM,7653
 vision_agent/type_defs.py,sha256=4LTnTL4HNsfYqCrDn9Ppjg9bSG2ZGcoKSSd9YeQf4Bw,1792
-vision_agent-0.2.4.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-vision_agent-0.2.4.dist-info/METADATA,sha256=2T1YLGMh2-n8F0gGf1P2BDhgzxmtmAiylpfW3E3Q4_c,7697
-vision_agent-0.2.4.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
-vision_agent-0.2.4.dist-info/RECORD,,
+vision_agent-0.2.6.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+vision_agent-0.2.6.dist-info/METADATA,sha256=84i_O_9o8Ro6PbR3bi0rLyWbSwjqoxO6n6V9Bk06tP4,7697
+vision_agent-0.2.6.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
+vision_agent-0.2.6.dist-info/RECORD,,