Improve: fix bugs; add back the cursor in screenshot; add pause in env.step

desktop_env/server/main.py

@@ -1,12 +1,14 @@
 import os
 import platform
 import subprocess
+import requests
 
 import Xlib.display
 import pyautogui
-from PIL import ImageGrab
+from PIL import ImageGrab, Image
 from flask import Flask, request, jsonify, send_file
 
+
 app = Flask(__name__)
 
 pyautogui.PAUSE = 0
@@ -43,7 +45,19 @@ def capture_screen_with_cursor():
     os.makedirs(os.path.dirname(file_path), exist_ok=True)
 
     if user_platform == "Windows":
+        def _download_image(url, path):
+            response = requests.get(url)
+            with open(path, 'wb') as file:
+                file.write(response.content)
+
+        cursor_path = os.path.join("screenshots", "cursor.png")
+        if not os.path.exists(cursor_path):
+            cursor_url = "https://vip.helloimg.com/images/2023/12/02/oQPzmt.png"
+            _download_image(cursor_url, cursor_path)
         screenshot = pyautogui.screenshot()
+        cursor_x, cursor_y = pyautogui.position()
+        cursor = Image.open(cursor_path)
+        screenshot.paste(cursor, (cursor_x, cursor_y), cursor)
         screenshot.save(file_path)
     elif user_platform == "Linux":
         # Use xlib to prevent scrot dependency for Linux
@@ -60,5 +74,15 @@ def capture_screen_with_cursor():
     return send_file(file_path, mimetype='image/png')
 
 
+@app.route('/platform', methods=['GET'])
+def get_platform():
+    return platform.system()
+
+
+@app.route('/cursor_position', methods=['GET'])
+def get_cursor_position():
+    return pyautogui.position().x, pyautogui.position().y
+
+
 if __name__ == '__main__':
     app.run(debug=True, host="0.0.0.0")