Merge branch 'main' into update-comments

README.md

@@ -11,7 +11,7 @@
 
 Vision Agent is a library that helps you utilize agent frameworks for your vision tasks.
 Many current vision problems can easily take hours or days to solve, you need to find the
-right model, figure out how to use it, possibly write programming logic around it to 
+right model, figure out how to use it, possibly write programming logic around it to
 accomplish the task you want or even more expensive, train your own model. Vision Agent
 aims to provide an in-seconds experience by allowing users to describe their problem in
 text and utilizing agent frameworks to solve the task for them. Check out our discord
@@ -110,6 +110,9 @@ you. For example:
 | BboxIoU | BboxIoU returns the intersection over union of two bounding boxes normalized to 2 decimal places. |
 | SegIoU | SegIoU returns the intersection over union of two segmentation masks normalized to 2 decimal places. |
 | BoxDistance | BoxDistance returns the minimum distance between two bounding boxes normalized to 2 decimal places. |
+| ExtractFrames | ExtractFrames extracts frames with motion from a video. |
+| ZeroShotCounting | ZeroShotCounting returns the total number of objects belonging to a single class in a given image |
+| VisualPromptCounting | VisualPromptCounting returns the total number of objects belonging to a single class given an image and visual prompt |
 
 
 It also has a basic set of calculate tools such as add, subtract, multiply and divide.

mkdocs.yml

@@ -37,7 +37,5 @@ nav:
  - vision_agent.tools: api/tools.md
  - vision_agent.llm: api/llm.md
  - vision_agent.lmm: api/lmm.md
- - vision_agent.data: api/data.md
- - vision_agent.emb: api/emb.md
  - vision_agent.image_utils: api/image_utils.md
  - Old documentation: old.md

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "vision-agent"
-version = "0.1.4"
+version = "0.2.1"
 description = "Toolset for Vision Agent"
 authors = ["Landing AI <[email protected]>"]
 readme = "README.md"
@@ -16,16 +16,12 @@ packages = [{include = "vision_agent"}]
 "documentation" = "https://github.com/landing-ai/vision-agent"
 
 [tool.poetry.dependencies] # main dependency group
-python = ">=3.9,<3.12"
-
+python = ">=3.9"
 numpy = ">=1.21.0,<2.0.0"
 pillow = "10.*"
 requests = "2.*"
 tqdm = ">=4.64.0,<5.0.0"
 pandas = "2.*"
-faiss-cpu = "1.*"
-torch = "2.1.*" # 2.2 causes sentence-transformers to seg fault
-sentence-transformers = "2.*"
 openai = "1.*"
 typing_extensions = "4.*"
 moviepy = "1.*"

vision_agent/__init__.py

@@ -1,5 +1,3 @@
 from .agent import Agent
-from .data import DataStore, build_data_store
-from .emb import Embedder, OpenAIEmb, SentenceTransformerEmb, get_embedder
 from .llm import LLM, OpenAILLM
 from .lmm import LMM, LLaVALMM, OpenAILMM, get_lmm

vision_agent/agent/vision_agent.py

@@ -8,7 +8,7 @@
 from PIL import Image
 from tabulate import tabulate
 
-from vision_agent.image_utils import overlay_bboxes, overlay_masks
+from vision_agent.image_utils import overlay_bboxes, overlay_masks, overlay_heat_map
 from vision_agent.llm import LLM, OpenAILLM
 from vision_agent.lmm import LMM, OpenAILMM
 from vision_agent.tools import TOOLS
@@ -33,6 +33,7 @@
 
 logging.basicConfig(stream=sys.stdout)
 _LOGGER = logging.getLogger(__name__)
+_MAX_TABULATE_COL_WIDTH = 80
 
 
 def parse_json(s: str) -> Any:
