From 577fe664cb5b931bef7031f3e91207d86fa3844d Mon Sep 17 00:00:00 2001
From: Dillon Laird <dillonalaird@gmail.com>
Date: Wed, 24 Apr 2024 20:55:58 -0700
Subject: [PATCH] updated return for visual prompt counting

---
 vision_agent/tools/tools.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/vision_agent/tools/tools.py b/vision_agent/tools/tools.py
index fd315bed..fa06a823 100644
--- a/vision_agent/tools/tools.py
+++ b/vision_agent/tools/tools.py
@@ -588,7 +588,9 @@ def __call__(self, image: Union[str, ImageType], prompt: str) -> Dict:
             "prompt": prompt,
             "tool": "few_shot_counting",
         }
-        return _send_inference_request(data, "tools")
+        resp_data = _send_inference_request(data, "tools")
+        resp_data["heat_map"] = np.array(b64_to_pil(resp_data["heat_map"][0]))
+        return resp_data
 
 
 class VisualQuestionAnswering(Tool):