From ebb5f1cbc5f7fb3afba506276da17fece9e4f964 Mon Sep 17 00:00:00 2001
From: Jing Hua <tohjinghua123@gmail.com>
Date: Thu, 30 Nov 2023 17:31:46 +0800
Subject: [PATCH] add mouse cursor to screenshot

---
 desktop_env/assets/cursor.png      | Bin 0 -> 4796 bytes
 desktop_env/controllers/mouse.py   |  17 ++++++++++++++++-
 desktop_env/controllers/python.py  |   1 +
 desktop_env/envs/desktop_env.py    |  22 +++++++++++++++++-----
 desktop_env/windows_server/main.py |   4 +++-
 5 files changed, 37 insertions(+), 7 deletions(-)
 create mode 100644 desktop_env/assets/cursor.png

diff --git a/desktop_env/assets/cursor.png b/desktop_env/assets/cursor.png
new file mode 100644
index 0000000000000000000000000000000000000000..8d9182a99ddafecce0ecac943a7fb5901ca6545b
GIT binary patch
literal 4796
zcmeHKc~lff8lT}1P&^O}c*IObBr5crJxou_h{%jEkwI<&iPF>41Fds8-N1m0Njwlt
zz%xh?uVB<oi~{>S@BlBoV?tIX&vPG0*cTObbr%r~Tix8dd2gTZZT@5C74=os@B3BN
zufF=2Du{{<8##Q!a1MtvQWLI<0rdn>{9T8EzXoH;1yKD+Bx(#A%?yq!XaFSFiNggJ
zgX8wt{=jzNI&vIAod?QTU^5sN$8~1w&p~@)kG%!h(|RzBkC^Mg842n#P$a<i0`)hb
zkZZf=ssT3hab>@92M&wqt;00hD4s+h773&<@E1v>s920jMLe+_mB>&T0_NxL>R%DV
zwED=)<zRiQ$<WX!O=u|3YO@$fBf;U=%W^j=!_Up0RJS;`?$Ekb&T5q;?t>F?9v73>
z@Z($$&-r$g-_{i&oAdV6)kRw&<W*&2hIQGkn`h$3I0es6`KZk6qdn)lrmlHMJb6-E
z)4fv{@6DXOe`|>$L&en|D!FWrh)7GQJGkH_=C1UaqcRkWKhU_@iw-S+^(x2vm!$hG
zzfF3x<d*J}`MyPxa|gy)#{6JRaKGMhJ|tm7d?VGU*(?a}yp&PbEWe(;)_KQxIQDXO
zBd@Z=f`46@bmm6>zJkzi_Quw1$V_TJS(oG=pY&o!=bBFs)Kp!|C@s&&J?3rT%?TN&
z-Vzm3lCa%tgy#18>{nrvYE^<YDQb}|V3+&IswS7NRS)hZSN}PGn^!(}yYR>x@hJL<
zLt&}u<M-}lHaR^Dc3hBib@4&RzwF%8kevK?SD&AU2@WnSt#V0e<(!+JXS`DMaNqQY
zw$mNsb7P13oOI4<xO=~}1IgvY+=y=+h9FLU|JZV>^?7USzwZC0{OjYVF4s5OuHrt6
zl^_P~WLzShsEt5zi%Ed#EjmJAH(8kj;{*lUtr)(Ppm{nXjWjFyPfyhHd8A&+Um(+p
zwAN4}oea;h5wSUuad^&BT%qR&&l(<NM*)F}pfR4^WHeK#UCC#8QD8IOLOzd$&`Xv4
zL=Z=z78}8X1+YK_sqJK@gg<LIFUY1hpfResJrtm&<fqfL6%`7zva$qOQh~*mCKM|a
z3ZY0Mlt>_efT(OUjoBeHHHD$*=1>t7ZX>NUX)*H{PE2RXpp|?+80Yn_&t%nV`{>P7
zj|xB!p&hdd#R8GgWD*YaplEd_An94q?|M*i*;YarLr|6s8&0S*2{S!qAcY?9^S5T$
zjO=#wxR5XsCIF?ttl~jaGB>WzgOMPOG+9|MAod_6O&a>e8niZMgx$`-g#h<H-a+Wz
zy|Z9|(rQta1<zoXr%@^S%>1a{f|GibHKkI75Q_-}k{j?q2*wZ*q>$k-gun`^PNC3A
zBvR=BDvg<<F*8muRDfJS0-iuT5XbaV7z)I77zD#QDWp(<b2Z3uolGiK08+&OiupDY
zL?vb%xGIK952&PiodP!~<d7bN0ac(}2O%QZ0O2BBBu5Yg(_`IKdK{f=v6(P%oTLd$
zBZO9S8aoFgI2sbAQSv1MQGZXA5u**jLCKFG%^CLofjH7c#L^g}rdTc#OQnHQ86tya
z3Ruz)N+4_$h(v~y&7JHBb1W#>3=kG$;uH|D{a`O>sExpAi!IJ#F)I0eS+{Rk3$BwM
zqcIgm6M(cY56AW9VSf>Vin^E2&}l7t(vbbX(9GrG1pz!b@^F#@^JlYN-8U+hSk~R@
zZW>AUD)D&ir9d&fdkYknN$A;i0<P{MJRLKq5g>c?gllg<`8%P2$wUYN6ADNt$7K*~
zz%dAf7Y->Tff9opMqs@}Hkh5V80aj_MuemR9f4LLJlX$HylLz~O&^TTN++080D?hc
zF(eiJr(nXK4GWo^(Yxay;eTlo!~zB+8Q|C52a*@ag~GmM*rOSfcK*Om&sqF|GXUzL
zN#07|A-RU+dMgFq3Op2DLvp>90&fK#imv~gT*LdHQ-m4(3d#bHrM!gdTJX>srkfY0
zVxL@`_Rtd_1Ix`Co<wmtBRrYH<&>8BfKC@$qgA`KjF{{;9*tRa2fWE0k~FH2xWsec
zH)!*VqP<+6J-GT1|4A>Nwdt!b)8A3SQ{NxKFWz+EPk8|WUa^5=o2l|?7$%;0ex-Aj
zk2+wlW>a*2OikV#_qoy+x66*2CrF#Nw_W*eUt4w8^49j?mfY6%oiAojY{kF2R_@af
z8}Y_lr9OG#+FeycQ=4)EHMa0{`?0QdUsva+B9Rw6T@sr4D;<iQZWR_*=YLo}{n?_&
z*ECLde)N2_);9J`a;N(?PoLt?ToW1>{wyn>;q4J{wx#0K%6i}}K+o1V%c}EtafSGN
z|0R^i8Lyy%+$SYvpSk*}vZpsX?*+`2$2>b8=n@3{{TB{Dd+zaq@^EXe-7u}IDn&K5
zIoPx2cm?Vmt!ZEHccoV09liKU?Y!#Jn?LL>8M9!f8OLrMsbAw;dY1QY<oegJThrK<
zn{OVs$3&_gRURIF(tpCXtYx|R?b|LL=02$ROVPrn%+9uZi5?!e+~Zf4Zu+OA+lna_
z_d?^JI*f8zG^6=lpK*`oIV>M{IXTmt?mT*aX<_zlkHgMKcHgaOO}RF+GU!>xr^<-P
zmfNQ%k3O<~nn&UJ=3`x}oeRFdgecVkk&}N-O7-u`+dnnin|fVw$NSM%Ph(i>(JSqx
zCC^<>Jg@g_+a-uz)mE8xKeb>-`3laRUn-QvsU40dkTLbzyAvOJFUDj3cENY!iuFzh
z5ABgW7+LkiEptZ2<}ZAcJtKGeUOX<~|LVT`yG#2&e0{vxZN>D8Q(t_(WIl4L_@%qc
z=9l{uw|`r>w^j~6kGJK#7k_LQ+OcrdtkcC&8`l-(h2Pp%zi3?N=qB%<RyP(rPOolY
O!cwD-R2`g?y7J$F-?mu*

literal 0
HcmV?d00001

diff --git a/desktop_env/controllers/mouse.py b/desktop_env/controllers/mouse.py
index 45961be..cd63175 100644
--- a/desktop_env/controllers/mouse.py
+++ b/desktop_env/controllers/mouse.py
@@ -2,6 +2,7 @@ from enum import Enum
 
 from abc import ABC, abstractmethod
 from fabric import Connection
+import re
 
 from .xdotool import XDoToolController
 from .python import PythonController
@@ -13,6 +14,10 @@ class MouseClick(Enum):
     WHEEL_DOWN = 5
 
 class AbstractMouseController(ABC):
+    @abstractmethod
+    def get_mouse(self):
+        raise NotImplementedError
+
     @abstractmethod
     def mouse_move(self, x: int, y: int):
         raise NotImplementedError
@@ -66,6 +71,10 @@ class XDoToolMouseController(AbstractMouseController, XDoToolController):
     def __init__(self, ssh_connection: Connection):
         super().__init__(ssh_connection=ssh_connection)
 
+    @abstractmethod
+    def get_mouse(self):
+        raise NotImplementedError
+
     def mouse_move(self, x: int, y: int):
         self._execute_xdotool_command(f"mousemove {x} {y}")
 
@@ -105,7 +114,13 @@ class XDoToolMouseController(AbstractMouseController, XDoToolController):
 class PythonMouseController(AbstractMouseController, PythonController):
     def __init__(self, http_server: str):
         super().__init__(http_server=http_server)
-        self.command = "python -c \"import mouse; {command}\""
+        self.command = "py -c \"import mouse; {command}\""
+
+    def get_mouse(self):
+        response = self._execute_python_command(self.command.format(command=f"print(mouse.get_position())"))
+        numbers = re.findall(r'-?\d+', response["output"])
+        x, y = map(int, numbers)
+        return x, y
 
     def mouse_move(self, x: int, y: int):
         self._execute_python_command(self.command.format(command=f"mouse.move({x}, {y})"))
diff --git a/desktop_env/controllers/python.py b/desktop_env/controllers/python.py
index 735bd44..1e6a751 100644
--- a/desktop_env/controllers/python.py
+++ b/desktop_env/controllers/python.py
@@ -20,6 +20,7 @@ class PythonController:
                 print("Command executed successfully:", response.text)
             else:
                 print("Failed to execute command. Status code:", response.status_code)
+            return response.json()
         except requests.exceptions.RequestException as e:
             print("An error occurred while trying to execute the command:", e)
 
diff --git a/desktop_env/envs/desktop_env.py b/desktop_env/envs/desktop_env.py
index 8755fdc..d032f87 100644
--- a/desktop_env/envs/desktop_env.py
+++ b/desktop_env/envs/desktop_env.py
@@ -58,6 +58,9 @@ class DesktopEnv(gym.Env):
         print("Initializing...")
         self._start_emulator()
 
+        # set up controllers
+        self.mouse_controller, self.keyboard_controller = self._create_controllers(vm_os)
+
         # Get the screen size
         self.screen_width, self.screen_height = self._get_screensize()
 
@@ -77,8 +80,6 @@ class DesktopEnv(gym.Env):
         # Additional setup
         self.metadata = {'render.modes': ['rgb_array']}
 
-        # set up controllers
-        self.mouse_controller, self.keyboard_controller = self._create_controllers(vm_os)
 
     def _get_screensize(self):
         screenshot_path = self._get_obs()
@@ -114,7 +115,9 @@ class DesktopEnv(gym.Env):
                 print(f"Error executing command: {e.output.decode().strip()}")
 
     def _execute_command(self, command: List[str]) -> None:
-        subprocess.run(command, shell=True, stderr=subprocess.STDOUT, timeout=60)
+        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=60, text=True)
+        if result.returncode != 0:
+            raise Exception("\033[91m" + result.stdout + result.stderr + "\033[0m")
 
     def _save_state(self):
         self._execute_command(["vmrun", "-T", "ws" "snapshot", self.path_to_vm, self.snapshot_path])
@@ -126,8 +129,7 @@ class DesktopEnv(gym.Env):
 
         if self.password:
             self._execute_command(
-                ["vmrun", "-T", "ws", "-gu", self.username, "-gp", self.password, "captureScreen", self.path_to_vm,
-                 image_path])
+                ["vmrun", "-T", "ws", "-gu", self.username, "-gp", self.password, "captureScreen", self.path_to_vm, image_path])
         else:
             self._execute_command(
                 ["vmrun", "-T", "ws", "-gu", self.username, "captureScreen", self.path_to_vm, image_path])
@@ -136,7 +138,17 @@ class DesktopEnv(gym.Env):
 
     def _get_obs(self):
         screenshot_image_path = self._get_screenshot()
+        self._add_cursor(screenshot_image_path)
         return screenshot_image_path
+    
+    def _add_cursor(self, img_path: str):
+        x, y = self.mouse_controller.get_mouse()
+        cursor_image = Image.open("./desktop_env/assets/cursor.png")
+        cursor_image = cursor_image.resize((int(cursor_image.width / 2), int(cursor_image.height / 2)))
+
+        screenshot = Image.open(img_path)
+        screenshot.paste(cursor_image, (x, y), cursor_image)
+        screenshot.save(img_path)
 
     def reset(self):
         print("Resetting environment...")
diff --git a/desktop_env/windows_server/main.py b/desktop_env/windows_server/main.py
index 467b40e..098dca5 100644
--- a/desktop_env/windows_server/main.py
+++ b/desktop_env/windows_server/main.py
@@ -11,9 +11,11 @@ def execute_command():
 
     # Execute the command without any safety checks.
     try:
-        subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
         return jsonify({
             'status': 'success',
+            'output': result.stdout,
+            'error': result.stderr
         })
     except Exception as e:
         return jsonify({