From 3873b4f53b125d042046cfb447939da49b1406e7 Mon Sep 17 00:00:00 2001 From: Dillon Laird Date: Thu, 25 Apr 2024 20:45:43 -0700 Subject: [PATCH] fix heat map append --- vision_agent/agent/vision_agent.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py index ef0f305c..7d937e3b 100644 --- a/vision_agent/agent/vision_agent.py +++ b/vision_agent/agent/vision_agent.py @@ -368,7 +368,8 @@ def _handle_viz_tools( image_to_data[image]["scores"].extend(call_result.get("scores", [])) image_to_data[image]["masks"].extend(call_result.get("masks", [])) # only single heatmap is returned - image_to_data[image]["heat_map"].append(call_result.get("heat_map", [])) + if "heat_map" in call_result: + image_to_data[image]["heat_map"].append(call_result["heat_map"]) if "mask_shape" in call_result: image_to_data[image]["mask_shape"] = call_result["mask_shape"]