vision-agent 0.1.6__py3-none-any.whl → 0.2.2__py3-none-any.whl

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,12 @@ from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 from PIL import Image
 from tabulate import tabulate
 
-from vision_agent.image_utils import overlay_bboxes, overlay_masks, overlay_heat_map
+from vision_agent.image_utils import (
+    convert_to_b64,
+    overlay_bboxes,
+    overlay_heat_map,
+    overlay_masks,
+)
 from vision_agent.llm import LLM, OpenAILLM
 from vision_agent.lmm import LMM, OpenAILMM
 from vision_agent.tools import TOOLS
@@ -481,6 +486,17 @@ class VisionAgent(Agent):
         if self.report_progress_callback:
             self.report_progress_callback(description)
 
+    def _report_visualization_via_callback(
+        self, images: Sequence[Union[str, Path]]
+    ) -> None:
+        """This is intended for streaming the visualization images via the callback to the client side."""
+        if self.report_progress_callback:
+            self.report_progress_callback("<VIZ>")
+            if images:
+                for img in images:
+                    self.report_progress_callback(f"<IMG>{convert_to_b64(img)}</IMG>")
+            self.report_progress_callback("</VIZ>")
+
     def chat_with_workflow(
         self,
         chat: List[Dict[str, str]],
@@ -577,9 +593,12 @@ class VisionAgent(Agent):
         )
 
         if visualize_output:
-            visualized_output = all_tool_results[-1]["visualized_output"]
-            for image in visualized_output:
-                Image.open(image).show()
+            viz_images: Sequence[Union[str, Path]] = all_tool_results[-1][
+                "visualized_output"
+            ]
+            self._report_visualization_via_callback(viz_images)
+            for img in viz_images:
+                Image.open(img).show()
 
         return final_answer, all_tool_results
 

vision_agent/image_utils.py

@@ -4,7 +4,7 @@ import base64
 from importlib import resources
 from io import BytesIO
 from pathlib import Path
-from typing import Dict, Tuple, Union, List
+from typing import Dict, List, Tuple, Union
 
 import numpy as np
 from PIL import Image, ImageDraw, ImageFont
@@ -108,7 +108,7 @@ def convert_to_b64(data: Union[str, Path, np.ndarray, ImageType]) -> str:
         data = Image.open(data)
     if isinstance(data, Image.Image):
         buffer = BytesIO()
-        data.convert("RGB").save(buffer, format="JPEG")
+        data.convert("RGB").save(buffer, format="PNG")
         return base64.b64encode(buffer.getvalue()).decode("utf-8")
     else:
         arr_bytes = data.tobytes()

{vision_agent-0.1.6.dist-info → vision_agent-0.2.2.dist-info}/METADATA

@@ -1,15 +1,14 @@
 Metadata-Version: 2.1
 Name: vision-agent
-Version: 0.1.6
+Version: 0.2.2
 Summary: Toolset for Vision Agent
 Author: Landing AI
 Author-email: dev@landing.ai
-Requires-Python: >=3.9,<3.12
+Requires-Python: >=3.9
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: faiss-cpu (>=1.0.0,<2.0.0)
 Requires-Dist: moviepy (>=1.0.0,<2.0.0)
 Requires-Dist: numpy (>=1.21.0,<2.0.0)
 Requires-Dist: openai (>=1.0.0,<2.0.0)
@@ -18,9 +17,7 @@ Requires-Dist: pandas (>=2.0.0,<3.0.0)
 Requires-Dist: pillow (>=10.0.0,<11.0.0)
 Requires-Dist: pydantic-settings (>=2.2.1,<3.0.0)
 Requires-Dist: requests (>=2.0.0,<3.0.0)
-Requires-Dist: sentence-transformers (>=2.0.0,<3.0.0)
 Requires-Dist: tabulate (>=0.9.0,<0.10.0)
-Requires-Dist: torch (>=2.1.0,<2.2.0)
 Requires-Dist: tqdm (>=4.64.0,<5.0.0)
 Requires-Dist: typing_extensions (>=4.0.0,<5.0.0)
 Project-URL: Homepage, https://landing.ai

{vision_agent-0.1.6.dist-info → vision_agent-0.2.2.dist-info}/RECORD

@@ -1,19 +1,15 @@
-vision_agent/__init__.py,sha256=wD1cssVTAJ55uTViNfBGooqJUV0p9fmVAuTMHHrmUBU,229
+vision_agent/__init__.py,sha256=GVLHCeK_R-zgldpbcPmOzJat-BkadvkuRCMxDvTIcXs,108
 vision_agent/agent/__init__.py,sha256=B4JVrbY4IRVCJfjmrgvcp7h1mTUEk8MZvL0Zmej4Ka0,127
 vision_agent/agent/agent.py,sha256=X7kON-g9ePUKumCDaYfQNBX_MEFE-ax5PnRp7-Cc5Wo,529
 vision_agent/agent/easytool.py,sha256=oMHnBg7YBtIPgqQUNcZgq7uMgpPThs99_UnO7ERkMVg,11511
 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=MTxeV5_Sghqoe2aOW9EbNgiq61sVCcF3ZndJ7BZl6x0,23588
