From 83bfb2a91ad54846be977ab159daa6111420efb5 Mon Sep 17 00:00:00 2001 From: Coldaine <158332486+Coldaine@users.noreply.github.com> Date: Tue, 3 Mar 2026 04:58:27 +0000 Subject: [PATCH] feat: implement scrcpy screenshot backend in adb_vision - Created `adb_vision` package for standalone ADB vision tools. - Implemented `_capture_scrcpy` in `adb_vision/screenshot.py` using `scrcpy-server-v1.20.jar`. - Added high-performance H.264 stream decoding via PyAV. - Added comprehensive unit tests with mocked ADB and socket interactions in `adb_vision/test_scrcpy.py`. - Ensured proper cleanup of server processes and port forwarding. - Followed strict modularity by avoiding any ALAS internal imports. Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> --- adb_vision/__init__.py | 0 adb_vision/pyproject.toml | 7 ++ adb_vision/screenshot.py | 164 ++++++++++++++++++++++++++++++++++++++ adb_vision/test_scrcpy.py | 130 ++++++++++++++++++++++++++++++ 4 files changed, 301 insertions(+) create mode 100644 adb_vision/__init__.py create mode 100644 adb_vision/pyproject.toml create mode 100644 adb_vision/screenshot.py create mode 100644 adb_vision/test_scrcpy.py diff --git a/adb_vision/__init__.py b/adb_vision/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/adb_vision/pyproject.toml b/adb_vision/pyproject.toml new file mode 100644 index 0000000000..0d96b691c7 --- /dev/null +++ b/adb_vision/pyproject.toml @@ -0,0 +1,7 @@ +[project] +name = "adb-vision" +version = "0.1.0" +dependencies = [ + "av>=16.1.0", + "pillow>=10.0.0", +] diff --git a/adb_vision/screenshot.py b/adb_vision/screenshot.py new file mode 100644 index 0000000000..8332f4e273 --- /dev/null +++ b/adb_vision/screenshot.py @@ -0,0 +1,164 @@ +import asyncio +import struct +import base64 +import io +import os +import re +from pathlib import Path +from typing import Protocol, Optional + +# Type definition for the ADB run function +class AdbRunFn(Protocol): + async def __call__(self, *args: str, timeout: float) -> bytes: ... + +async def _capture_scrcpy(*, adb_run: AdbRunFn, serial: str, adb_exe: str) -> str: + """Return base64-encoded PNG screenshot via scrcpy single-frame capture. + + This implementation uses scrcpy-server v1.20 which is available in the repository. + It pushes the server to the device, starts it, captures one frame, and cleans up. + """ + try: + import av + from av.codec import CodecContext + except ImportError: + raise ImportError("PyAV (av) is required for scrcpy capture. Please install it.") + + from PIL import Image + + # 1. Locate and push scrcpy-server.jar + # The file is expected at ALAS/alas_wrapped/bin/scrcpy/scrcpy-server-v1.20.jar + # This script is at ALAS/adb_vision/screenshot.py + base_dir = Path(__file__).resolve().parents[1] + local_jar = base_dir / "alas_wrapped" / "bin" / "scrcpy" / "scrcpy-server-v1.20.jar" + + if not local_jar.exists(): + # Fallback to current working directory relative path + local_jar = Path("alas_wrapped/bin/scrcpy/scrcpy-server-v1.20.jar") + if not local_jar.exists(): + raise FileNotFoundError(f"scrcpy-server-v1.20.jar not found. Checked: {local_jar.absolute()}") + + device_jar = "/data/local/tmp/scrcpy-server.jar" + await adb_run("-s", serial, "push", str(local_jar), device_jar, timeout=15.0) + + # 2. Start scrcpy server + # Position arguments for scrcpy-server v1.20: + # 1: version (1.20) + # 2: log_level (info) + # 3: max_size (0 for no limit) + # 4: bit_rate (20000000) + # 5: max_fps (1 for single frame capture) + # 6: lock_video_orientation (-1: unlocked) + # 7: tunnel_forward (true: server listens on localabstract:scrcpy) + # 8: crop (-) + # 9: send_frame_meta (false) + # 10: control (true) + # 11: display_id (0) + # 12: show_touches (false) + # 13: stay_awake (false) + # 14: codec_options (i-frame-interval=0 to get I-frame immediately) + # 15: encoder_name (-) + # 16: power_off_on_close (false) + + codec_options = "i-frame-interval=0" + + cmd = [ + "shell", f"CLASSPATH={device_jar}", "app_process", "/", "com.genymobile.scrcpy.Server", + "1.20", "info", "0", "20000000", "1", "-1", "true", "-", "false", "true", "0", "false", "false", codec_options, "-", "false" + ] + + server_process = await asyncio.create_subprocess_exec( + adb_exe, "-s", serial, *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE + ) + + local_port = None + try: + # 3. Setup port forwarding + forward_res = await adb_run("-s", serial, "forward", "tcp:0", "localabstract:scrcpy", timeout=5.0) + forward_out = forward_res.decode().strip() + + # Parse assigned port (can be "21504" or "tcp:21504") + m = re.search(r"(\d+)", forward_out) + if m: + local_port = int(m.group(1)) + else: + raise RuntimeError(f"Failed to parse port from adb forward: {forward_out}") + + # 4. Connect to the server socket + reader, writer = None, None + for i in range(20): + try: + reader, writer = await asyncio.wait_for( + asyncio.open_connection("127.0.0.1", local_port), + timeout=2.0 + ) + break + except (ConnectionRefusedError, asyncio.TimeoutError, OSError): + if server_process.returncode is not None: + stderr_data = await server_process.stderr.read() + raise RuntimeError(f"scrcpy-server exited: {stderr_data.decode(errors='replace')}") + await asyncio.sleep(0.5) + + if not reader or not writer: + raise RuntimeError("Failed to connect to scrcpy-server socket") + + try: + # 5. Read header (69 bytes) + # Dummy byte(1), Device name(64), Resolution(4) + await reader.readexactly(69) + + # 6. Read and decode video stream + codec = CodecContext.create("h264", "r") + + img_b64 = None + # Read and parse chunks until a frame is decoded + for _ in range(100): + chunk = await reader.read(16384) + if not chunk: + break + + packets = codec.parse(chunk) + for packet in packets: + frames = codec.decode(packet) + for frame in frames: + # Success! Convert to PNG + img = frame.to_image() + buf = io.BytesIO() + img.save(buf, format="PNG") + img_b64 = base64.b64encode(buf.getvalue()).decode("ascii") + break + if img_b64: break + if img_b64: break + + if not img_b64: + raise RuntimeError("Failed to decode a frame from scrcpy stream") + + return img_b64 + + finally: + if writer: + writer.close() + try: + await writer.wait_closed() + except Exception: + pass + + finally: + # 7. Cleanup + if server_process: + try: + server_process.terminate() + await asyncio.wait_for(server_process.wait(), timeout=2.0) + except Exception: + try: + server_process.kill() + await server_process.wait() + except Exception: + pass + + if local_port: + try: + await adb_run("-s", serial, "forward", "--remove", f"tcp:{local_port}", timeout=5.0) + except Exception: + pass diff --git a/adb_vision/test_scrcpy.py b/adb_vision/test_scrcpy.py new file mode 100644 index 0000000000..8cea7c39bd --- /dev/null +++ b/adb_vision/test_scrcpy.py @@ -0,0 +1,130 @@ +import pytest +import asyncio +import base64 +import io +from unittest.mock import MagicMock, patch, AsyncMock, ANY +from adb_vision.screenshot import _capture_scrcpy + +@pytest.mark.asyncio +async def test_capture_scrcpy_happy_path(): + # Mock adb_run + adb_run = AsyncMock() + adb_run.side_effect = [ + b"pushed", # push jar + b"21504", # forward tcp:0 + b"removed" # forward --remove + ] + + # Mock asyncio.create_subprocess_exec + mock_process = AsyncMock() + mock_process.returncode = None + mock_process.terminate = MagicMock() + mock_process.wait = AsyncMock() + + # Mock asyncio.open_connection + mock_reader = AsyncMock() + mock_writer = AsyncMock() + mock_writer.close = MagicMock() + mock_writer.wait_closed = AsyncMock() + + # Header: 69 bytes + mock_reader.readexactly.return_value = b"\x00" + b"x" * 64 + b"\x05\x00\x02\xd0" + + # Mock frame.to_image() to return a real PIL image so img.save works + from PIL import Image + real_img = Image.new("RGB", (1280, 720), color="red") + + with patch("asyncio.create_subprocess_exec", return_value=mock_process), \ + patch("asyncio.open_connection", return_value=(mock_reader, mock_writer)), \ + patch("av.codec.CodecContext") as mock_codec_context_class: + + mock_codec = MagicMock() + mock_codec_context_class.create.return_value = mock_codec + + mock_packet = MagicMock() + mock_codec.parse.return_value = [mock_packet] + + mock_frame = MagicMock() + mock_codec.decode.return_value = [mock_frame] + mock_frame.to_image.return_value = real_img + + mock_reader.read.side_effect = [b"some h264 data", b""] + + serial = "127.0.0.1:5555" + result = await _capture_scrcpy( + adb_run=adb_run, + serial=serial, + adb_exe="adb" + ) + + assert result is not None + assert len(result) > 0 + # Verify it's valid base64 + base64.b64decode(result) + + # Verify ADB calls + assert adb_run.call_count >= 3 + adb_run.assert_any_call("-s", serial, "push", ANY, "/data/local/tmp/scrcpy-server.jar", timeout=15.0) + adb_run.assert_any_call("-s", serial, "forward", "tcp:0", "localabstract:scrcpy", timeout=5.0) + adb_run.assert_any_call("-s", serial, "forward", "--remove", "tcp:21504", timeout=5.0) + +@pytest.mark.asyncio +async def test_capture_scrcpy_server_fail(): + adb_run = AsyncMock() + adb_run.side_effect = [ + b"pushed", # push jar + b"21504", # forward tcp:0 + b"removed" # forward --remove + ] + + mock_process = AsyncMock() + mock_process.returncode = 1 + mock_process.stderr.read.return_value = b"some error" + mock_process.terminate = MagicMock() + mock_process.wait = AsyncMock() + + with patch("asyncio.create_subprocess_exec", return_value=mock_process), \ + patch("asyncio.open_connection", side_effect=ConnectionRefusedError()): + + with pytest.raises(RuntimeError, match="scrcpy-server exited: some error"): + await _capture_scrcpy( + adb_run=adb_run, + serial="127.0.0.1:5555", + adb_exe="adb" + ) + +@pytest.mark.asyncio +async def test_capture_scrcpy_decode_fail(): + adb_run = AsyncMock() + adb_run.side_effect = [ + b"pushed", # push jar + b"21504", # forward tcp:0 + b"removed" # forward --remove + ] + + mock_process = AsyncMock() + mock_process.returncode = None + mock_process.terminate = MagicMock() + mock_process.wait = AsyncMock() + + mock_reader = AsyncMock() + mock_writer = AsyncMock() + mock_writer.close = MagicMock() + mock_writer.wait_closed = AsyncMock() + mock_reader.readexactly.return_value = b"\x00" + b"x" * 64 + b"\x05\x00\x02\xd0" + mock_reader.read.return_value = b"" # No data + + with patch("asyncio.create_subprocess_exec", return_value=mock_process), \ + patch("asyncio.open_connection", return_value=(mock_reader, mock_writer)), \ + patch("av.codec.CodecContext") as mock_codec_context_class: + + mock_codec = MagicMock() + mock_codec_context_class.create.return_value = mock_codec + mock_codec.parse.return_value = [] + + with pytest.raises(RuntimeError, match="Failed to decode a frame from scrcpy stream"): + await _capture_scrcpy( + adb_run=adb_run, + serial="127.0.0.1:5555", + adb_exe="adb" + )