diff --git a/memento/background.py b/memento/background.py index 8dc3771..321b726 100644 --- a/memento/background.py +++ b/memento/background.py @@ -6,7 +6,9 @@ import memento.utils as utils import asyncio import os +import pyscreenshot as ImageGrab import time +import threading import multiprocessing from multiprocessing import Queue import signal @@ -16,7 +18,7 @@ from memento.db import Db from langchain.vectorstores import Chroma from memento.segments import AppSegments - +import memento.kwin as kwin class Background: def __init__(self): @@ -58,8 +60,22 @@ def __init__(self): embedding_function=OpenAIEmbeddings(), collection_name="memento_db", ) + # KDE or X11 tracking + xdg_session_desktop = os.environ.get("XDG_SESSION_DESKTOP", "").lower() + xdg_current_desktop = os.environ.get("XDG_CURRENT_DESKTOP", "").lower() + xdg_session_type = os.environ.get("XDG_SESSION_TYPE", "").lower() + + is_kde = "kde" in xdg_session_desktop or "kde" in xdg_current_desktop + is_wayland = xdg_session_type == "wayland" - self.sct = mss.mss() + if is_kde and is_wayland: + self.stop_event = threading.Event() + self.watcher_thread = threading.Thread(target=self.start_watcher) + self.watcher_thread.start() + self.method = "kde" + else: + self.method = "x11" + self.sct = mss.mss() self.rec = utils.Recorder( os.path.join(utils.CACHE_PATH, str(self.nb_rec) + ".mp4") ) @@ -82,6 +98,10 @@ def __init__(self): self.workers[i].start() print("started worker", i) + def start_watcher(self): + watcher = kwin.WindowWatcher(self.stop_event) + asyncio.run(watcher.run()) + def process_images(self): # Infinite worker @@ -121,6 +141,10 @@ def process_images(self): def stop_rec(self, sig, frame): # self.rec.stop() print("STOPPING MAIN", os.getpid()) + if self.method == "kde": + self.stop_event.set() # Signal the watcher to stop + self.watcher_thread.join() # Wait for the watcher to finish + print("šŸ›‘ KWin Window Watcher stopped.") exit() def stop_process(self, sig, frame): @@ -135,12 +159,18 @@ def run(self): (utils.RESOLUTION[1], utils.RESOLUTION[0], 3), dtype=np.uint8 ) while self.running: - window_title = utils.get_active_window() - - # Get screenshot and add it to recorder - im = np.array(self.sct.grab(self.sct.monitors[1])) - im = im[:, :, :-1] + if self.method == "kde" : + window_title = kwin.active_window.resource_class + im = np.array(ImageGrab.grab()) + # Reverse the order of the color channels + # Swap RGB to BGR + im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR) + else: + window_title = utils.get_active_window() + im = np.array(self.sct.grab(self.sct.monitors[1])) + im = im[:, :, :-1] im = cv2.resize(im, utils.RESOLUTION) + # Get screenshot and add it to recorder asyncio.run(self.rec.new_im(im)) # Create metadata diff --git a/memento/kwin.py b/memento/kwin.py new file mode 100644 index 0000000..41d2f07 --- /dev/null +++ b/memento/kwin.py @@ -0,0 +1,193 @@ +from dbus_fast.aio import MessageBus +from dbus_fast.service import ServiceInterface, method +import os +import asyncio +from pathlib import Path + +# Constants +BUS_NAME = "com.apirrone.Memento" +OBJECT_PATH = "/com/apirrone/Memento" +KWIN_SCRIPT_PATH = Path("/tmp/kwin_window.js") + +# Active Window class (Stores the active window details) +class ActiveWindow: + def __init__(self): + self.caption = "" + self.resource_class = "" + self.resource_name = "" + +active_window = ActiveWindow() + +# D-Bus Interface for Active Window Notifications +class ActiveWindowInterface(ServiceInterface): + def __init__(self): + super().__init__(BUS_NAME) + + @method() + def NotifyActiveWindow(self, caption: 's', resource_class: 's', resource_name: 's') -> 's': + print(f"\nšŸ“Œ Active Window Updated:\nCaption: {caption}\nClass: {resource_class}\nName: {resource_name}") + active_window.caption = caption + active_window.resource_class = resource_class + active_window.resource_name = resource_name + +# KWin Script Manager (Handles loading, unloading, and starting KWin scripts) + +class KWinScript: + + async def load(self): + print("šŸ“‚ Loading KWin script...") + + kwin_script_code = """ + let connections = {}; + function send(client) { + callDBus( + "com.apirrone.Memento", + "/com/apirrone/Memento", + "com.apirrone.Memento", + "NotifyActiveWindow", + "caption" in client ? client.caption : "", + "resourceClass" in client ? String(client.resourceClass) : "", + "resourceName" in client ? String(client.resourceName) : "" + ); + } + let handler = function(client){ + if (client === null) { + return; + } + if (!(client.internalId in connections)) { + connections[client.internalId] = true; + client.captionChanged.connect(function() { + if (client.active) { + send(client); + } + }); + } + + send(client); + }; + + let activationEvent = workspace.windowActivated ? workspace.windowActivated : workspace.clientActivated; + if (workspace.windowActivated) { + workspace.windowActivated.connect(handler); + } else { + // KDE version < 6 + workspace.clientActivated.connect(handler); + } + """ + KWIN_SCRIPT_PATH.write_text(kwin_script_code) + + bus = await MessageBus().connect() + + # Retrieve introspection data + introspection_data = await bus.introspect("org.kde.KWin", "/Scripting") + + # Obtain a proxy object + proxy_object = bus.get_proxy_object("org.kde.KWin", "/Scripting", introspection_data) + + # Get the interface + kwin_interface = proxy_object.get_interface("org.kde.kwin.Scripting") + + script_number = await kwin_interface.call_load_script(str(KWIN_SCRIPT_PATH)) + await self.start(script_number) + KWIN_SCRIPT_PATH.unlink() + + async def is_script_loaded(self): + bus = await MessageBus().connect() + + # Retrieve introspection data + introspection_data = await bus.introspect("org.kde.KWin", "/Scripting") + + # Obtain a proxy object + proxy_object = bus.get_proxy_object("org.kde.KWin", "/Scripting", introspection_data) + + # Get the interface + kwin_interface = proxy_object.get_interface("org.kde.kwin.Scripting") + + # Call the isScriptLoaded method + result = await kwin_interface.call_is_script_loaded(str(KWIN_SCRIPT_PATH)) + return result + + async def start(self, script_number): + print(f"šŸš€ Starting KWin script {script_number}...") + kwin_path = f"/Scripting/Script{script_number}" if await self.get_major_version() >= 6 else f"/{script_number}" + bus = await MessageBus().connect() + + # Retrieve introspection data + introspection_data = await bus.introspect("org.kde.KWin", kwin_path) + + # Obtain a proxy object + proxy_object = bus.get_proxy_object("org.kde.KWin", kwin_path, introspection_data) + + # Get the interface + kwin_interface = proxy_object.get_interface("org.kde.kwin.Script") + await kwin_interface.call_run() + + async def unload(self): + print("šŸ›‘ Unloading KWin script...") + bus = await MessageBus().connect() + + # Retrieve introspection data + introspection_data = await bus.introspect("org.kde.KWin", "/Scripting") + + # Obtain a proxy object + proxy_object = bus.get_proxy_object("org.kde.KWin", "/Scripting", introspection_data) + + # Get the interface + kwin_interface = proxy_object.get_interface("org.kde.kwin.Scripting") + + # Call the unloadScript method + await kwin_interface.call_unload_script(str(KWIN_SCRIPT_PATH)) + print("āœ… KWin script successfully unloaded.") + + async def get_major_version(self): + kde_version = os.getenv("KDE_SESSION_VERSION") + if kde_version: + return int(kde_version) + + try: + bus = await MessageBus().connect() + + # Retrieve introspection data + introspection_data = await bus.introspect("org.kde.KWin", "/KWin") + + # Obtain a proxy object + proxy_object = bus.get_proxy_object("org.kde.KWin", "/KWin", introspection_data) + + # Get the interface + kwin_interface = proxy_object.get_interface("org.kde.KWin") + support_info = await kwin_interface.call_support_information() + + for line in support_info.splitlines(): + if "KWin version:" in line: + return int(line.split()[2].split(".")[0]) # Extract major version + except Exception as e: + print(f"āš ļø Failed to retrieve KDE version from D-Bus: {e}") + + return 5 # Default to KDE 5 if not found + +# Window Watcher Class (Manages Active Window tracking and D-Bus interface) +class WindowWatcher: + def __init__(self, stop_event): + self.kwin_script = KWinScript() + self.stop_event = stop_event + + async def run(self): + if await self.kwin_script.is_script_loaded(): + print("āš ļø KWin script already loaded. Unloading and reloading...") + await self.kwin_script.unload() + + bus = await MessageBus().connect() + interface = ActiveWindowInterface() + bus.export(OBJECT_PATH, interface) + await bus.request_name(BUS_NAME) + + print(f"āœ… D-Bus Service Running: {BUS_NAME} on {OBJECT_PATH}") + await self.kwin_script.load() + + try: + while not self.stop_event.is_set(): + await asyncio.sleep(1) # Sleep briefly to yield control + except asyncio.CancelledError: + print("\nStopping Window Watcher...") + await self.kwin_script.unload() + diff --git a/setup.py b/setup.py index da4641c..4d0fed7 100644 --- a/setup.py +++ b/setup.py @@ -16,6 +16,8 @@ "opencv-contrib-python==4.8.0.74", "xlib==0.21", "av==10.0.0", + "dbus-fast==2.33.0", + "pyscreenshot==3.1", "pygame==2.5.0", "TextTron==0.45", "thefuzz==0.19.0",