__init__.py

import base64
import ctypes
import os
import platform
import signal
import subprocess
import sys
import threading
from collections import deque
from functools import cache
import av
from time import sleep as sleep_
from math import floor


def sleep(secs):
    try:
        if secs == 0:
            return
        maxrange = 50 * secs
        if isinstance(maxrange, float):
            sleeplittle = floor(maxrange)
            sleep_((maxrange - sleeplittle) / 50)
            maxrange = int(sleeplittle)
        if maxrange > 0:
            for _ in range(maxrange):
                sleep_(0.016)
    except KeyboardInterrupt:
        return

iswindows = "win" in platform.platform().lower()
if iswindows:
    startupinfo = subprocess.STARTUPINFO()
    startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
    startupinfo.wShowWindow = subprocess.SW_HIDE
    creationflags = subprocess.CREATE_NO_WINDOW
    invisibledict = {
        "startupinfo": startupinfo,
        "creationflags": creationflags,
        "start_new_session": True,
    }
    from ctypes import wintypes

    windll = ctypes.LibraryLoader(ctypes.WinDLL)
    kernel32 = windll.kernel32
    _GetShortPathNameW = kernel32.GetShortPathNameW
    _GetShortPathNameW.argtypes = [wintypes.LPCWSTR, wintypes.LPWSTR, wintypes.DWORD]
    _GetShortPathNameW.restype = wintypes.DWORD
else:
    invisibledict = {}


@cache
def get_short_path_name(long_name):
    try:
        if not iswindows:
            return long_name
        output_buf_size = 4096
        output_buf = ctypes.create_unicode_buffer(output_buf_size)
        _ = _GetShortPathNameW(long_name, output_buf, output_buf_size)
        return output_buf.value
    except Exception as e:
        sys.stderr.write(f"{e}\n")
        return long_name


def killthread(threadobject):
    # based on https://pypi.org/project/kthread/
    if not threadobject.is_alive():
        return True
    tid = -1
    for tid1, tobj in threading._active.items():
        if tobj is threadobject:
            tid = tid1
            break
    if tid == -1:
        sys.stderr.write(f"{threadobject} not found")
        return False
    res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
        ctypes.c_long(tid), ctypes.py_object(SystemExit)
    )
    if res == 0:
        return False
    elif res != 1:
        ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, 0)
        return False
    return True


def send_ctrl_commands(pid, command=0):
    if iswindows:
        commandstring = r"""import ctypes, sys; CTRL_C_EVENT, CTRL_BREAK_EVENT, CTRL_CLOSE_EVENT, CTRL_LOGOFF_EVENT, CTRL_SHUTDOWN_EVENT = 0, 1, 2, 3, 4; kernel32 = ctypes.WinDLL("kernel32", use_last_error=True); (lambda pid, cmdtosend=CTRL_C_EVENT: [kernel32.FreeConsole(), kernel32.AttachConsole(pid), kernel32.SetConsoleCtrlHandler(None, 1), kernel32.GenerateConsoleCtrlEvent(cmdtosend, 0), sys.exit(0) if isinstance(pid, int) else None])(int(sys.argv[1]), int(sys.argv[2]) if len(sys.argv) > 2 else None) if __name__ == '__main__' else None"""
        subprocess.Popen(
            [sys.executable, "-c", commandstring, str(pid), str(command)],
            **invisibledict,
        )
    else:
        os.kill(pid, signal.SIGINT)


class StopDescriptor:
    def __get__(self, instance, owner):
        return instance.__dict__[self.name]

    def __set__(self, instance, value):
        if not value:
            instance.__dict__[self.name] = False
        else:
            instance.__dict__[self.name] = True
            instance.stop_capture()

    def __delete__(self, instance):
        sys.stderr.write("Cannot be deleted!")

    def __set_name__(self, owner, name):
        self.name = name


