vision-agent 0.2.10__py3-none-any.whl → 0.2.22__py3-none-any.whl

vision_agent/tools/tools.py

@@ -9,8 +9,12 @@ import numpy as np
 import requests
 from PIL import Image
 from PIL.Image import Image as ImageType
+from scipy.spatial import distance  # type: ignore
 
-from vision_agent.image_utils import (
+from vision_agent.lmm import OpenAILMM
+from vision_agent.tools.tool_utils import _send_inference_request
+from vision_agent.utils import extract_frames_from_video
+from vision_agent.utils.image_utils import (
     b64_to_pil,
     convert_to_b64,
     denormalize_bbox,
@@ -18,13 +22,8 @@ from vision_agent.image_utils import (
     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
 
 _LOGGER = logging.getLogger(__name__)
-_LND_API_KEY = LandingaiAPIKey().api_key
-_LND_API_URL = "https://api.dev.landing.ai/v1/agent"
 
 
 class Tool(ABC):
@@ -175,15 +174,15 @@ class GroundingDINO(Tool):
     """
 
     name = "grounding_dino_"
-    description = "'grounding_dino_' is a tool that can detect and count objects given a text prompt such as category names or referring expressions. It returns a list and count of bounding boxes, label names and associated probability scores."
+    description = "'grounding_dino_' is a tool that can detect and count multiple objects given a text prompt such as category names or referring expressions. It returns a list and count of bounding boxes, label names and associated probability scores."
     usage = {
         "required_parameters": [
             {"name": "prompt", "type": "str"},
             {"name": "image", "type": "str"},
         ],
         "optional_parameters": [
-            {"name": "box_threshold", "type": "float"},
-            {"name": "iou_threshold", "type": "float"},
+            {"name": "box_threshold", "type": "float", "min": 0.1, "max": 0.5},
+            {"name": "iou_threshold", "type": "float", "min": 0.01, "max": 0.99},
         ],
         "examples": [
             {
@@ -210,7 +209,7 @@ class GroundingDINO(Tool):
                     "prompt": "red shirt. green shirt",
                     "image": "shirts.jpg",
                     "box_threshold": 0.20,
-                    "iou_threshold": 0.75,
+                    "iou_threshold": 0.20,
                 },
             },
         ],
@@ -222,7 +221,7 @@ class GroundingDINO(Tool):
         prompt: str,
         image: Union[str, Path, ImageType],
         box_threshold: float = 0.20,
-        iou_threshold: float = 0.75,
+        iou_threshold: float = 0.20,
     ) -> Dict:
         """Invoke the Grounding DINO model.
 
@@ -250,7 +249,7 @@ class GroundingDINO(Tool):
             data["scores"] = [round(score, 2) for score in data["scores"]]
         if "labels" in data:
             data["labels"] = list(data["labels"])
-        data["size"] = (image_size[1], image_size[0])
+        data["image_size"] = image_size
         return data
 
 
@@ -278,15 +277,15 @@ class GroundingSAM(Tool):
     """
 
     name = "grounding_sam_"
-    description = "'grounding_sam_' is a tool that can detect and segment objects given a text prompt such as category names or referring expressions. It returns a list of bounding boxes, label names and masks file names and associated probability scores."
+    description = "'grounding_sam_' is a tool that can detect and segment multiple objects given a text prompt such as category names or referring expressions. It returns a list of bounding boxes, label names and masks file names and associated probability scores."
     usage = {
         "required_parameters": [
             {"name": "prompt", "type": "str"},
             {"name": "image", "type": "str"},
         ],
         "optional_parameters": [
-            {"name": "box_threshold", "type": "float"},
-            {"name": "iou_threshold", "type": "float"},
+            {"name": "box_threshold", "type": "float", "min": 0.1, "max": 0.5},
+            {"name": "iou_threshold", "type": "float", "min": 0.01, "max": 0.99},
         ],
         "examples": [
             {
@@ -313,7 +312,7 @@ class GroundingSAM(Tool):
                     "prompt": "red shirt, green shirt",
                     "image": "shirts.jpg",
                     "box_threshold": 0.20,
-                    "iou_threshold": 0.75,
+                    "iou_threshold": 0.20,
                 },
             },
         ],
@@ -325,7 +324,7 @@ class GroundingSAM(Tool):
         prompt: str,
         image: Union[str, ImageType],
         box_threshold: float = 0.2,
-        iou_threshold: float = 0.75,
+        iou_threshold: float = 0.2,
     ) -> Dict:
         """Invoke the Grounding SAM model.
 
@@ -354,6 +353,7 @@ class GroundingSAM(Tool):
                 rle_decode(mask_rle=mask, shape=data["mask_shape"])
                 for mask in data["masks"]
             ]
+        data["image_size"] = image_size
         data.pop("mask_shape", None)
         return data
 
@@ -423,7 +423,6 @@ class DINOv(Tool):
         request_data = {
             "prompt": prompt,
             "image": image_b64,
-            "tool": "dinov",
         }
         data: Dict[str, Any] = _send_inference_request(request_data, "dinov")
         if "bboxes" in data:
@@ -436,6 +435,8 @@ class DINOv(Tool):
                 for mask in data["masks"]
             ]
         data["labels"] = ["visual prompt" for _ in range(len(data["masks"]))]
+        mask_shape = data.pop("mask_shape", None)
+        data["image_size"] = (mask_shape[0], mask_shape[1]) if mask_shape else None
         return data
 
 
@@ -544,7 +545,7 @@ class VisualPromptCounting(Tool):
     -------
         >>> import vision_agent as va
         >>> prompt_count = va.tools.VisualPromptCounting()
-        >>> prompt_count(image="image1.jpg", prompt="0.1, 0.1, 0.4, 0.42")
+        >>> prompt_count(image="image1.jpg", prompt={"bbox": [0.1, 0.1, 0.4, 0.42]})
         {'count': 23}
     """
 
@@ -554,52 +555,60 @@ class VisualPromptCounting(Tool):
     usage = {
         "required_parameters": [
             {"name": "image", "type": "str"},
-            {"name": "prompt", "type": "str"},
+            {"name": "prompt", "type": "Dict[str, List[float]"},
         ],
         "examples": [
             {
                 "scenario": "Here is an example of a lid '0.1, 0.1, 0.14, 0.2', Can you count the items in the image ? Image name: lids.jpg",
-                "parameters": {"image": "lids.jpg", "prompt": "0.1, 0.1, 0.14, 0.2"},
+                "parameters": {
+                    "image": "lids.jpg",
+                    "prompt": {"bbox": [0.1, 0.1, 0.14, 0.2]},
+                },
             },
             {
-                "scenario": "Can you count the total number of objects in this image ? Image name: tray.jpg",
-                "parameters": {"image": "tray.jpg", "prompt": "0.1, 0.1, 0.2, 0.25"},
+                "scenario": "Can you count the total number of objects in this image ? Image name: tray.jpg, reference_data: {'bbox': [0.1, 0.1, 0.2, 0.25]}",
+                "parameters": {
+                    "image": "tray.jpg",
+                    "prompt": {"bbox": [0.1, 0.1, 0.2, 0.25]},
+                },
             },
             {
-                "scenario": "Can you count this item based on an example, reference_data: '0.1, 0.15, 0.2, 0.2' ? Image name: shirts.jpg",
+                "scenario": "Can you count this item based on an example, reference_data: {'bbox': [100, 115, 200, 200]} ? Image name: shirts.jpg",
                 "parameters": {
                     "image": "shirts.jpg",
-                    "prompt": "0.1, 0.15, 0.2, 0.2",
+                    "prompt": {"bbox": [100, 115, 200, 200]},
                 },
             },
             {
-                "scenario": "Can you build me a counting tool based on an example prompt ? Image name: shoes.jpg",
+                "scenario": "Can you build me a counting tool based on an example prompt ? Image name: shoes.jpg, reference_data: {'bbox': [0.1, 0.1, 0.6, 0.65]}",
                 "parameters": {
                     "image": "shoes.jpg",
-                    "prompt": "0.1, 0.1, 0.6, 0.65",
+                    "prompt": {"bbox": [0.1, 0.1, 0.6, 0.65]},
                 },
             },
         ],
     }
 
-    # TODO: Add support for input multiple images, which aligns with the output type.
-    def __call__(self, image: Union[str, ImageType], prompt: str) -> Dict:
+    def __call__(
+        self, image: Union[str, ImageType], prompt: Dict[str, List[float]]
+    ) -> Dict:
         """Invoke the few shot counting model.
 
         Parameters:
             image: the input image.
+            prompt: the visual prompt which is a bounding box describing the object.
 
         Returns:
             A dictionary containing the key 'count' and the count as value. E.g. {count: 12}
         """
         image_size = get_image_size(image)
-        bbox = [float(x) for x in prompt.split(",")]
-        prompt = ", ".join(map(str, denormalize_bbox(bbox, image_size)))
+        bbox = prompt["bbox"]
+        bbox_str = ", ".join(map(str, denormalize_bbox(bbox, image_size)))
         image_b64 = convert_to_b64(image)
 
         data = {
             "image": image_b64,
-            "prompt": prompt,
+            "prompt": bbox_str,
             "tool": "few_shot_counting",
         }
         resp_data = _send_inference_request(data, "tools")
@@ -783,33 +792,49 @@ class Crop(Tool):
         return {"image": tmp.name}
 
 
-class BboxArea(Tool):
-    r"""BboxArea returns the area of the bounding box in pixels normalized to 2 decimal places."""
+class BboxStats(Tool):
+    r"""BboxStats returns the height, width and area of the bounding box in pixels to 2 decimal places."""
 
-    name = "bbox_area_"
-    description = "'bbox_area_' returns the area of the given bounding box in pixels normalized to 2 decimal places."
+    name = "bbox_stats_"
+    description = "'bbox_stats_' returns the height, width and area of the given bounding box in pixels to 2 decimal places."
     usage = {
-        "required_parameters": [{"name": "bboxes", "type": "List[int]"}],
+        "required_parameters": [
+            {"name": "bboxes", "type": "List[int]"},
+            {"name": "image_size", "type": "Tuple[int]"},
+        ],
         "examples": [
             {
-                "scenario": "If you want to calculate the area of the bounding box [0.2, 0.21, 0.34, 0.42]",
-                "parameters": {"bboxes": [0.2, 0.21, 0.34, 0.42]},
-            }
+                "scenario": "Calculate the width and height of the bounding box [0.2, 0.21, 0.34, 0.42]",
+                "parameters": {
+                    "bboxes": [[0.2, 0.21, 0.34, 0.42]],
+                    "image_size": (500, 1200),
+                },
+            },
+            {
+                "scenario": "Calculate the area of the bounding box [0.2, 0.21, 0.34, 0.42]",
+                "parameters": {
+                    "bboxes": [[0.2, 0.21, 0.34, 0.42]],
+                    "image_size": (640, 480),
+                },
+            },
         ],
     }
 
-    def __call__(self, bboxes: List[Dict]) -> List[Dict]:
+    def __call__(
+        self, bboxes: List[List[int]], image_size: Tuple[int, int]
+    ) -> List[Dict]:
         areas = []
-        for elt in bboxes:
-            height, width = elt["size"]
-            for label, bbox in zip(elt["labels"], elt["bboxes"]):
-                x1, y1, x2, y2 = bbox
-                areas.append(
-                    {
-                        "area": round((x2 - x1) * (y2 - y1) * width * height, 2),
-                        "label": label,
-                    }
-                )
+        height, width = image_size
+        for bbox in bboxes:
+            x1, y1, x2, y2 = bbox
+            areas.append(
+                {
+                    "width": round((x2 - x1) * width, 2),
+                    "height": round((y2 - y1) * height, 2),
+                    "area": round((x2 - x1) * (y2 - y1) * width * height, 2),
+                }
+            )
+
         return areas
 
 
@@ -878,7 +903,7 @@ class SegIoU(Tool):
         ],
         "examples": [
             {
-                "scenario": "If you want to calculate the intersection over union of the segmentation masks for mask_file1.jpg and mask_file2.jpg",
+                "scenario": "Calculate the intersection over union of the segmentation masks for mask_file1.jpg and mask_file2.jpg",
                 "parameters": {"mask1": "mask_file1.png", "mask2": "mask_file2.png"},
             }
         ],
@@ -947,6 +972,46 @@ class BboxContains(Tool):
         }
 
 
+class ObjectDistance(Tool):
+    name = "object_distance_"
+    description = "'object_distance_' calculates the distance between two objects in an image. It returns the minimum distance between the two objects."
+    usage = {
+        "required_parameters": [
+            {"name": "object1", "type": "Dict[str, Any]"},
+            {"name": "object2", "type": "Dict[str, Any]"},
+        ],
+        "examples": [
+            {
+                "scenario": "Calculate the distance between these two objects {bboxes: [0.2, 0.21, 0.34, 0.42], masks: 'mask_file1.png'}, {bboxes: [0.3, 0.31, 0.44, 0.52], masks: 'mask_file2.png'}",
+                "parameters": {
+                    "object1": {
+                        "bboxes": [0.2, 0.21, 0.34, 0.42],
+                        "scores": 0.54,
+                        "masks": "mask_file1.png",
+                    },
+                    "object2": {
+                        "bboxes": [0.3, 0.31, 0.44, 0.52],
+                        "scores": 0.66,
+                        "masks": "mask_file2.png",
+                    },
+                },
+            }
+        ],
+    }
+
+    def __call__(self, object1: Dict[str, Any], object2: Dict[str, Any]) -> float:
+        if "masks" in object1 and "masks" in object2:
+            mask1 = object1["masks"]
+            mask2 = object2["masks"]
+            return MaskDistance()(mask1, mask2)
+        elif "bboxes" in object1 and "bboxes" in object2:
+            bbox1 = object1["bboxes"]
+            bbox2 = object2["bboxes"]
+            return BoxDistance()(bbox1, bbox2)
+        else:
+            raise ValueError("Either of the objects should have masks or bboxes")
+
+
 class BoxDistance(Tool):
     name = "box_distance_"
     description = "'box_distance_' calculates distance between two bounding boxes. It returns the minumum distance between the given bounding boxes"
@@ -957,7 +1022,7 @@ class BoxDistance(Tool):
         ],
         "examples": [
             {
-                "scenario": "Calculate the distance between the bounding boxes [0.2, 0.21, 0.34, 0.42] and [0.3, 0.31, 0.44, 0.52]",
+                "scenario": "Calculate the distance between these two bounding boxes [0.2, 0.21, 0.34, 0.42] and [0.3, 0.31, 0.44, 0.52]",
                 "parameters": {
                     "bbox1": [0.2, 0.21, 0.34, 0.42],
                     "bbox2": [0.3, 0.31, 0.44, 0.52],
@@ -976,26 +1041,57 @@ class BoxDistance(Tool):
         return cast(float, round(np.sqrt(horizontal_dist**2 + vertical_dist**2), 2))
 
 
+class MaskDistance(Tool):
+    name = "mask_distance_"
+    description = "'mask_distance_' calculates distance between two masks. It is helpful in checking proximity of two objects. It returns the minumum distance between the given masks"
+    usage = {
+        "required_parameters": [
+            {"name": "mask1", "type": "str"},
+            {"name": "mask2", "type": "str"},
+        ],
+        "examples": [
+            {
+                "scenario": "Calculate the distance between the segmentation masks for mask_file1.jpg and mask_file2.jpg",
+                "parameters": {"mask1": "mask_file1.png", "mask2": "mask_file2.png"},
+            }
+        ],
+    }
+
+    def __call__(self, mask1: Union[str, Path], mask2: Union[str, Path]) -> float:
+        pil_mask1 = Image.open(str(mask1))
+        pil_mask2 = Image.open(str(mask2))
+        np_mask1 = np.clip(np.array(pil_mask1), 0, 1)
+        np_mask2 = np.clip(np.array(pil_mask2), 0, 1)
+
+        mask1_points = np.transpose(np.nonzero(np_mask1))
+        mask2_points = np.transpose(np.nonzero(np_mask2))
+        dist_matrix = distance.cdist(mask1_points, mask2_points, "euclidean")
+        return cast(float, np.round(np.min(dist_matrix), 2))
+
+
 class ExtractFrames(Tool):
     r"""Extract frames from a video."""
 
     name = "extract_frames_"
-    description = "'extract_frames_' extracts frames from a video, returns a list of tuples (frame, timestamp), where timestamp is the relative time in seconds where the frame was captured. The frame is a local image file path."
+    description = "'extract_frames_' extracts frames from a video every 2 seconds, returns a list of tuples (frame, timestamp), where timestamp is the relative time in seconds where the frame was captured. The frame is a local image file path."
     usage = {
         "required_parameters": [{"name": "video_uri", "type": "str"}],
+        "optional_parameters": [{"name": "frames_every", "type": "float"}],
         "examples": [
             {
                 "scenario": "Can you extract the frames from this video? Video: www.foobar.com/video?name=test.mp4",
                 "parameters": {"video_uri": "www.foobar.com/video?name=test.mp4"},
             },
             {
-                "scenario": "Can you extract the images from this video file? Video path: tests/data/test.mp4",
-                "parameters": {"video_uri": "tests/data/test.mp4"},
+                "scenario": "Can you extract the images from this video file at every 2 seconds ? Video path: tests/data/test.mp4",
+                "parameters": {"video_uri": "tests/data/test.mp4", "frames_every": 2},
             },
         ],
     }
 
-    def __call__(self, video_uri: str) -> List[Tuple[str, float]]:
+    def __call__(
+        self, video_uri: str, frames_every: float = 2
+    ) -> List[Tuple[str, float]]:
         """Extract frames from a video.
 
 
@@ -1005,7 +1101,7 @@ class ExtractFrames(Tool):
         Returns:
             a list of tuples containing the extracted frame and the timestamp in seconds. E.g. [(path_to_frame1, 0.0), (path_to_frame2, 0.5), ...]. The timestamp is the time in seconds from the start of the video. E.g. 12.125 means 12.125 seconds from the start of the video. The frames are sorted by the timestamp in ascending order.
         """
-        frames = extract_frames_from_video(video_uri)
+        frames = extract_frames_from_video(video_uri, fps=round(1 / frames_every, 2))
         result = []
         _LOGGER.info(
             f"Extracted {len(frames)} frames from video {video_uri}. Temporarily saving them as images to disk for downstream tasks."
@@ -1108,12 +1204,11 @@ TOOLS = {
             AgentDINOv,
             ExtractFrames,
             Crop,
-            BboxArea,
+            BboxStats,
             SegArea,
-            BboxIoU,
-            SegIoU,
+            ObjectDistance,
             BboxContains,
-            BoxDistance,
+            SegIoU,
             OCR,
             Calculator,
         ]
@@ -1145,20 +1240,3 @@ def register_tool(tool: Type[Tool]) -> Type[Tool]:
         "class": tool,
     }
     return tool
-
-
-def _send_inference_request(
-    payload: Dict[str, Any], endpoint_name: str
-) -> Dict[str, Any]:
-    res = requests.post(
-        f"{_LND_API_URL}/model/{endpoint_name}",
-        headers={
-            "Content-Type": "application/json",
-            "apikey": _LND_API_KEY,
-        },
-        json=payload,
-    )
-    if res.status_code != 200:
-        _LOGGER.error(f"Request failed: {res.text}")
-        raise ValueError(f"Request failed: {res.text}")
-    return res.json()["data"]  # type: ignore