Remove delay of getting device screen

alune/adb.py

@@ -2,33 +2,41 @@
 Module for all ADB (Android Debug Bridge) related methods.
 """
 
+import asyncio
+import atexit
 import os.path
 import random
 
 from adb_shell.adb_device import AdbDeviceTcp
 from adb_shell.adb_device_async import AdbDeviceTcpAsync
 from adb_shell.auth.keygen import keygen
 from adb_shell.auth.sign_pythonrsa import PythonRSASigner
+from adb_shell.exceptions import TcpTimeoutException
+import av
+from av.error import InvalidDataError  # pylint: disable=no-name-in-module
 import cv2
 from loguru import logger
 import numpy
 from numpy import ndarray
 import psutil
 
 from alune import helpers
+from alune.config import AluneConfig
 from alune.images import ClickButton
 from alune.images import ImageButton
 from alune.screen import BoundingBox
 from alune.screen import ImageSearchResult
 
 
-class ADB:
+# The amount of attributes is fine in my opinion.
+# We could split off screen recording into its own class, but I don't see the need to.
+class ADB:  # pylint: disable=too-many-instance-attributes
     """
     Class to hold the connection to an ADB connection via TCP.
     USB connection is possible, but not supported at the moment.
     """
 
-    def __init__(self):
+    def __init__(self, config: AluneConfig):
         """
         Initiates base values for the ADB instance.
         """
@@ -37,16 +45,23 @@ def __init__(self):
         self._random = random.Random()
         self._rsa_signer = None
         self._device = None
+        self._config = config
+        self._default_port = config.get_adb_port()
 
-    async def load(self, port: int):
+        if not config.should_use_screen_record():
+            return
+
+        self._video_codec = av.codec.CodecContext.create("h264", "r")
+        self._is_screen_recording = False
+        self._should_stop_screen_recording = False
+        self._latest_frame = None
+
+    async def load(self):
         """
         Load the RSA signer and attempt to connect to a device via ADB TCP.
-
-        Args:
-            port: The port to attempt to connect to.
         """
         await self._load_rsa_signer()
-        await self._connect_to_device(port)
+        await self._connect_to_device(self._default_port)
 
     async def _load_rsa_signer(self):
         """
@@ -97,6 +112,25 @@ async def scan_localhost_devices(self) -> int | None:
         logger.warning("No local device was found. Make sure ADB is enabled in your emulator's settings.")
         return None
 
+    def mark_screen_record_for_close(self):
+        """
+        Tells the screen recording to close itself when possible.
+        """
+        if self._is_screen_recording:
+            self._should_stop_screen_recording = True
+            self._is_screen_recording = False
+
+    def create_screen_record_task(self):
+        """
+        Create the screen recording task. Will not start recording if there's already a recording.
+        """
+        if self._is_screen_recording:
+            return
+
+        self._should_stop_screen_recording = False
+        asyncio.create_task(self.__screen_record())
+        atexit.register(self.mark_screen_record_for_close)
+
     async def _connect_to_device(self, port: int, retry_with_scan: bool = True):
         """
         Connect to the device via TCP.
@@ -136,7 +170,7 @@ async def get_screen_size(self) -> str:
         Returns:
              A string containing 'WIDTHxHEIGHT'.
         """
-        shell_output = await self._device.shell("wm size | awk 'END{print $3}'")
+        shell_output = await self._wrap_shell_call("wm size | awk 'END{print $3}'")
         return shell_output.replace("\n", "")
 
     async def get_screen_density(self) -> str:
@@ -146,20 +180,20 @@ async def get_screen_density(self) -> str:
         Returns:
              A string containing the pixel density.
         """
-        shell_output = await self._device.shell("wm density | awk 'END{print $3}'")
+        shell_output = await self._wrap_shell_call("wm density | awk 'END{print $3}'")
         return shell_output.replace("\n", "")
 
     async def set_screen_size(self):
         """
         Set the screen size to 1280x720.
         """
-        await self._device.shell("wm size 1280x720")
+        await self._wrap_shell_call("wm size 1280x720")
 
     async def set_screen_density(self):
         """
         Set the screen pixel density to 240.
         """
-        await self._device.shell("wm density 240")
+        await self._wrap_shell_call("wm density 240")
 
     async def get_memory(self) -> int:
         """
