From be6d1ceaa9608108d79ea277636724d94d343a67 Mon Sep 17 00:00:00 2001
From: Dillon Laird <dillonalaird@gmail.com>
Date: Wed, 17 Apr 2024 14:34:09 -0700
Subject: [PATCH] added box distance

---
 tests/tools/test_tools.py      | 25 ++++++++++++++++++++++++-
 vision_agent/tools/__init__.py |  1 +
 vision_agent/tools/tools.py    | 15 ++++++++++++---
 3 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/tests/tools/test_tools.py b/tests/tools/test_tools.py
index c5d7c6cc..12c21347 100644
--- a/tests/tools/test_tools.py
+++ b/tests/tools/test_tools.py
@@ -4,7 +4,7 @@
 import numpy as np
 from PIL import Image
 
-from vision_agent.tools.tools import BboxIoU, SegArea, SegIoU
+from vision_agent.tools.tools import BboxIoU, BoxDistance, SegArea, SegIoU
 
 
 def test_bbox_iou():
@@ -42,3 +42,26 @@ def test_seg_area_2():
         mask_path = os.path.join(tmpdir, "mask.png")
         Image.fromarray(mask).save(mask_path)
         assert SegArea()(mask_path) == 4.0
+
+
+def test_box_distance():
+    box_dist = BoxDistance()
+    # horizontal dist
+    box1 = [0, 0, 2, 2]
+    box2 = [4, 1, 6, 3]
+    assert box_dist(box1, box2) == 2.0
+
+    # vertical dist
+    box1 = [0, 0, 2, 2]
+    box2 = [1, 4, 3, 6]
+    assert box_dist(box1, box2) == 2.0
+
+    # vertical and horizontal
+    box1 = [0, 0, 2, 2]
+    box2 = [3, 3, 5, 5]
+    assert box_dist(box1, box2) == 1.41
+
+    # overlap
+    box1 = [0, 0, 2, 2]
+    box2 = [1, 1, 3, 3]
+    assert box_dist(box1, box2) == 0.0
diff --git a/vision_agent/tools/__init__.py b/vision_agent/tools/__init__.py
index 63a410aa..63931c9f 100644
--- a/vision_agent/tools/__init__.py
+++ b/vision_agent/tools/__init__.py
@@ -4,6 +4,7 @@
     TOOLS,
     BboxArea,
     BboxIoU,
+    BoxDistance,
     Crop,
     ExtractFrames,
     GroundingDINO,
diff --git a/vision_agent/tools/tools.py b/vision_agent/tools/tools.py
index b7f4fab4..3a5c8a4f 100644
--- a/vision_agent/tools/tools.py
+++ b/vision_agent/tools/tools.py
@@ -547,7 +547,9 @@ def __call__(self, mask1: Union[str, Path], mask2: Union[str, Path]) -> float:
 
 class BoxDistance(Tool):
     name = "box_distance_"
-    description = "'box_distance_' returns the distance between two bounding boxes."
+    description = (
+        "'box_distance_' returns the minimum distance between two bounding boxes."
+    )
     usage = {
         "required_parameters": [
             {"name": "bbox1", "type": "List[int]"},
@@ -564,8 +566,14 @@ class BoxDistance(Tool):
         ],
     }
 
-    def __call__(self, box1: List[int], box2: List[int]) -> float:
-        raise NotImplementedError("Not implemented yet.")
+    def __call__(self, bbox1: List[int], bbox2: List[int]) -> float:
+        x11, y11, x12, y12 = bbox1
+        x21, y21, x22, y22 = bbox2
+
+        horizontal_dist = np.max([0, x21 - x12, x11 - x22])
+        vertical_dist = np.max([0, y21 - y12, y11 - y22])
+
+        return cast(float, round(np.sqrt(horizontal_dist**2 + vertical_dist**2), 2))
 
 
 class ExtractFrames(Tool):
@@ -650,6 +658,7 @@ def __call__(self, equation: str) -> float:
             SegArea,
             BboxIoU,
             SegIoU,
+            BoxDistance,
             Calculator,
         ]
     )