Support downsampling; Fix bugs in windows a11y tree; Add a11y_tree trim

2024-03-25 18:02:48 +08:00
parent 635b6717b3
commit 172123ab2c
4 changed files with 104 additions and 43 deletions
--- a/desktop_env/envs/desktop_env.py
+++ b/desktop_env/envs/desktop_env.py
@@ -146,7 +146,13 @@ class DesktopEnv(gym.Env):
        image_path: str = os.path.join(self.tmp_dir, "screenshots", "{:d}.png".format(self._step_no))

        # Get the screenshot and save to the image_path
-        screenshot = self.controller.get_screenshot()
+        max_retries = 20
+        for _ in range(max_retries):
+            screenshot = self.controller.get_screenshot()
+            if screenshot is not None:
+                break
+            time.sleep(1)
+
        with open(image_path, "wb") as f:
            f.write(screenshot)