Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 24 additions & 3 deletions phone_agent/adb/screenshot.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,11 @@
import uuid
from dataclasses import dataclass
from io import BytesIO
from typing import Tuple

from PIL import Image

MAX_MODEL_IMAGE_DIMENSION = 2048


@dataclass
class Screenshot:
Expand Down Expand Up @@ -67,6 +68,7 @@ def get_screenshot(device_id: str | None = None, timeout: int = 10) -> Screensho

# Read and encode image
img = Image.open(temp_path)
img = _resize_image_to_model_limit(img)
width, height = img.size

buffered = BytesIO()
Expand All @@ -92,18 +94,37 @@ def _get_adb_prefix(device_id: str | None) -> list:
return ["adb"]


def _resize_image_to_model_limit(
img: Image.Image, max_dimension: int = MAX_MODEL_IMAGE_DIMENSION
) -> Image.Image:
"""Resize an image so neither dimension exceeds the model input limit."""
width, height = img.size
largest_dimension = max(width, height)
if largest_dimension <= max_dimension:
return img

scale = max_dimension / largest_dimension
resized_size = (
max(1, int(width * scale)),
max(1, int(height * scale)),
)
return img.resize(resized_size, Image.Resampling.LANCZOS)


def _create_fallback_screenshot(is_sensitive: bool) -> Screenshot:
"""Create a black fallback image when screenshot fails."""
default_width, default_height = 1080, 2400

black_img = Image.new("RGB", (default_width, default_height), color="black")
black_img = _resize_image_to_model_limit(black_img)
width, height = black_img.size
buffered = BytesIO()
black_img.save(buffered, format="PNG")
base64_data = base64.b64encode(buffered.getvalue()).decode("utf-8")

return Screenshot(
base64_data=base64_data,
width=default_width,
height=default_height,
width=width,
height=height,
is_sensitive=is_sensitive,
)
33 changes: 33 additions & 0 deletions tests/test_screenshot_resize.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
import base64
from io import BytesIO

from PIL import Image

from phone_agent.adb.screenshot import (
_create_fallback_screenshot,
_resize_image_to_model_limit,
)


def _decode_image_size(base64_data: str) -> tuple[int, int]:
image_data = base64.b64decode(base64_data)
image = Image.open(BytesIO(image_data))
return image.size


def test_resize_image_to_model_limit_preserves_aspect_ratio() -> None:
image = Image.new("RGB", (1240, 2772), color="white")

resized = _resize_image_to_model_limit(image)

assert resized.size == (916, 2048)


def test_fallback_screenshot_stays_within_model_limit() -> None:
screenshot = _create_fallback_screenshot(is_sensitive=False)

assert max(screenshot.width, screenshot.height) <= 2048
assert _decode_image_size(screenshot.base64_data) == (
screenshot.width,
screenshot.height,
)