@@ -168,13 +202,25 @@ async def get_memory(self) -> int:
         Returns:
             The memory of the device in kB.
         """
-        shell_output = await self._device.shell("grep MemTotal /proc/meminfo | awk '{print $2}'")
+        shell_output = await self._wrap_shell_call("grep MemTotal /proc/meminfo | awk '{print $2}'")
         return int(shell_output)
 
     async def get_screen(self) -> ndarray | None:
         """
         Gets a ndarray which contains the values of the gray-scaled pixels
-        currently on the screen.
+        currently on the screen. Uses buffered frames from screen recording, available instantly.
+
+        Returns:
+            The ndarray containing the gray-scaled pixels. Is None until the first screen record frame is processed.
+        """
+        if self._config.should_use_screen_record():
+            return self._latest_frame
+        return await self._get_screen_capture()
+
+    async def _get_screen_capture(self) -> ndarray | None:
+        """
+        Gets a ndarray which contains the values of the gray-scaled pixels
+        currently on the screen. Uses screencap, so will take some processing time.
 
         Returns:
             The ndarray containing the gray-scaled pixels.
@@ -237,7 +283,7 @@ async def click(self, x: int, y: int):
         """
         # input tap x y comes with the downtime of tapping too fast for the game sometimes,
         # so we swipe on the same coordinate to simulate a longer press with a random duration.
-        await self._device.shell(f"input swipe {x} {y} {x} {y} {self._random.randint(60, 120)}")
+        await self._wrap_shell_call(f"input swipe {x} {y} {x} {y} {self._random.randint(60, 120)}")
 
     async def is_tft_installed(self) -> bool:
         """
@@ -246,7 +292,7 @@ async def is_tft_installed(self) -> bool:
         Returns:
             Whether the TFT package is in the list of the installed packages.
         """
-        shell_output = await self._device.shell(f"pm list packages | grep {self.tft_package_name}")
+        shell_output = await self._wrap_shell_call(f"pm list packages | grep {self.tft_package_name}")
         if not shell_output:
             return False
 
@@ -273,14 +319,14 @@ async def is_tft_active(self) -> bool:
         Returns:
              Whether TFT is the currently active window.
         """
-        shell_output = await self._device.shell("dumpsys window | grep -E 'mCurrentFocus' | awk '{print $3}'")
+        shell_output = await self._wrap_shell_call("dumpsys window | grep -E 'mCurrentFocus' | awk '{print $3}'")
         return shell_output.split("/")[0].replace("\n", "") == self.tft_package_name
 
     async def start_tft_app(self):
         """
         Start TFT using the activity manager (am).
         """
-        await self._device.shell(f"am start -n {self.tft_package_name}/{self._tft_activity_name}")
+        await self._wrap_shell_call(f"am start -n {self.tft_package_name}/{self._tft_activity_name}")
 
     async def get_tft_version(self) -> str:
         """
@@ -289,12 +335,91 @@ async def get_tft_version(self) -> str:
         Returns:
             The versionName of the tft package.
         """
