From e166106b6ad13251fcb61d328dae67288649e7df Mon Sep 17 00:00:00 2001
From: David Chang <zdy004007@126.com>
Date: Fri, 15 Mar 2024 22:46:14 +0800
Subject: [PATCH] ver Mar15th

added an option to keep buttons without text information but with an
image for SoM setting
---
 desktop_env/server/main.py                         |  9 +++++++++
 .../accessibility_tree_wrap/heuristic_retrieve.py  | 14 ++++++++------
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/desktop_env/server/main.py b/desktop_env/server/main.py
index 57649ac..efa62c7 100644
--- a/desktop_env/server/main.py
+++ b/desktop_env/server/main.py
@@ -284,6 +284,15 @@ def _create_atspi_node(node: Accessible, depth: int = 0, flag: Optional[str] = N
         text = text.replace("\ufffc", "").replace("\ufffd", "")
     #  }}} Text # 
 
+    #  Image {{{ # 
+    try:
+        node.queryImage()
+    except NotImplementedError:
+        pass
+    else:
+        attribute_dict["image"] = "true"
+    #  }}} Image # 
+
     #  Selection {{{ # 
     try:
         node.querySelection()
diff --git a/mm_agents/accessibility_tree_wrap/heuristic_retrieve.py b/mm_agents/accessibility_tree_wrap/heuristic_retrieve.py
index 34a1d76..191eaa7 100644
--- a/mm_agents/accessibility_tree_wrap/heuristic_retrieve.py
+++ b/mm_agents/accessibility_tree_wrap/heuristic_retrieve.py
@@ -26,7 +26,7 @@ def find_leaf_nodes(xlm_file_str):
 
 state_ns = "uri:deskat:state.at-spi.gnome.org"
 component_ns = "uri:deskat:component.at-spi.gnome.org"
-def judge_node(node: ET, platform="ubuntu") -> bool:
+def judge_node(node: ET, platform="ubuntu", check_image=False) -> bool:
     keeps: bool = node.tag.startswith("document")\
                or node.tag.endswith("item")\
                or node.tag.endswith("button")\
@@ -60,7 +60,9 @@ def judge_node(node: ET, platform="ubuntu") -> bool:
                      or node.get("{{{:}}}expandable".format(state_ns), "false")=="true"\
                      or node.get("{{{:}}}checkable".format(state_ns), "false")=="true"
                       )\
-                  and (node.get("name", "") != "" or node.text is not None and len(node.text)>0)
+                  and ( node.get("name", "") != "" or node.text is not None and len(node.text)>0\
+                     or check_image and node.get("image", "false")=="true"
+                      )
 
     coordinates: Tuple[int, int] = eval(node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)"))
     sizes: Tuple[int, int] = eval(node.get("{{{:}}}size".format(component_ns), "(-1, -1)"))
@@ -155,12 +157,12 @@ def print_nodes_with_indent(nodes, indent=0):
 
 if __name__ == '__main__':
     import json
-    with open('4.json', 'r', encoding='utf-8') as f:
-        xml_file_str = json.load(f)["AT"]
+    with open('selection_sorted(imaged).xml', 'r', encoding='utf-8') as f:
+        xml_file_str = f.read()
     filtered_nodes = filter_nodes(ET.fromstring(xml_file_str))
     print(len(filtered_nodes))
-    masks = draw_bounding_boxes( filtered_nodes, '4.png'
-                               , '4.a.png'
+    masks = draw_bounding_boxes( filtered_nodes, 'selection_sorted(imaged).png'
+                               , 'selection_sorted(imaged).ai.png'
                                )
 
     # print(masks)