Add detector tool

vision_agent/__init__.py

@@ -1,3 +1,4 @@
-from .lmm import LMM, LLaVALMM, OpenAILMM, get_lmm
-from .emb import Embedder, SentenceTransformerEmb, OpenAIEmb, get_embedder
 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/data/data.py

@@ -2,7 +2,7 @@
 
 import uuid
 from pathlib import Path
-from typing import Dict, List, Optional, Union, cast, Callable
+from typing import Callable, Dict, List, Optional, Union, cast
 
 import faiss
 import numpy as np

vision_agent/image_utils.py

@@ -1,20 +1,31 @@
 import base64
 from io import BytesIO
 from pathlib import Path
-from typing import Union
+from typing import Tuple, Union
 
 import numpy as np
 from PIL import Image
+from PIL.Image import Image as ImageType
 
 
-def b64_to_pil(b64_str: str) -> Image.Image:
+def b64_to_pil(b64_str: str) -> ImageType:
  # , can't be encoded in b64 data so must be part of prefix
  if "," in b64_str:
  b64_str = b64_str.split(",")[1]
  return Image.open(BytesIO(base64.b64decode(b64_str)))
 
 
-def convert_to_b64(data: Union[str, Path, np.ndarray, Image.Image]) -> str:
+def get_image_size(data: Union[str, Path, np.ndarray, ImageType]) -> Tuple[int, ...]:
+ if isinstance(data, (str, Path)):
+ data = Image.open(data)
+
+ if isinstance(data, Image.Image):
+ return data.size[::-1]
+ else:
+ return data.shape[:2]
+
+
+def convert_to_b64(data: Union[str, Path, np.ndarray, ImageType]) -> str:
  if data is None:
  raise ValueError(f"Invalid input image: {data}. Input image can't be None.")
  if isinstance(data, (str, Path)):

vision_agent/llm/__init__.py

@@ -0,0 +1 @@
+from .llm import LLM, OpenAILLM

vision_agent/llm/llm.py

@@ -0,0 +1,86 @@
+import json
+from abc import ABC, abstractmethod
+from typing import Mapping, cast
+
+from vision_agent.tools import (
+ CHOOSE_PARAMS,
+ CLIP,
+ SYSTEM_PROMPT,
+ GroundingDINO,
+ GroundingSAM,
+ ImageTool,
+)
+
+
+class LLM(ABC):
+ @abstractmethod
+ def generate(self, prompt: str) -> str:
+ pass
+
+
+class OpenAILLM(LLM):
+ r"""An LLM class for any OpenAI LLM model."""
+
+ def __init__(self, model_name: str = "gpt-4-turbo-preview"):
+ from openai import OpenAI
+
+ self.model_name = model_name
+ self.client = OpenAI()
+
+ def generate(self, prompt: str) -> str:
+ response = self.client.chat.completions.create(
+ model=self.model_name,
+ messages=[
+ {"role": "user", "content": prompt},
+ ],
+ )
+
+ return cast(str, response.choices[0].message.content)
+
+ def generate_classifier(self, prompt: str) -> ImageTool:
+ prompt = CHOOSE_PARAMS.format(api_doc=CLIP.doc, question=prompt)
+ response = self.client.chat.completions.create(
+ model=self.model_name,
+ response_format={"type": "json_object"},
+ messages=[
+ {"role": "system", "content": SYSTEM_PROMPT},
+ {"role": "user", "content": prompt},
+ ],
+ )
+
+ params = json.loads(cast(str, response.choices[0].message.content))[
+ "Parameters"
+ ]
+ return CLIP(**cast(Mapping, params))
+
+ def generate_detector(self, params: str) -> ImageTool:
+ params = CHOOSE_PARAMS.format(api_doc=GroundingDINO.doc, question=params)
+ response = self.client.chat.completions.create(
+ model=self.model_name,
+ response_format={"type": "json_object"},
+ messages=[
+ {"role": "system", "content": SYSTEM_PROMPT},
+ {"role": "user", "content": params},
+ ],
+ )
+
+ params = json.loads(cast(str, response.choices[0].message.content))[
+ "Parameters"
+ ]
+ return GroundingDINO(**cast(Mapping, params))
+
+ def generate_segmentor(self, params: str) -> ImageTool:
+ params = CHOOSE_PARAMS.format(api_doc=GroundingSAM.doc, question=params)
+ response = self.client.chat.completions.create(
+ model=self.model_name,
+ response_format={"type": "json_object"},
+ messages=[
+ {"role": "system", "content": SYSTEM_PROMPT},
+ {"role": "user", "content": params},
+ ],
+ )
+
+ params = json.loads(cast(str, response.choices[0].message.content))[
+ "Parameters"
+ ]
+ return GroundingSAM(**cast(Mapping, params))

