updated return for visual prompt counting

landing-ai · Apr 25, 2024 · 577fe66 · 577fe66
1 parent 8e3c525
commit 577fe66
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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):