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