+vision_agent/agent/vision_agent.py,sha256=2VUMRVI6KAnmaUK-34wrgyfSQ2DAUm4g4QQcpqa2zao,24235
 vision_agent/agent/vision_agent_prompts.py,sha256=W3Z72FpUt71UIJSkjAcgtQqxeMqkYuATqHAN5fYY26c,7342
-vision_agent/data/__init__.py,sha256=YU-5g3LbEQ6a4drz0RLGTagXMVU2Z4Xr3RlfWE-R0jU,46
-vision_agent/data/data.py,sha256=Z2l76OrT0GgyuN52OeJqDitUcP0q1rhfdXd1of3GsVo,5128
-vision_agent/emb/__init__.py,sha256=YmCkGrJBtXb6X6Z3lnKiFoQYKXMgHMJp8JJyMLVvqcI,75
-vision_agent/emb/emb.py,sha256=la9lhEzk7jqUCjYYQ5oRgVNSnC9_EJBJIpE_B9c6PJo,1375
 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=Cg4aKO1tQiETT1gdsZ50XzORBtJnBFfMG2cKJyjaY6Q,7555
+vision_agent/image_utils.py,sha256=YvP5KE9NrWdgJKuHW2NR1glzfObkxtcXBknpmj3Gsbs,7554
 vision_agent/llm/__init__.py,sha256=BoUm_zSAKnLlE8s-gKTSQugXDqVZKPqYlWwlTLdhcz4,48
 vision_agent/llm/llm.py,sha256=gwDQ9-p9wEn24xi1019e5jzTGQg4xWDSqBCsqIqGcU4,5168
 vision_agent/lmm/__init__.py,sha256=nnNeKD1k7q_4vLb1x51O_EUTYaBgGfeiCx5F433gr3M,67
@@ -23,7 +19,7 @@ vision_agent/tools/prompts.py,sha256=V1z4YJLXZuUl_iZ5rY0M5hHc_2tmMEUKr0WocXKGt4E
 vision_agent/tools/tools.py,sha256=gCjHs5vJuGNBFsnJWFT7PX3wTyfHgtrgX1Eq9vqknN0,34979
 vision_agent/tools/video.py,sha256=xTElFSFp1Jw4ulOMnk81Vxsh-9dTxcWUO6P9fzEi3AM,7653
 vision_agent/type_defs.py,sha256=4LTnTL4HNsfYqCrDn9Ppjg9bSG2ZGcoKSSd9YeQf4Bw,1792
-vision_agent-0.1.6.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-vision_agent-0.1.6.dist-info/METADATA,sha256=Ig2tSKyeH8a2A8xZRq72M9XnKyi4_03UM4hDiFpT-eU,6574
-vision_agent-0.1.6.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
-vision_agent-0.1.6.dist-info/RECORD,,
+vision_agent-0.2.2.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+vision_agent-0.2.2.dist-info/METADATA,sha256=dOZ9KWmhuVb5wvschxYBis8x79HwgOD3MmTKqyupggg,6434
+vision_agent-0.2.2.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
+vision_agent-0.2.2.dist-info/RECORD,,

vision_agent/data/__init__.py

@@ -1 +0,0 @@
-from .data import DataStore, build_data_store

vision_agent/data/data.py