-        return await self._device.shell(
+        return await self._wrap_shell_call(
             f"dumpsys package {self.tft_package_name} | grep versionName | sed s/[[:space:]]*versionName=//g"
         )
 
     async def go_back(self):
         """
         Send a back key press event to the device.
         """
-        await self._device.shell("input keyevent 4")
+        await self._wrap_shell_call("input keyevent 4")
+
+    async def _wrap_shell_call(self, shell_command: str, retries: int = 0):
+        """
+        Wrapper for shell commands to catch timeout exceptions.
+        Retries 3 times with incremental backoff.
+
+        Args:
+            shell_command: The shell command to call.
+            retries: Optional, the amount of attempted retries so far.
+
+        Returns:
+            The output of the shell command.
+        """
+        try:
+            return await self._device.exec_out(shell_command)
+        except TcpTimeoutException:
+            if retries == 3:
+                raise
+            logger.debug(f"Timed out while calling '{shell_command}', retrying {3 - retries} times.")
+            await asyncio.sleep(1 + (1 * retries))
+            return await self._wrap_shell_call(shell_command, retries=retries + 1)
+
+    async def __convert_frame_to_cv2(self, frame_bytes: bytes):
+        """
+        Convert frame bytes to a CV2 compatible gray image.
+
+        Args:
+            frame_bytes: Byte output of the screen record session.
+        """
+        packets = self._video_codec.parse(frame_bytes)
+        if not packets:
+            return
+
+        try:
+            frames = self._video_codec.decode(packets[0])
+        except InvalidDataError:
+            return
+        if not frames:
+            return
+
+        self._latest_frame = frames[0].to_ndarray(format="gray8").copy()  # Change to bgr24 if color is ever needed.
+
+    async def __write_frame_data(self):
+        """
+        Start a streaming shell that outputs screenrecord frame bytes and store it as a cv2 compatible image.
+        """
+        # output-format h264 > H264 is the only format that outputs to console which we can work with.
+        # time-limit 10 > Restarts screen recording every 10 seconds instead of every 180. Fixes compression artifacts.
+        # bit-rate 16M > 16_000_000 Mbps, could probably be lowered or made configurable, but works well.
+        # - at the end makes screenrecord output to console, if format is h264.
+        async for data in self._device.streaming_shell(
+            command="screenrecord --time-limit 8 --output-format h264 --bit-rate 16M --size 1280x720 -", decode=False
+        ):
+            if self._should_stop_screen_recording:
+                break
+
+            await self.__convert_frame_to_cv2(data)
+
+    async def __screen_record(self):
+        """
+        Start the screen record session. Restarts itself until an external value stops it.
+        """
+        if self._is_screen_recording:
+            return
+
+        await self._device.exec_out("pkill -2 screenrecord")
+        logger.debug("Screen record starting.")
+
+        self._is_screen_recording = True
+        while not self._should_stop_screen_recording:
+            try:
+                await self.__write_frame_data()
+            except TcpTimeoutException:
+                logger.warning("Timed out while re-/starting screen record, waiting 5 seconds.")
+                await asyncio.sleep(5)
+
+        logger.debug("Screen record stopped.")
+        await self._device.exec_out("pkill -2 screenrecord")
+        self._is_screen_recording = False

alune/config.py

@@ -212,3 +212,12 @@ def get_queue_timeout(self) -> int:
             The queue timeout in seconds.
         """
         return self._config["queue_timeout"]
+
+    def should_use_screen_record(self) -> bool:
+        """
+        Get the screen record setting the user wants.
+
+        Returns:
+            Whether we should use screen recording.
+        """
+        return self._config["screen_record"]["enabled"]

alune/resources/config.yaml

@@ -31,6 +31,14 @@ surrender_early: false
 # Default value : 0 (disabled = surrender as fast as possible)
 surrender_random_delay: 0
 
+# Screen recording settings.
+screen_record:
+  # Whether we should use screen recording.
+  # If false, the bot will fall back to screen capture.
+  # Screen recording saves ~600ms per screen check, but may occasionally cause a timeout.
+  # The timeout will be caught and handled.
+  enabled: true
+
 # Configuration for chance events in the bot, in percent.
 chances:
   # The chance % to buy experience per check - configuring 50 means it's a 50% chance.
@@ -44,6 +52,6 @@ adb_port: 5555
 
 # Changing these below values manually can potentially break the bot, so don't!
 # Version of the YAML.
-version: 7
+version: 9
 # Version of the TFT set.
 set: 13

main.py

@@ -291,7 +291,9 @@ async def loop_disconnect_wrapper(adb_instance: ADB, config: AluneConfig):
         await loop(adb_instance, config)
     except TcpTimeoutException:
         logger.warning("ADB device was disconnected, attempting one reconnect...")
-        await adb_instance.load(config.get_adb_port())
+        adb_instance.mark_screen_record_for_close()
+        await adb_instance.load()
+        adb_instance.create_screen_record_task()
         if not adb_instance.is_connected():
             raise_and_exit("Could not reconnect. Please check your emulator for any errors. Exiting.")
         logger.info("Reconnected to device, continuing main loop.")
@@ -539,15 +541,24 @@ async def main():
 
     await check_version()
 
-    adb_instance = ADB()
+    adb_instance = ADB(config)
 
-    await adb_instance.load(config.get_adb_port())
+    await adb_instance.load()
     if not adb_instance.is_connected():
         logger.error("There is no ADB device ready. Exiting.")
         return
 
     logger.debug("ADB is connected, checking phone and app details")
     await check_phone_preconditions(adb_instance)
+
+    if config.should_use_screen_record():
+        logger.info("The bot will use live screen recording for image searches.")
+        adb_instance.create_screen_record_task()
+        while await adb_instance.get_screen() is None:
+            logger.debug("Waiting for frame data to become available...")
+            await asyncio.sleep(0.5)
+        logger.debug("Frames are now available.")
+
     logger.info("Connected to ADB and device is set up correctly, starting main loop.")
 
     if config.should_surrender():

pyproject.toml

@@ -23,6 +23,7 @@ dependencies = [
     "psutil==6.0.0",
     "keyboard==0.13.5",
     "imutils==0.5.4",
+    "av==14.0.0",
 ]
 
 [project.optional-dependencies]