@@ -335,7 +336,9 @@ def _handle_viz_tools(
 
  for param, call_result in zip(parameters, tool_result["call_results"]):
  # calls can fail, so we need to check if the call was successful
- if not isinstance(call_result, dict) or "bboxes" not in call_result:
+ if not isinstance(call_result, dict) or (
+ "bboxes" not in call_result and "masks" not in call_result
+ ):
  return image_to_data
 
  # if the call was successful, then we can add the image data
@@ -348,11 +351,12 @@ def _handle_viz_tools(
  "scores": [],
  }
 
- image_to_data[image]["bboxes"].extend(call_result["bboxes"])
- image_to_data[image]["labels"].extend(call_result["labels"])
- image_to_data[image]["scores"].extend(call_result["scores"])
- if "masks" in call_result:
- image_to_data[image]["masks"].extend(call_result["masks"])
+ image_to_data[image]["bboxes"].extend(call_result.get("bboxes", []))
+ 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", []))
+ if "mask_shape" in call_result:
+ image_to_data[image]["mask_shape"] = call_result["mask_shape"]
 
  return image_to_data
 
@@ -366,6 +370,8 @@ def visualize_result(all_tool_results: List[Dict]) -> Sequence[Union[str, Path]]
  "grounding_dino_",
  "extract_frames_",
  "dinov_",
+ "zero_shot_counting_",
+ "visual_prompt_counting_",
  ]:
  continue
 
@@ -378,8 +384,11 @@ def visualize_result(all_tool_results: List[Dict]) -> Sequence[Union[str, Path]]
  for image_str in image_to_data:
  image_path = Path(image_str)
  image_data = image_to_data[image_str]
- image = overlay_masks(image_path, image_data)
- image = overlay_bboxes(image, image_data)
+ if "_counting_" in tool_result["tool_name"]:
+ image = overlay_heat_map(image_path, image_data)
+ else:
+ image = overlay_masks(image_path, image_data)
+ image = overlay_bboxes(image, image_data)
  with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
  image.save(f.name)
  visualized_images.append(f.name)
@@ -498,11 +507,21 @@ def chat_with_workflow(
  if image:
  question += f" Image name: {image}"
  if reference_data:
- if not ("image" in reference_data and "mask" in reference_data):
+ if not (
+ "image" in reference_data
+ and ("mask" in reference_data or "bbox" in reference_data)
+ ):
  raise ValueError(
- f"Reference data must contain 'image' and 'mask'. but got {reference_data}"
+ f"Reference data must contain 'image' and a visual prompt which can be 'mask' or 'bbox'. but got {reference_data}"
  )
- question += f" Reference image: {reference_data['image']}, Reference mask: {reference_data['mask']}"
+ visual_prompt_data = (
+ f"Reference mask: {reference_data['mask']}"
+ if "mask" in reference_data
+ else f"Reference bbox: {reference_data['bbox']}"
+ )
+ question += (
+ f" Reference image: {reference_data['image']}, {visual_prompt_data}"
+ )
 
  reflections = ""
  final_answer = ""
@@ -545,7 +564,6 @@ def chat_with_workflow(
  final_answer = answer_summarize(
  self.answer_model, question, answers, reflections
  )
-
  visualized_output = visualize_result(all_tool_results)
  all_tool_results.append({"visualized_output": visualized_output})
  if len(visualized_output) > 0:
@@ -629,7 +647,7 @@ def retrieval(
 
  self.log_progress(
  f"""Going to run the following tool(s) in sequence:
-{tabulate([tool_results], headers="keys", tablefmt="mixed_grid")}"""
+{tabulate(tabular_data=[tool_results], headers="keys", tablefmt="mixed_grid", maxcolwidths=_MAX_TABULATE_COL_WIDTH)}"""
  )
 
  def parse_tool_results(result: Dict[str, Union[Dict, List]]) -> Any:
@@ -675,6 +693,6 @@ def create_tasks(
  task_list = []
  self.log_progress(
  f"""Planned tasks:
-{tabulate(task_list, headers="keys", tablefmt="mixed_grid")}"""
+{tabulate(task_list, headers="keys", tablefmt="mixed_grid", maxcolwidths=_MAX_TABULATE_COL_WIDTH)}"""
  )
  return task_list