add mask to necessary keys

landing-ai · Apr 25, 2024 · 3daa4f0 · 3daa4f0
1 parent 1305e05
commit 3daa4f0
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py
@@ -346,7 +346,9 @@ def _handle_viz_tools(
  # 2. return a dictionary but not have the necessary keys
 
  if not isinstance(call_result, dict) or (
- "bboxes" not in call_result and "heat_map" not in call_result
+ "bboxes" not in call_result
+ and "mask" not in call_result
+ and "heat_map" not in call_result
  ):
  return image_to_data