vision_agent/lmm/lmm.py

@@ -3,7 +3,7 @@
 import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Mapping, Optional, Union, cast
 
 import requests
 
@@ -38,8 +38,8 @@ def generate(self, prompt: str, image: Optional[Union[str, Path]] = None) -> str
 class LLaVALMM(LMM):
  r"""An LMM class for the LLaVA-1.6 34B model."""
 
- def __init__(self, name: str):
- self.name = name
+ def __init__(self, model_name: str):
+ self.model_name = model_name
 
  def generate(
  self,
@@ -67,10 +67,10 @@ def generate(
 class OpenAILMM(LMM):
  r"""An LMM class for the OpenAI GPT-4 Vision model."""
 
- def __init__(self, name: str):
+ def __init__(self, model_name: str = "gpt-4-vision-preview"):
  from openai import OpenAI
 
- self.name = name
+ self.model_name = model_name
  self.client = OpenAI()
 
  def generate(self, prompt: str, image: Optional[Union[str, Path]] = None) -> str:
@@ -96,15 +96,14 @@ def generate(self, prompt: str, image: Optional[Union[str, Path]] = None) -> str
  )
 
  response = self.client.chat.completions.create(
- model="gpt-4-vision-preview", messages=message # type: ignore
+ model=self.model_name, messages=message # type: ignore
  )
  return cast(str, response.choices[0].message.content)
 
  def generate_classifier(self, prompt: str) -> ImageTool:
  prompt = CHOOSE_PARAMS.format(api_doc=CLIP.doc, question=prompt)
  response = self.client.chat.completions.create(
- model="gpt-4-turbo-preview", # no need to use vision model here
- response_format={"type": "json_object"},
+ model=self.model_name,
  messages=[
  {"role": "system", "content": SYSTEM_PROMPT},
  {"role": "user", "content": prompt},
@@ -113,44 +112,42 @@ def generate_classifier(self, prompt: str) -> ImageTool:
 
  try:
  prompt = json.loads(cast(str, response.choices[0].message.content))[
- "prompt"
+ "Parameters"
  ]
  except json.JSONDecodeError:
  _LOGGER.error(
  f"Failed to decode response: {response.choices[0].message.content}"
  )
  raise ValueError("Failed to decode response")
 
- return CLIP(prompt)
+ return CLIP(**cast(Mapping, prompt))
 
- def generate_detector(self, prompt: str) -> ImageTool:
- prompt = CHOOSE_PARAMS.format(api_doc=GroundingDINO.doc, question=prompt)
+ def generate_detector(self, params: str) -> ImageTool:
+ params = CHOOSE_PARAMS.format(api_doc=GroundingDINO.doc, question=params)
  response = self.client.chat.completions.create(
- model="gpt-4-turbo-preview", # no need to use vision model here
- response_format={"type": "json_object"},
+ model=self.model_name,
  messages=[
  {"role": "system", "content": SYSTEM_PROMPT},
- {"role": "user", "content": prompt},
+ {"role": "user", "content": params},
  ],
  )
 
  try:
- prompt = json.loads(cast(str, response.choices[0].message.content))[
- "prompt"
+ params = json.loads(cast(str, response.choices[0].message.content))[
+ "Parameters"
  ]
  except json.JSONDecodeError:
  _LOGGER.error(
  f"Failed to decode response: {response.choices[0].message.content}"
  )
  raise ValueError("Failed to decode response")
 
- return GroundingDINO(prompt)
+ return GroundingDINO(**cast(Mapping, params))
 
  def generate_segmentor(self, prompt: str) -> ImageTool:
  prompt = CHOOSE_PARAMS.format(api_doc=GroundingSAM.doc, question=prompt)
  response = self.client.chat.completions.create(
- model="gpt-4-turbo-preview", # no need to use vision model here
- response_format={"type": "json_object"},
+ model=self.model_name,
  messages=[
  {"role": "system", "content": SYSTEM_PROMPT},
  {"role": "user", "content": prompt},
@@ -159,15 +156,15 @@ def generate_segmentor(self, prompt: str) -> ImageTool:
 
  try:
  prompt = json.loads(cast(str, response.choices[0].message.content))[
- "prompt"
+ "Parameters"
  ]
  except json.JSONDecodeError:
  _LOGGER.error(
  f"Failed to decode response: {response.choices[0].message.content}"
  )
  raise ValueError("Failed to decode response")
 
- return GroundingSAM(prompt)
+ return GroundingSAM(**cast(Mapping, prompt))
 
 
 def get_lmm(name: str) -> LMM:

vision_agent/tools/__init__.py

@@ -1,2 +1,2 @@
-from .prompts import SYSTEM_PROMPT, CHOOSE_PARAMS
-from .tools import ImageTool, CLIP, GroundingDINO, GroundingSAM
+from .prompts import CHOOSE_PARAMS, SYSTEM_PROMPT
+from .tools import CLIP, GroundingDINO, GroundingSAM, ImageTool

vision_agent/tools/tools.py

@@ -1,16 +1,28 @@
 import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, Dict, List, Union, cast
+from typing import Any, Dict, List, Tuple, Union, cast
 
 import requests
 from PIL.Image import Image as ImageType
 
-from vision_agent.image_utils import convert_to_b64
+from vision_agent.image_utils import convert_to_b64, get_image_size
 
 _LOGGER = logging.getLogger(__name__)
 
 
+def normalize_bbox(
+ bbox: List[Union[int, float]], image_size: Tuple[int, ...]
+) -> List[float]:
+ r"""Normalize the bounding box coordinates to be between 0 and 1."""
+ x1, y1, x2, y2 = bbox
+ x1 = x1 / image_size[1]
+ y1 = y1 / image_size[0]
+ x2 = x2 / image_size[1]
+ y2 = y2 / image_size[0]
+ return [x1, y1, x2, y2]
+
+
 class ImageTool(ABC):
  @abstractmethod
  def __call__(self, image: Union[str, ImageType]) -> List[Dict]:
@@ -42,12 +54,18 @@ class GroundingDINO(ImageTool):
  'Example 1: User Question: "Can you build me a car detector?" {{"Parameters":{{"prompt": "car"}}}}\n'
  'Example 2: User Question: "Can you detect the person on the left?" {{"Parameters":{{"prompt": "person on the left"}}\n'
  'Exmaple 3: User Question: "Can you build me a tool that detects red shirts and green shirts?" {{"Parameters":{{"prompt": "red shirt. green shirt"}}}}\n'
+ "The tool returns a list of dictionaries, each containing the following keys:\n"
+ " - 'lable': The label of the detected object.\n"
+ " - 'score': The confidence score of the detection.\n"
+ " - 'bbox': The bounding box of the detected object. The box coordinates are normalize to [0, 1]\n"
+ "An example output would be: [{'label': ['car'], 'score': [0.99], 'bbox': [[0.1, 0.2, 0.3, 0.4]]}]\n"
  )
 
  def __init__(self, prompt: str):
  self.prompt = prompt
 
  def __call__(self, image: Union[str, Path, ImageType]) -> List[Dict]:
+ image_size = get_image_size(image)
  image_b64 = convert_to_b64(image)
  data = {
  "prompt": self.prompt,
@@ -59,9 +77,18 @@ def __call__(self, image: Union[str, Path, ImageType]) -> List[Dict]:
  json=data,
  )
  resp_json: Dict[str, Any] = res.json()
- if resp_json["statusCode"] != 200:
+ if (
+ "statusCode" in resp_json and resp_json["statusCode"] != 200
+ ) or "statusCode" not in resp_json:
  _LOGGER.error(f"Request failed: {resp_json}")
- return cast(List[Dict], resp_json["data"])
+ raise ValueError(f"Request failed: {resp_json}")
+ resp_data = resp_json["data"]
+ for elt in resp_data:
+ if "bboxes" in elt:
+ elt["bboxes"] = [
+ normalize_bbox(box, image_size) for box in elt["bboxes"]
+ ]
+ return cast(List[Dict], resp_data)
 
 
 class GroundingSAM(ImageTool):