Fix the cursor logo placement for ubuntu

desktop_env/evaluators/metrics/table.py

@@ -1,14 +1,13 @@
 import pandas as pd
-#import pylightxl
 import openpyxl
-#from openpyxl import Workbook
 from openpyxl.worksheet.worksheet import Worksheet
 
-from utils import load_charts, load_sparklines
+from .utils import load_charts, load_sparklines
 
 from typing import Dict, List
 from typing import Any
 
+
 def compare_table(actual, expected):
     df1 = pd.read_excel(expected)
     df2 = pd.read_excel(actual)
@@ -16,6 +15,7 @@ def compare_table(actual, expected):
     # Compare the DataFrames
     return 1 if df1.equals(df2) else 0
 
+
 def compare_with_sparklines(actual: str, expected: str) -> float:
     df1 = pd.read_excel(actual)
     df2 = pd.read_excel(expected)
@@ -29,6 +29,7 @@ def compare_with_sparklines(actual: str, expected: str) -> float:
 
     return float(normal_content_metric and sparkline_metric)
 
+
 def compare_with_charts(actual: str, expected: str, **options) -> float:
     """
     Args:
@@ -50,6 +51,7 @@ def compare_with_charts(actual: str, expected: str, **options) -> float:
 
     return float(normal_content_metric and chart_metric)
 
+
 def check_sheet_list(result: str, rules: List[Dict[str, Any]]) -> float:
     # workbook: Workbook = openpyxl.load_workbook(filename=result)
     workbook = pd.ExcelFile(result)
@@ -88,10 +90,12 @@ def check_sheet_list(result: str, rules: List[Dict[str, Any]]) -> float:
 
     return float(passes)
 
+
 def check_xlsx_freeze(result: str, rules: Dict[str, str]) -> float:
     worksheet: Worksheet = openpyxl.load_workbook(filename=result).active
     return float(worksheet.freeze_panes == rules["position"])
 
+
 if __name__ == '__main__':
     # path1 = ""
     # path2 = ""

desktop_env/server/main.py

@@ -3,7 +3,7 @@ from pathlib import Path
 import platform
 import subprocess
 import requests
-
+from .pyxcursor import Xcursor
 # import Xlib.display
 import pyautogui
 # from PIL import ImageGrab, Image
@@ -48,7 +48,7 @@ def capture_screen_with_cursor():
     os.makedirs(os.path.dirname(file_path), exist_ok=True)
 
     # fixme: This is a temporary fix for the cursor not being captured on Windows and Linux
-    if user_platform == "Windows" or user_platform == "Linux":
+    if user_platform == "Windows":
         def _download_image(url, path):
             response = requests.get(url)
             with open(path, 'wb') as file:
@@ -65,12 +65,14 @@ def capture_screen_with_cursor():
         cursor = cursor.resize((int(cursor.width / 1.5), int(cursor.height / 1.5)))
         screenshot.paste(cursor, (cursor_x, cursor_y), cursor)
         screenshot.save(file_path)
-    # elif user_platform == "Linux":
+    elif user_platform == "Linux":
-    #     # Use xlib to prevent scrot dependency for Linux
+        cursor_obj = Xcursor()
-    #     screen = Xlib.display.Display().screen()
+        imgarray = cursor_obj.getCursorImageArrayFast()
-    #     size = screen.width_in_pixels, screen.height_in_pixels
+        cursor_img = Image.fromarray(imgarray)
-    #     screenshot = ImageGrab.grab(bbox=(0, 0, size[0], size[1]))
+        screenshot = pyautogui.screenshot()
-    #     screenshot.save(file_path)
+        cursor_x, cursor_y = pyautogui.position()
+        screenshot.paste(cursor_img, (cursor_x, cursor_y), cursor_img)
+        screenshot.save(file_path)
     elif user_platform == "Darwin":  # (Mac OS)
         # Use the screencapture utility to capture the screen with the cursor
         subprocess.run(["screencapture", "-C", file_path])

desktop_env/server/pyxcursor.py

@@ -0,0 +1,146 @@
+import os
+import ctypes
+import ctypes.util
+import numpy as np
+
+# A helper function to convert data from Xlib to byte array.
+import struct, array
+
+# Define ctypes version of XFixesCursorImage structure.
+PIXEL_DATA_PTR = ctypes.POINTER(ctypes.c_ulong)
+Atom = ctypes.c_ulong
+
+
+class XFixesCursorImage(ctypes.Structure):
+    """
+    See /usr/include/X11/extensions/Xfixes.h
+
+    typedef struct {
+        short	    x, y;
+        unsigned short  width, height;
+        unsigned short  xhot, yhot;
+        unsigned long   cursor_serial;
+        unsigned long   *pixels;
+    if XFIXES_MAJOR >= 2
+        Atom	    atom;	/* Version >= 2 only */
+        const char	*name;	/* Version >= 2 only */
+    endif
+    } XFixesCursorImage;
+    """
+    _fields_ = [('x', ctypes.c_short),
+                ('y', ctypes.c_short),
+                ('width', ctypes.c_ushort),
+                ('height', ctypes.c_ushort),
+                ('xhot', ctypes.c_ushort),
+                ('yhot', ctypes.c_ushort),
+                ('cursor_serial', ctypes.c_ulong),
+                ('pixels', PIXEL_DATA_PTR),
+                ('atom', Atom),
+                ('name', ctypes.c_char_p)]
+
+
+class Display(ctypes.Structure):
+    pass
+
+
+class Xcursor:
+    display = None
+
+    def __init__(self, display=None):
+        if not display:
+            try:
+                display = os.environ["DISPLAY"].encode("utf-8")
+            except KeyError:
+                raise Exception("$DISPLAY not set.")
+
+        # XFixeslib = ctypes.CDLL('libXfixes.so')
+        XFixes = ctypes.util.find_library("Xfixes")
+        if not XFixes:
+            raise Exception("No XFixes library found.")
+        self.XFixeslib = ctypes.cdll.LoadLibrary(XFixes)
+
+        # xlib = ctypes.CDLL('libX11.so.6')
+        x11 = ctypes.util.find_library("X11")
+        if not x11:
+            raise Exception("No X11 library found.")
+        self.xlib = ctypes.cdll.LoadLibrary(x11)
+
+        # Define ctypes' version of XFixesGetCursorImage function
+        XFixesGetCursorImage = self.XFixeslib.XFixesGetCursorImage
+        XFixesGetCursorImage.restype = ctypes.POINTER(XFixesCursorImage)
+        XFixesGetCursorImage.argtypes = [ctypes.POINTER(Display)]
+        self.XFixesGetCursorImage = XFixesGetCursorImage
+
+        XOpenDisplay = self.xlib.XOpenDisplay
+        XOpenDisplay.restype = ctypes.POINTER(Display)
+        XOpenDisplay.argtypes = [ctypes.c_char_p]
+
+        if not self.display:
+            self.display = self.xlib.XOpenDisplay(display)  # (display) or (None)
+
+    def argbdata_to_pixdata(self, data, len):
+        if data == None or len < 1: return None
+
+        # Create byte array
+        b = array.array('b', b'\x00' * 4 * len)
+
+        offset, i = 0, 0
+        while i < len:
+            argb = data[i] & 0xffffffff
+            rgba = (argb << 8) | (argb >> 24)
+            b1 = (rgba >> 24) & 0xff
+            b2 = (rgba >> 16) & 0xff
+            b3 = (rgba >> 8) & 0xff
+            b4 = rgba & 0xff
+
+            struct.pack_into("=BBBB", b, offset, b1, b2, b3, b4)
+            offset = offset + 4
+            i = i + 1
+
+        return b
+
+    def getCursorImageData(self):
+        # Call the function. Read data of cursor/mouse-pointer.
+        cursor_data = self.XFixesGetCursorImage(self.display)
+
+        if not (cursor_data and cursor_data[0]):
+            raise Exception("Cannot read XFixesGetCursorImage()")
+
+        # Note: cursor_data is a pointer, take cursor_data[0]
+        return cursor_data[0]
+
+    def getCursorImageArray(self):
+        data = self.getCursorImageData()
+        # x, y = data.x, data.y
+        height, width = data.height, data.width
+
+        bytearr = self.argbdata_to_pixdata(data.pixels, height * width)
+
+        imgarray = np.array(bytearr, dtype=np.uint8)
+        imgarray = imgarray.reshape(height, width, 4)
+        del bytearr
+
+        return imgarray
+
+    def getCursorImageArrayFast(self):
+        data = self.getCursorImageData()
+        # x, y = data.x, data.y
+        height, width = data.height, data.width
+
+        bytearr = ctypes.cast(data.pixels, ctypes.POINTER(ctypes.c_ulong * height * width))[0]
+        imgarray = np.array(bytearray(bytearr))
+        imgarray = imgarray.reshape(height, width, 8)[:, :, (0, 1, 2, 3)]
+        del bytearr
+
+        return imgarray
+
+    def saveImage(self, imgarray, text):
+        from PIL import Image
+        img = Image.fromarray(imgarray)
+        img.save(text)
+
+
+if __name__ == "__main__":
+    cursor = Xcursor()
+    imgarray = cursor.getCursorImageArrayFast()
+    cursor.saveImage(imgarray, 'cursor_image.png')

desktop_env/server/requirements.txt

@@ -4,3 +4,4 @@ Pillow==10.1.0
 git+https://github.com/moses-palmer/pynput.git@refs/pull/541/head # to make sure that it works on Apple Silicon
 requests
 flask
+numpy

main.py

@@ -9,7 +9,7 @@ def human_agent():
 
     with open("evaluation_examples/examples/libreoffice_calc/f9584479-3d0d-4c79-affa-9ad7afdd8850.json", "r") as f:
         example = json.load(f)
-    example["snapshot"] = "base_setup2"
+    example["snapshot"] = "base_setup3"
 
     env = DesktopEnv(
         path_to_vm=r"C:\Users\tianbaox\Documents\Virtual Machines\Ubuntu\Ubuntu.vmx",