Add detector tool (#13)

* added tool usage

* added llm

* added exception

* isort

* fix typing errors

* remove extra import

* fix typing errors

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):