class AdbFastScreenshots:
    stop_recording = StopDescriptor()

    def __init__(
        self,
        adb_path,
        device_serial,
        time_interval=179,
        width=1600,
        height=900,
        bitrate="20M",
        use_busybox=False,
        connect_to_device=True,
        screenshotbuffer=10,
        go_idle=0,
    ):
        r"""Capture Android device screen using ADB's screenrecord with high frame rate.

            This class allows capturing the screen of an Android device using ADB's screenrecord
            command with an improved frame rate. It continuously captures frames from the device
            and provides them as NumPy arrays to the caller.

            Args:
                adb_path (str): The path to the ADB executable.
                device_serial (str): The serial number of the target Android device.
                time_interval (int): The maximum duration, in seconds, for each screen recording session (up to a maximum of 180 seconds). After reaching this time limit, a new recording session automatically starts without causing interruptions to the user experience.
                width (int): The width of the captured screen.
                height (int): The height of the captured screen.
                bitrate (str): The bitrate for screen recording (e.g., "20M" for 20Mbps).
                use_busybox (bool): Whether to use BusyBox for base64 encoding.
                connect_to_device (bool): Whether to connect to the device using ADB.
                screenshotbuffer (int): The size of the frame buffer to store the last captured frames.
                go_idle (float): The idle time (in seconds) when no new frames are available. # higher value -> less fps, but also less CPU usage.

            Attributes:
                stop_recording (bool): Control attribute to stop the screen capture.

            Methods:
                stop_capture(): Stops the screen capture.

            Usage:
                import cv2
                from adbnativeblitz import AdbFastScreenshots

                with AdbFastScreenshots(
                    adb_path=r"C:\Android\android-sdk\platform-tools\adb.exe",
                    device_serial="127.0.0.1:5555",
                    time_interval=179,
                    width=1600,
                    height=900,
                    bitrate="20M",
                    use_busybox=False,
                    connect_to_device=True,
                    screenshotbuffer=10,
                    go_idle=0,
                ) as adbscreen:
                    for image in adbscreen:
                        cv2.imshow("CV2 WINDOW", image)
                        if cv2.waitKey(1) & 0xFF == ord("q"):
                            break
                cv2.destroyAllWindows()


            Note:
            - The `AdbFastScreenshots` class should be used in a context manager (`with` statement).
            - The `stop_capture()` method can be called to stop the screen capture.
            - The frames are continuously captured and provided in the form of NumPy arrays.
            - The class aims to achieve a higher frame rate by avoiding slow subprocess creation
              for each screen capture session.
            """

        self.stop_recording = False
        self.size = f"{width}x{height}"
        self.width = width
        self.height = height
        self.timelimit = time_interval
        self.bitrate = bitrate
        self.use_busybox = use_busybox
        self.adb_path = get_short_path_name(adb_path)
        self.device_serial = device_serial
        if connect_to_device:
            subprocess.run(
                [self.adb_path, "connect", self.device_serial], **invisibledict
            )
        self.threadlock = threading.Lock()
        self.codec = av.codec.CodecContext.create("h264", "r")
        self.lastframes = deque([], screenshotbuffer)
        self.command_to_execute = (
            f"""#!/bin/bash
startscreenrecord() {{

    screenrecord --output-format=h264 --time-limit "$1" --size "$2" --bit-rate "$3" -
}}

time_interval={self.timelimit}
size="{self.size}"
bitrate="{self.bitrate}"
#screenrecord --output-format=h264 --time-limit 1 --size "$size" --bit-rate "$bitrate" -
while true; do
    startscreenrecord $time_interval "$size" "$bitrate" 
done"""
            + "\n"
        )
        self.base64cmd = self.format_adb_command(
            self.command_to_execute,
            su=False,
            exitcommand="",
            errors="strict",
        )
        self.p = None
        self.threadstdout = None
        self.framecounter = 0
        self.go_idle = go_idle

    def format_adb_command(
        self,
        cmd,
        su=False,
        exitcommand="DONE",
        errors="strict",
    ):
        if su:
            cmd = f"su -- {cmd}"
        if exitcommand:
            cmd = cmd.rstrip() + f"\necho {exitcommand}\n"
        nolimitcommand = []
        base64_command = base64.standard_b64encode(cmd.encode("utf-8", errors)).decode(
            "utf-8", errors
        )
        nolimitcommand.extend(["echo", base64_command, "|"])
        if self.use_busybox:
            nolimitcommand.extend(["busybox"])
        nolimitcommand.extend(["base64", "-d", "|", "sh"])

        return " ".join(nolimitcommand) + "\n"

    def _start_capturing(self):
        def _execute_stdout_read():
            try:
                for q in iter(self.p.stdout.readline, b""):
                    if iswindows:
                        q = q.replace(b"\r\n", b"\n")
                    if q:
                        alldata.append(q)
                    if alldata:
                        joineddata = b"".join(alldata)
                        try:
                            packets = self.codec.parse(joineddata)
                            if packets:
                                for pack in packets:
                                    frames = self.codec.decode(pack)
                                    for frame in frames:
                                        nparray = (
                                            frame.to_rgb()
                                            .reformat(
                                                width=self.width,
                                                height=self.height,
                                                format="bgr24",
                                            )
                                            .to_ndarray()
                                        )
                                        try:
                                            self.threadlock.acquire()
                                            self.lastframes.append(nparray)
                                            self.framecounter += 1

                                        finally:
                                            try:
                                                self.threadlock.release()
                                            except Exception as e:
                                                sys.stderr.write(f"{e}\n")
                            alldata.clear()
                        except Exception as e:
                            sys.stderr.write(f"{e}\n")
            except Exception as e:
                sys.stderr.write(f"{e}\n")

        self.p = subprocess.Popen(
            [self.adb_path, "-s", self.device_serial, "shell", self.base64cmd],
            stderr=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stdin=subprocess.DEVNULL,
            bufsize=0,
            **invisibledict,
        )
        alldata = []
        self.threadstdout = threading.Thread(target=_execute_stdout_read)
        self.threadstdout.daemon = True
        self.threadstdout.start()

    def _stop_capture(self):
        try:
            if iswindows:
                subprocess.Popen(f"taskkill /F /PID {self.p.pid} /T", **invisibledict)
        except:
            pass
        try:
            self.p.stdout.close()
        except:
            pass
        try:
            killthread(self.threadstdout)
        except:
            pass

    def stop_capture(self):
        send_ctrl_commands(self.p.pid, command=0)
        try:
            sleep(1)
        except KeyboardInterrupt:
            pass
        self._stop_capture()

    def __iter__(self):
        oldframecounter = 0

        self._start_capturing()
        sleep(0.05)
        while not self.stop_recording:
            if not self.lastframes:
                sleep(0.005)
                continue
            yield self.lastframes[-1].copy()
            if oldframecounter == self.framecounter:
                if self.go_idle:
                    sleep(self.go_idle)
            oldframecounter = self.framecounter

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.stop_recording = True