fix heat map append

landing-ai · Apr 26, 2024 · 3873b4f · 3873b4f
1 parent 3daa4f0
commit 3873b4f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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"]