@@ -1,142 +0,0 @@
-from __future__ import annotations
-
-import uuid
-from pathlib import Path
-from typing import Callable, Dict, List, Optional, Union, cast
-
-import faiss
-import numpy as np
-import numpy.typing as npt
-import pandas as pd
-from faiss import read_index, write_index
-from tqdm import tqdm
-from typing_extensions import Self
-
-from vision_agent.emb import Embedder
-from vision_agent.lmm import LMM
-
-tqdm.pandas()
-
-
-class DataStore:
-    r"""A class to store and manage image data along with its generated metadata from an LMM."""
-
-    def __init__(self, df: pd.DataFrame):
-        r"""Initializes the DataStore with a DataFrame containing image paths and image
-        IDs. If the image IDs are not present, they are generated using UUID4. The
-        DataFrame must contain an 'image_paths' column.
-
-        Args:
-            df: The DataFrame containing "image_paths" and "image_id" columns.
-        """
-        self.df = df
-        self.lmm: Optional[LMM] = None
-        self.emb: Optional[Embedder] = None
-        self.index: Optional[faiss.IndexFlatIP] = None  # type: ignore
-        if "image_paths" not in self.df.columns:
-            raise ValueError("image_paths column must be present in DataFrame")
-        if "image_id" not in self.df.columns:
-            self.df["image_id"] = [str(uuid.uuid4()) for _ in range(len(df))]
-
-    def add_embedder(self, emb: Embedder) -> Self:
-        self.emb = emb
-        return self
-
-    def add_lmm(self, lmm: LMM) -> Self:
-        self.lmm = lmm
-        return self
-
-    def add_column(
-        self, name: str, prompt: str, func: Optional[Callable[[str], str]] = None
-    ) -> Self:
-        r"""Adds a new column to the DataFrame containing the generated metadata from
-        the LMM.
-
-        Args:
-            name: The name of the column to be added.
-            prompt: The prompt to be used to generate the metadata.
-            func: A Python function to be applied on the output of `lmm.generate`.
-                Defaults to None.
-        """
-        if self.lmm is None:
-            raise ValueError("LMM not set yet")
-
-        self.df[name] = self.df["image_paths"].progress_apply(  # type: ignore
-            lambda x: (
-                func(self.lmm.generate(prompt, images=[x]))
-                if func
-                else self.lmm.generate(prompt, images=[x])
-            )
-        )
-        return self
-
-    def build_index(self, target_col: str) -> Self:
-        r"""This will generate embeddings for the `target_col` and build a searchable
-        index over them, so next time you run search it will search over this index.
-
-        Args:
-            target_col: The column name containing the data to be indexed."""
-        if self.emb is None:
-            raise ValueError("Embedder not set yet")
-
-        embeddings: pd.Series = self.df[target_col].progress_apply(lambda x: self.emb.embed(x))  # type: ignore
-        embeddings_np = np.array(embeddings.tolist()).astype(np.float32)
-        self.index = faiss.IndexFlatIP(embeddings_np.shape[1])
-        self.index.add(embeddings_np)
-        return self
-
-    def get_embeddings(self) -> npt.NDArray[np.float32]:
-        if self.index is None:
-            raise ValueError("Index not built yet")
-
-        ntotal = self.index.ntotal
-        d: int = self.index.d
-        return cast(
-            npt.NDArray[np.float32],
-            faiss.rev_swig_ptr(self.index.get_xb(), ntotal * d).reshape(ntotal, d),
-        )
-
-    def search(self, query: str, top_k: int = 10) -> List[Dict]:
-        r"""Searches the index for the most similar images to the query and returns
-        the top_k results.
-
-        Args:
-            query: The query to search for.
-            top_k: The number of results to return. Defaults to 10."""
-        if self.index is None:
-            raise ValueError("Index not built yet")
-        if self.emb is None:
-            raise ValueError("Embedder not set yet")
-
-        query_embedding: npt.NDArray[np.float32] = self.emb.embed(query)
-        _, idx = self.index.search(query_embedding.reshape(1, -1), top_k)
-        return cast(List[Dict], self.df.iloc[idx[0]].to_dict(orient="records"))
-
-    def save(self, path: Union[str, Path]) -> None:
-        path = Path(path)
-        path.mkdir(parents=True)
-        self.df.to_csv(path / "data.csv")
-        if self.index is not None:
-            write_index(self.index, str(path / "data.index"))
-
-    @classmethod
-    def load(cls, path: Union[str, Path]) -> DataStore:
-        path = Path(path)
-        df = pd.read_csv(path / "data.csv", index_col=0)
-        ds = DataStore(df)
-        if Path(path / "data.index").exists():
-            ds.index = read_index(str(path / "data.index"))
-        return ds
-
-
-def build_data_store(data: Union[str, Path, list[Union[str, Path]]]) -> DataStore:
-    if isinstance(data, Path) or isinstance(data, str):
-        data = Path(data)
-        data_files = list(Path(data).glob("*"))
-    elif isinstance(data, list):
-        data_files = [Path(d) for d in data]
-
-    df = pd.DataFrame()
-    df["image_paths"] = data_files
-    df["image_id"] = [uuid.uuid4() for _ in range(len(data_files))]
-    return DataStore(df)

vision_agent/emb/__init__.py

@@ -1 +0,0 @@
-from .emb import Embedder, OpenAIEmb, SentenceTransformerEmb, get_embedder

vision_agent/emb/emb.py

@@ -1,47 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import cast
-
-import numpy as np
-import numpy.typing as npt
-
-
-class Embedder(ABC):
-    @abstractmethod
-    def embed(self, text: str) -> npt.NDArray[np.float32]:
-        pass
-
-
-class SentenceTransformerEmb(Embedder):
-    def __init__(self, model_name: str = "BAAI/bge-small-en-v1.5"):
-        from sentence_transformers import SentenceTransformer
-
-        self.model = SentenceTransformer(model_name)
-
-    def embed(self, text: str) -> npt.NDArray[np.float32]:
-        return cast(
-            npt.NDArray[np.float32],
-            self.model.encode([text]).flatten().astype(np.float32),
-        )
-
-
-class OpenAIEmb(Embedder):
-    def __init__(self, model_name: str = "text-embedding-3-small"):
-        from openai import OpenAI
-
-        self.client = OpenAI()
-        self.model_name = model_name
-
-    def embed(self, text: str) -> npt.NDArray[np.float32]:
-        response = self.client.embeddings.create(input=text, model=self.model_name)
-        return np.array(response.data[0].embedding).astype(np.float32)
-
-
-def get_embedder(name: str) -> Embedder:
-    if name == "sentence-transformer":
-        return SentenceTransformerEmb()
-    elif name == "openai":
-        return OpenAIEmb()
-    else:
-        raise ValueError(
-            f"Unknown embedder name: {name}, currently support sentence-transformer, openai."
-        )