From b7bfc5fecb81de1349af3cf5b83c9011ef1b5af6 Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Tue, 24 Feb 2026 00:13:23 -0800 Subject: [PATCH 01/20] Initial implementation Currently, this shares memory with the .raw bytearray in some cases. We should consider whether we want to change that to always copy, or just wait until we have the full shared buffer support. This implementation always does the array manipulation on the CPU. For PyTorch or TensorFlow, the user is likely to want to do the array manipulation on the GPU for performance. This is very easy to implement. --- demos/tinytv-stream-simple.py | 2 +- demos/tinytv-stream.py | 2 +- demos/video-capture-simple.py | 4 +- demos/video-capture.py | 9 +- docs/source/examples.rst | 8 +- docs/source/examples/opencv_numpy.py | 7 +- docs/source/examples/pil.py | 8 +- docs/source/examples/pil_pixels.py | 21 +-- src/mss/base.py | 3 +- src/mss/screenshot.py | 169 +++++++++++++++++- src/tests/test_setup.py | 4 + src/tests/third_party/test_numpy_method.py | 57 ++++++ src/tests/third_party/test_pil_method.py | 40 +++++ .../third_party/test_tensorflow_method.py | 36 ++++ src/tests/third_party/test_torch_method.py | 32 ++++ 15 files changed, 366 insertions(+), 36 deletions(-) create mode 100644 src/tests/third_party/test_numpy_method.py create mode 100644 src/tests/third_party/test_pil_method.py create mode 100644 src/tests/third_party/test_tensorflow_method.py create mode 100644 src/tests/third_party/test_torch_method.py diff --git a/demos/tinytv-stream-simple.py b/demos/tinytv-stream-simple.py index c0dc1c75..f776415c 100755 --- a/demos/tinytv-stream-simple.py +++ b/demos/tinytv-stream-simple.py @@ -97,7 +97,7 @@ def main() -> None: # The next step is to resize the image to fit the TinyTV's screen. There's a great image # manipulation library called PIL, or Pillow, that can do that. Let's transfer the raw pixels in # the ScreenShot object into a PIL Image. - original_image = Image.frombytes("RGB", screenshot.size, screenshot.bgra, "raw", "BGRX") + original_image = screenshot.to_pil("RGB") # Now, we can resize it. The resize method may stretch the image to make it match the TinyTV's # screen; the advanced demo gives other options. Using a reducing gap is optional, but speeds up diff --git a/demos/tinytv-stream.py b/demos/tinytv-stream.py index a3993892..51bd0a0d 100755 --- a/demos/tinytv-stream.py +++ b/demos/tinytv-stream.py @@ -348,7 +348,7 @@ def capture_image( while True: sct_img = sct.grab(rect) - pil_img = Image.frombytes("RGB", sct_img.size, sct_img.bgra, "raw", "BGRX") + pil_img = sct_img.to_pil("RGB") yield pil_img diff --git a/demos/video-capture-simple.py b/demos/video-capture-simple.py index 1e7f7f94..cb172c48 100755 --- a/demos/video-capture-simple.py +++ b/demos/video-capture-simple.py @@ -72,7 +72,7 @@ def main() -> None: monitor = sct.monitors[1] # Because of how H.264 video stores color information, libx264 requires the video size to be a multiple of - # two. + # two. monitor["width"] = (monitor["width"] // 2) * 2 monitor["height"] = (monitor["height"] // 2) * 2 @@ -129,7 +129,7 @@ def main() -> None: # use PIL: you can create an Image from the screenshot, and create a VideoFrame from that. That said, # if you want to boost the fps rate by about 50%, check out the full demo, and search for # from_numpy_buffer. - img = Image.frombytes("RGB", screenshot.size, screenshot.bgra, "raw", "BGRX") + img = screenshot.to_pil("RGB") frame = av.VideoFrame.from_image(img) # When we encode frames, we get back a list of packets. Often, we'll get no packets at first: the diff --git a/demos/video-capture.py b/demos/video-capture.py index cd7b8ac8..878a3d74 100755 --- a/demos/video-capture.py +++ b/demos/video-capture.py @@ -227,20 +227,19 @@ def video_process( # Many Python objects expose their underlying memory via the "buffer protocol". A buffer is just a view of # raw bytes that other libraries can interpret without copying. # - # Common buffer objects include: `bytes`, `bytearray`, `memoryview`, and `array.array`. `screenshot.bgra` is - # also a buffer (currently it is a `bytes` object, though that detail may change in the future). + # Common buffer objects include: `bytes`, `bytearray`, `memoryview`, and `array.array`. A ScreenShot object + # stores its pixel data in a buffer. # # Minimum-copy path: ScreenShot -> NumPy -> VideoFrame # ---------------------------------------------------- # # `np.frombuffer()` creates an ndarray *view* of an existing buffer (no copy). Reshaping also stays as a - # view. + # view. `ScreenShot.to_numpy()` uses the same approach internally and keeps the zero-copy behavior. # # PyAV's `VideoFrame.from_ndarray()` always copies the data into a new frame-owned buffer. For this demo we # use the undocumented `VideoFrame.from_numpy_buffer()`, which creates a `VideoFrame` that shares memory with # the ndarray. - ndarray = np.frombuffer(screenshot.bgra, dtype=np.uint8) - ndarray = ndarray.reshape(screenshot.height, screenshot.width, 4) + ndarray = screenshot.to_numpy(channels="BGRA") frame = av.VideoFrame.from_numpy_buffer(ndarray, format="bgra") # Set the PTS and time base for the frame. diff --git a/docs/source/examples.rst b/docs/source/examples.rst index ce540141..fa6cba11 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -123,7 +123,7 @@ PIL === You can use the Python Image Library (aka Pillow) to do whatever you want with raw pixels. -This is an example using `frombytes() `_: +This is an example using :py:meth:`mss.screenshot.ScreenShot.to_pil`: .. literalinclude:: examples/pil.py :lines: 7- @@ -133,7 +133,7 @@ This is an example using `frombytes() `_: +This is an example using :py:meth:`mss.screenshot.ScreenShot.to_pil` and direct pixel edits: .. literalinclude:: examples/pil_pixels.py :lines: 7- @@ -147,6 +147,10 @@ See how fast you can record the screen. You can easily view a HD movie with VLC and see it too in the OpenCV window. And with __no__ lag please. +When using :py:meth:`mss.screenshot.ScreenShot.to_numpy`, pass ``channels="BGR"`` +for OpenCV, and use ``channels="RGB"`` (the default) for scikit-image and most +other frameworks. + .. literalinclude:: examples/opencv_numpy.py :lines: 7- diff --git a/docs/source/examples/opencv_numpy.py b/docs/source/examples/opencv_numpy.py index 9275de2b..282f1981 100644 --- a/docs/source/examples/opencv_numpy.py +++ b/docs/source/examples/opencv_numpy.py @@ -7,7 +7,6 @@ import time import cv2 -import numpy as np import mss @@ -18,15 +17,15 @@ while "Screen capturing": last_time = time.time() - # Get raw pixels from the screen, save it to a Numpy array - img = np.array(sct.grab(monitor)) + # Get raw pixels from the screen, save it to a NumPy array + img = sct.grab(monitor).to_numpy(channels="BGR") # Display the picture cv2.imshow("OpenCV/Numpy normal", img) # Display the picture in grayscale # cv2.imshow('OpenCV/Numpy grayscale', - # cv2.cvtColor(img, cv2.COLOR_BGRA2GRAY)) + # cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)) print(f"fps: {1 / (time.time() - last_time)}") diff --git a/docs/source/examples/pil.py b/docs/source/examples/pil.py index 03ff778c..92b79999 100644 --- a/docs/source/examples/pil.py +++ b/docs/source/examples/pil.py @@ -1,11 +1,9 @@ """This is part of the MSS Python's module. Source: https://github.com/BoboTiG/python-mss. -PIL example using frombytes(). +PIL example using ScreenShot.to_pil(). """ -from PIL import Image - import mss with mss.mss() as sct: @@ -15,9 +13,7 @@ sct_img = sct.grab(monitor) # Create the Image - img = Image.frombytes("RGB", sct_img.size, sct_img.bgra, "raw", "BGRX") - # The same, but less efficient: - # img = Image.frombytes('RGB', sct_img.size, sct_img.rgb) + img = sct_img.to_pil("RGB") # And save it! output = f"monitor-{num}.png" diff --git a/docs/source/examples/pil_pixels.py b/docs/source/examples/pil_pixels.py index d1264bc6..b9349449 100644 --- a/docs/source/examples/pil_pixels.py +++ b/docs/source/examples/pil_pixels.py @@ -4,8 +4,6 @@ PIL examples to play with pixels. """ -from PIL import Image - import mss with mss.mss() as sct: @@ -13,19 +11,16 @@ sct_img = sct.grab(sct.monitors[1]) # Create an Image - img = Image.new("RGB", sct_img.size) - - # Best solution: create a list(tuple(R, G, B), ...) for putdata() - pixels = zip(sct_img.raw[2::4], sct_img.raw[1::4], sct_img.raw[::4]) - img.putdata(list(pixels)) + img = sct_img.to_pil("RGB") - # But you can set individual pixels too (slower) - """ + # Set individual pixels (slower) pixels = img.load() - for x in range(sct_img.width): - for y in range(sct_img.height): - pixels[x, y] = sct_img.pixel(x, y) - """ + max_x = min(100, sct_img.width) + max_y = min(100, sct_img.height) + for x in range(max_x): + for y in range(max_y): + r, g, b = pixels[x, y] + pixels[x, y] = (255 - r, 255 - g, 255 - b) # Show it! img.show() diff --git a/src/mss/base.py b/src/mss/base.py index b738b4df..64335ece 100644 --- a/src/mss/base.py +++ b/src/mss/base.py @@ -238,7 +238,8 @@ def primary_monitor(self) -> Monitor: """ monitors = self.monitors if len(monitors) <= 1: # Only the "all monitors" entry or empty - raise ScreenShotError("No monitor found.") + msg = "No monitor found." + raise ScreenShotError(msg) for monitor in monitors[1:]: # Skip the "all monitors" entry at index 0 if monitor.get("is_primary", False): diff --git a/src/mss/screenshot.py b/src/mss/screenshot.py index 9443aba4..118e8079 100644 --- a/src/mss/screenshot.py +++ b/src/mss/screenshot.py @@ -3,7 +3,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast from mss.exception import ScreenShotError from mss.models import Monitor, Pixel, Pixels, Pos, Size @@ -11,6 +11,14 @@ if TYPE_CHECKING: # pragma: nocover from collections.abc import Iterator + import numpy as np + import PIL.Image + import tensorflow as tf + import torch + +Channels: TypeAlias = Literal["BGRA", "BGR", "RGB", "RGBA"] +Layout: TypeAlias = Literal["HWC", "CHW"] + class ScreenShot: """Screenshot object. @@ -127,6 +135,165 @@ def rgb(self) -> bytes: return self.__rgb + def to_pil(self, mode: str = "RGB") -> PIL.Image.Image: + """Convert the screenshot to a Pillow image. + + Args: + mode: The requested image mode. Must be ``"RGB"`` + (default) or ``"RGBA"``. + + Returns: + A :class:`PIL.Image.Image` instance. + + Notes: + When requesting ``"RGBA"``, the alpha channel may not be + meaningful on all platforms/backends. + + Use ``channels="BGR"`` for OpenCV, and ``channels="RGB"`` + (the default) for scikit-image and most other frameworks. + """ + mode = mode.upper() + if mode not in {"RGB", "RGBA"}: + msg = "Mode must be 'RGB' or 'RGBA'" + raise ValueError(msg) + + from PIL import Image # noqa: PLC0415 + + raw_mode = "BGRX" if mode == "RGB" else "BGRA" + return Image.frombuffer(mode, self.size, self.raw, "raw", raw_mode, 0, 1) + + def to_numpy(self, channels: Channels = "RGB", layout: Layout = "HWC") -> np.ndarray: + """Convert the screenshot to a NumPy array. + + Args: + channels: The requested channel order. Must be + ``"BGRA"``, ``"BGR"``, ``"RGB"`` (default), or + ``"RGBA"``. + layout: The requested layout. Must be ``"HWC"`` (default) + or ``"CHW"``. + + Returns: + A NumPy array of dtype ``uint8``. + + Notes: + When requesting ``"RGBA"``, the alpha channel may not be + meaningful on all platforms/backends. + """ + channels = cast("Channels", channels.upper()) + layout = cast("Layout", layout.upper()) + + if channels not in {"BGRA", "BGR", "RGB", "RGBA"}: + msg = "Channels must be 'BGRA', 'BGR', 'RGB', or 'RGBA'" + raise ValueError(msg) + if layout not in {"HWC", "CHW"}: + msg = "Layout must be 'HWC' or 'CHW'" + raise ValueError(msg) + + import numpy as np # noqa: PLC0415 + + frame = np.frombuffer(self.raw, dtype=np.uint8).reshape((self.height, self.width, 4)) + if channels == "BGRA": + data = frame + elif channels == "BGR": + data = frame[:, :, :3] + elif channels == "RGB": + data = frame[:, :, [2, 1, 0]] + else: # RGBA + data = frame[:, :, [2, 1, 0, 3]] + + if layout == "CHW": + data = np.transpose(data, (2, 0, 1)) + + return data + + def to_torch( + self, + channels: Channels = "RGB", + layout: Layout = "CHW", + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + """Convert the screenshot to a PyTorch tensor. + + Args: + channels: The requested channel order. Must be + ``"BGRA"``, ``"BGR"``, ``"RGB"`` (default), or + ``"RGBA"``. + layout: The requested layout. Must be ``"CHW"`` (default) + or ``"HWC"``. + dtype: The requested dtype as a ``torch.dtype``. Defaults to + ``torch.float32``. + + Returns: + A PyTorch tensor. + + Notes: + The default layout is ``"CHW"`` because it is more + commonly used in PyTorch models. This is different than + in :py:meth:`to_numpy` or :py:meth:`to_tensorflow`, which + default to ``"HWC"``. + + When requesting ``"RGBA"``, the alpha channel may not be + meaningful on all platforms/backends. + + Floating point dtypes are scaled to the ``[0, 1]`` range. + """ + import torch # noqa: PLC0415 + + frame = self.to_numpy(channels=channels, layout=layout) + + if dtype is None: + dtype = torch.float32 + elif not isinstance(dtype, torch.dtype): + msg = "Dtype must be a torch dtype" + raise ValueError(msg) + + tensor = torch.from_numpy(frame) + tensor = tensor.to(dtype=dtype) + if dtype.is_floating_point: + tensor = tensor / 255.0 + return tensor + + def to_tensorflow( + self, + channels: Channels = "RGB", + layout: Layout = "HWC", + dtype: tf.DType | str = "float32", + ) -> tf.Tensor: + """Convert the screenshot to a TensorFlow tensor. + + Args: + channels: The requested channel order. Must be + ``"BGRA"``, ``"BGR"``, ``"RGB"`` (default), or + ``"RGBA"``. + layout: The requested layout. Must be ``"HWC"`` (default) + or ``"CHW"``. + dtype: The requested dtype. Can be a string like + ``"float32"`` (default) or a ``tf.DType``. + + Returns: + A TensorFlow tensor. + + Notes: + When requesting ``"RGBA"``, the alpha channel may not be + meaningful on all platforms/backends. + + Floating point dtypes are scaled to the ``[0, 1]`` range. + """ + import tensorflow as tf # noqa: PLC0415 + + frame = self.to_numpy(channels=channels, layout=layout) + + try: + tf_dtype = tf.as_dtype(dtype) + except (TypeError, ValueError) as exc: + msg = "Dtype must be a TensorFlow DType or valid string" + raise ValueError(msg) from exc + + tensor = tf.convert_to_tensor(frame, dtype=tf_dtype) + if tf_dtype.is_floating: + tensor = tensor / tf.constant(255.0, dtype=tf_dtype) + return tensor + @property def top(self) -> int: """Convenient accessor to the top position.""" diff --git a/src/tests/test_setup.py b/src/tests/test_setup.py index 29be9b62..1e4b400e 100644 --- a/src/tests/test_setup.py +++ b/src/tests/test_setup.py @@ -105,7 +105,11 @@ def test_sdist() -> None: f"mss-{__version__}/src/tests/test_xcb.py", f"mss-{__version__}/src/tests/third_party/__init__.py", f"mss-{__version__}/src/tests/third_party/test_numpy.py", + f"mss-{__version__}/src/tests/third_party/test_numpy_method.py", f"mss-{__version__}/src/tests/third_party/test_pil.py", + f"mss-{__version__}/src/tests/third_party/test_pil_method.py", + f"mss-{__version__}/src/tests/third_party/test_tensorflow_method.py", + f"mss-{__version__}/src/tests/third_party/test_torch_method.py", f"mss-{__version__}/src/xcbproto/README.md", f"mss-{__version__}/src/xcbproto/gen_xcb_to_py.py", f"mss-{__version__}/src/xcbproto/randr.xml", diff --git a/src/tests/third_party/test_numpy_method.py b/src/tests/third_party/test_numpy_method.py new file mode 100644 index 00000000..620a9aeb --- /dev/null +++ b/src/tests/third_party/test_numpy_method.py @@ -0,0 +1,57 @@ +"""This is part of the MSS Python's module. +Source: https://github.com/BoboTiG/python-mss. +""" + +from __future__ import annotations + +import pytest + +from mss.screenshot import ScreenShot + +np = pytest.importorskip("numpy") + + +def test_to_numpy_default_rgb_hwc() -> None: + raw = bytearray([10, 20, 30, 40]) + shot = ScreenShot.from_size(raw, 1, 1) + + arr = shot.to_numpy() + assert arr.shape == (1, 1, 3) + assert arr.dtype == np.uint8 + assert np.array_equal(arr[0, 0], np.array([30, 20, 10], dtype=np.uint8)) + + +def test_to_numpy_bgra_chw() -> None: + raw = bytearray([1, 2, 3, 4]) + shot = ScreenShot.from_size(raw, 1, 1) + + arr = shot.to_numpy(channels="BGRA", layout="CHW") + assert arr.shape == (4, 1, 1) + assert arr.dtype == np.uint8 + assert np.array_equal(arr[:, 0, 0], np.array([1, 2, 3, 4], dtype=np.uint8)) + + +def test_to_numpy_rgba_hwc() -> None: + raw = bytearray([5, 6, 7, 8]) + shot = ScreenShot.from_size(raw, 1, 1) + + arr = shot.to_numpy(channels="RGBA") + assert arr.shape == (1, 1, 4) + assert arr.dtype == np.uint8 + assert np.array_equal(arr[0, 0], np.array([7, 6, 5, 8], dtype=np.uint8)) + + +def test_to_numpy_bad_channels() -> None: + raw = bytearray([0, 0, 0, 0]) + shot = ScreenShot.from_size(raw, 1, 1) + + with pytest.raises(ValueError, match="Channels must be 'BGRA', 'BGR', 'RGB', or 'RGBA'"): + shot.to_numpy(channels="gray") # type: ignore[arg-type] + + +def test_to_numpy_bad_layout() -> None: + raw = bytearray([0, 0, 0, 0]) + shot = ScreenShot.from_size(raw, 1, 1) + + with pytest.raises(ValueError, match="Layout must be 'HWC' or 'CHW'"): + shot.to_numpy(layout="NHWC") # type: ignore[arg-type] diff --git a/src/tests/third_party/test_pil_method.py b/src/tests/third_party/test_pil_method.py new file mode 100644 index 00000000..90220f70 --- /dev/null +++ b/src/tests/third_party/test_pil_method.py @@ -0,0 +1,40 @@ +"""This is part of the MSS Python's module. +Source: https://github.com/BoboTiG/python-mss. +""" + +from __future__ import annotations + +import pytest + +from mss.screenshot import ScreenShot + +pytest.importorskip("PIL.Image") + + +def test_to_pil_rgb_default() -> None: + # The raw format is BGRA/BGRX (B, G, R, X) + raw = bytearray([1, 2, 3, 4]) + shot = ScreenShot.from_size(raw, 1, 1) + + img = shot.to_pil() + assert img.mode == "RGB" + assert img.size == shot.size + assert img.getpixel((0, 0)) == (3, 2, 1) + + +def test_to_pil_rgba() -> None: + raw = bytearray([5, 6, 7, 8]) + shot = ScreenShot.from_size(raw, 1, 1) + + img = shot.to_pil("rgba") + assert img.mode == "RGBA" + assert img.size == shot.size + assert img.getpixel((0, 0)) == (7, 6, 5, 8) + + +def test_to_pil_bad_mode() -> None: + raw = bytearray([0, 0, 0, 0]) + shot = ScreenShot.from_size(raw, 1, 1) + + with pytest.raises(ValueError, match="Mode must be 'RGB' or 'RGBA'"): + shot.to_pil("L") diff --git a/src/tests/third_party/test_tensorflow_method.py b/src/tests/third_party/test_tensorflow_method.py new file mode 100644 index 00000000..f8dee855 --- /dev/null +++ b/src/tests/third_party/test_tensorflow_method.py @@ -0,0 +1,36 @@ +"""This is part of the MSS Python's module. +Source: https://github.com/BoboTiG/python-mss. +""" + +from __future__ import annotations + +import pytest + +from mss.screenshot import ScreenShot + +np = pytest.importorskip("numpy") +tf = pytest.importorskip("tensorflow") + + +def test_to_tensorflow_default() -> None: + raw = bytearray([1, 2, 3, 4]) + shot = ScreenShot.from_size(raw, 1, 1) + + tensor = shot.to_tensorflow() + assert tuple(tensor.shape) == (1, 1, 3) + assert tensor.dtype == tf.float32 + np.testing.assert_allclose( + tensor.numpy()[0, 0], + np.array([3.0 / 255.0, 2.0 / 255.0, 1.0 / 255.0], dtype=np.float32), + rtol=1e-6, + atol=1e-7, + ) + + +def test_to_tensorflow_dtype_string() -> None: + raw = bytearray([9, 8, 7, 6]) + shot = ScreenShot.from_size(raw, 1, 1) + + tensor = shot.to_tensorflow(dtype="uint8") + assert tensor.dtype == tf.uint8 + assert (tensor.numpy()[0, 0] == [7, 8, 9]).all() diff --git a/src/tests/third_party/test_torch_method.py b/src/tests/third_party/test_torch_method.py new file mode 100644 index 00000000..89fd0cb3 --- /dev/null +++ b/src/tests/third_party/test_torch_method.py @@ -0,0 +1,32 @@ +"""This is part of the MSS Python's module. +Source: https://github.com/BoboTiG/python-mss. +""" + +from __future__ import annotations + +import pytest + +from mss.screenshot import ScreenShot + +np = pytest.importorskip("numpy") +torch = pytest.importorskip("torch") + + +def test_to_torch_default() -> None: + raw = bytearray([1, 2, 3, 4]) + shot = ScreenShot.from_size(raw, 1, 1) + + tensor = shot.to_torch() + assert tuple(tensor.shape) == (3, 1, 1) + assert tensor.dtype == torch.float32 + expected = torch.tensor([3 / 255.0, 2 / 255.0, 1 / 255.0], dtype=torch.float32) + assert torch.allclose(tensor[:, 0, 0], expected) + + +def test_to_torch_dtype_uint8() -> None: + raw = bytearray([5, 6, 7, 8]) + shot = ScreenShot.from_size(raw, 1, 1) + + tensor = shot.to_torch(dtype=torch.uint8) + assert tensor.dtype == torch.uint8 + assert torch.equal(tensor[:, 0, 0], torch.tensor([7, 6, 5], dtype=torch.uint8)) From 8acdf969ee45ec23af560f104eb60380114df9a2 Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Wed, 25 Feb 2026 21:57:48 -0800 Subject: [PATCH 02/20] Fix Python 3.9 compatibility Python 3.9 doesn't define TypeAlias, and it's not required in this case, so we can just remove it. --- src/mss/screenshot.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mss/screenshot.py b/src/mss/screenshot.py index 118e8079..fab8c153 100644 --- a/src/mss/screenshot.py +++ b/src/mss/screenshot.py @@ -3,7 +3,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast +from typing import TYPE_CHECKING, Any, Literal, cast from mss.exception import ScreenShotError from mss.models import Monitor, Pixel, Pixels, Pos, Size @@ -16,8 +16,8 @@ import tensorflow as tf import torch -Channels: TypeAlias = Literal["BGRA", "BGR", "RGB", "RGBA"] -Layout: TypeAlias = Literal["HWC", "CHW"] +Channels = Literal["BGRA", "BGR", "RGB", "RGBA"] +Layout = Literal["HWC", "CHW"] class ScreenShot: From 70f89ed8d4bf1faf158e0520dfb634d752a4a6d3 Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Mon, 1 Jun 2026 21:21:05 -0700 Subject: [PATCH 03/20] Update to 11.0 conventions --- docs/source/examples.rst | 11 +++++------ src/mss/screenshot.py | 7 +++---- src/tests/third_party/test_numpy_method.py | 2 +- src/tests/third_party/test_pil_method.py | 2 +- src/tests/third_party/test_tensorflow_method.py | 2 +- src/tests/third_party/test_torch_method.py | 2 +- 6 files changed, 12 insertions(+), 14 deletions(-) diff --git a/docs/source/examples.rst b/docs/source/examples.rst index 03efa352..7f2c27ba 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -122,8 +122,8 @@ falling back to XGetImage: PIL === -You can use the Python Image Library (aka Pillow) to do whatever you want with raw pixels. -This is an example using :py:meth:`mss.screenshot.ScreenShot.to_pil`: +You can use the Python Image Library (aka Pillow) to do whatever you want with raw pixels. This is an example using +:py:meth:`mss.ScreenShot.to_pil`: .. literalinclude:: examples/pil.py :lines: 7- @@ -133,7 +133,7 @@ This is an example using :py:meth:`mss.screenshot.ScreenShot.to_pil`: Playing with pixels ------------------- -This is an example using :py:meth:`mss.screenshot.ScreenShot.to_pil` and direct pixel edits: +This is an example using :py:meth:`mss.ScreenShot.to_pil` and direct pixel edits: .. literalinclude:: examples/pil_pixels.py :lines: 7- @@ -147,9 +147,8 @@ See how fast you can record the screen. You can easily view a HD movie with VLC and see it too in the OpenCV window. And with __no__ lag please. -When using :py:meth:`mss.screenshot.ScreenShot.to_numpy`, pass ``channels="BGR"`` -for OpenCV, and use ``channels="RGB"`` (the default) for scikit-image and most -other frameworks. +When using :py:meth:`mss.ScreenShot.to_numpy`, pass ``channels="BGR"`` for OpenCV, and use ``channels="RGB"`` (the +default) for scikit-image and most other frameworks. .. literalinclude:: examples/opencv_numpy.py :lines: 7- diff --git a/src/mss/screenshot.py b/src/mss/screenshot.py index 1d8f7652..d4f34a8c 100644 --- a/src/mss/screenshot.py +++ b/src/mss/screenshot.py @@ -12,12 +12,11 @@ from collections.abc import Iterator from typing import Any - from typing_extensions import Buffer - import numpy as np import PIL.Image import tensorflow as tf import torch + from typing_extensions import Buffer Channels = Literal["BGRA", "BGR", "RGB", "RGBA"] Layout = Literal["HWC", "CHW"] @@ -191,7 +190,7 @@ def to_pil(self, mode: str = "RGB") -> PIL.Image.Image: from PIL import Image # noqa: PLC0415 raw_mode = "BGRX" if mode == "RGB" else "BGRA" - return Image.frombuffer(mode, self.size, self.raw, "raw", raw_mode, 0, 1) + return Image.frombuffer(mode, self.size, self._raw, "raw", raw_mode, 0, 1) def to_numpy(self, channels: Channels = "RGB", layout: Layout = "HWC") -> np.ndarray: """Convert the screenshot to a NumPy array. @@ -222,7 +221,7 @@ def to_numpy(self, channels: Channels = "RGB", layout: Layout = "HWC") -> np.nda import numpy as np # noqa: PLC0415 - frame = np.frombuffer(self.raw, dtype=np.uint8).reshape((self.height, self.width, 4)) + frame = np.frombuffer(self._raw, dtype=np.uint8).reshape((self.height, self.width, 4)) if channels == "BGRA": data = frame elif channels == "BGR": diff --git a/src/tests/third_party/test_numpy_method.py b/src/tests/third_party/test_numpy_method.py index 620a9aeb..c8e67405 100644 --- a/src/tests/third_party/test_numpy_method.py +++ b/src/tests/third_party/test_numpy_method.py @@ -6,7 +6,7 @@ import pytest -from mss.screenshot import ScreenShot +from mss import ScreenShot np = pytest.importorskip("numpy") diff --git a/src/tests/third_party/test_pil_method.py b/src/tests/third_party/test_pil_method.py index 90220f70..04042dd3 100644 --- a/src/tests/third_party/test_pil_method.py +++ b/src/tests/third_party/test_pil_method.py @@ -6,7 +6,7 @@ import pytest -from mss.screenshot import ScreenShot +from mss import ScreenShot pytest.importorskip("PIL.Image") diff --git a/src/tests/third_party/test_tensorflow_method.py b/src/tests/third_party/test_tensorflow_method.py index f8dee855..6b2036d7 100644 --- a/src/tests/third_party/test_tensorflow_method.py +++ b/src/tests/third_party/test_tensorflow_method.py @@ -6,7 +6,7 @@ import pytest -from mss.screenshot import ScreenShot +from mss import ScreenShot np = pytest.importorskip("numpy") tf = pytest.importorskip("tensorflow") diff --git a/src/tests/third_party/test_torch_method.py b/src/tests/third_party/test_torch_method.py index 89fd0cb3..7f7bba65 100644 --- a/src/tests/third_party/test_torch_method.py +++ b/src/tests/third_party/test_torch_method.py @@ -6,7 +6,7 @@ import pytest -from mss.screenshot import ScreenShot +from mss import ScreenShot np = pytest.importorskip("numpy") torch = pytest.importorskip("torch") From cd2ce5800b8bb10077a75d12a99a5ec5eb42739e Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Wed, 24 Jun 2026 16:35:05 -0700 Subject: [PATCH 04/20] Fix an inadvertent change The reverted change was from a temporary test, but I accidentally included it in the commit. It's not relevant to the issue at hand. --- docs/source/examples/pil_pixels.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/docs/source/examples/pil_pixels.py b/docs/source/examples/pil_pixels.py index e82b7b3f..57f50952 100644 --- a/docs/source/examples/pil_pixels.py +++ b/docs/source/examples/pil_pixels.py @@ -20,12 +20,9 @@ # But you can set individual pixels too (slower) """ pixels = img.load() - max_x = min(100, sct_img.width) - max_y = min(100, sct_img.height) - for x in range(max_x): - for y in range(max_y): - r, g, b = pixels[x, y] - pixels[x, y] = (255 - r, 255 - g, 255 - b) + for x in range(sct_img.width): + for y in range(sct_img.height): + pixels[x, y] = sct_img.pixel(x, y) """ # Show it! From 37238bef2bd28eebd252e2ce49556c486471560c Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Wed, 24 Jun 2026 17:42:56 -0700 Subject: [PATCH 05/20] Docstring cleanups, and widen TensorFlow dtype support a bit The docstrings weren't all in Sphinx format, and could use a little clarification in places. TensorFlow itself allows NumPy dtypes and DataType enums, so expand to_tensorflow to accept those too. --- src/mss/screenshot.py | 126 ++++++++++++++++++++---------------------- 1 file changed, 59 insertions(+), 67 deletions(-) diff --git a/src/mss/screenshot.py b/src/mss/screenshot.py index d4f34a8c..983e32a5 100644 --- a/src/mss/screenshot.py +++ b/src/mss/screenshot.py @@ -67,7 +67,7 @@ def __array_interface__(self) -> dict[str, Any]: ``torch.from_numpy``), TensorFlow (via ``tf.convert_to_tensor``), JAX (via ``jax.numpy.asarray``), Pandas, scikit-learn, Matplotlib, some OpenCV functions, and others. This allows you to pass a - :class:`ScreenShot` instance directly to these libraries without + :py:class:`ScreenShot` instance directly to these libraries without needing to convert it first. This is in HWC order, with 4 channels (BGRA). @@ -102,9 +102,8 @@ def bgra(self) -> memoryview: when given a read-only buffer, but will still work. However, actually modifying the data may cause undefined behavior. - .. note:: - While the name is ``bgra``, the alpha channel may or may - not be valid. + While the name is ``bgra``, the alpha channel may not represent + meaningful transparency on all platforms/backends. """ # Making a read-only copy of a memoryview is very cheap. But # we still always return the same memoryview: somebody using a @@ -150,7 +149,7 @@ def rgb(self) -> memoryview: when given a read-only buffer, but will still work. However, actually modifying the data may cause undefined behavior. - :: note:: + .. note:: This is a computed property. If possible, using the :py:attr:`bgra` property directly is usually more efficient. @@ -168,19 +167,14 @@ def rgb(self) -> memoryview: def to_pil(self, mode: str = "RGB") -> PIL.Image.Image: """Convert the screenshot to a Pillow image. - Args: - mode: The requested image mode. Must be ``"RGB"`` - (default) or ``"RGBA"``. - - Returns: - A :class:`PIL.Image.Image` instance. + :param mode: The requested image mode. Must be ``"RGB"`` + (default) or ``"RGBA"``. - Notes: - When requesting ``"RGBA"``, the alpha channel may not be - meaningful on all platforms/backends. + When requesting ``"RGBA"``, the alpha channel may not represent + meaningful transparency on all platforms/backends. - Use ``channels="BGR"`` for OpenCV, and ``channels="RGB"`` - (the default) for scikit-image and most other frameworks. + TODO(jholveck): After #536 is resolved, add a note about the + buffer sharing semantics. """ mode = mode.upper() if mode not in {"RGB", "RGBA"}: @@ -195,19 +189,20 @@ def to_pil(self, mode: str = "RGB") -> PIL.Image.Image: def to_numpy(self, channels: Channels = "RGB", layout: Layout = "HWC") -> np.ndarray: """Convert the screenshot to a NumPy array. - Args: - channels: The requested channel order. Must be - ``"BGRA"``, ``"BGR"``, ``"RGB"`` (default), or - ``"RGBA"``. - layout: The requested layout. Must be ``"HWC"`` (default) - or ``"CHW"``. + :param channels: The requested channel order. Must be + ``"BGRA"``, ``"BGR"``, ``"RGB"`` (default), or ``"RGBA"``. + :param layout: The requested layout. Must be ``"HWC"`` + (default) or ``"CHW"``. + :returns: A NumPy array of dtype ``uint8``. - Returns: - A NumPy array of dtype ``uint8``. + Use ``channels="BGR"`` for OpenCV, and ``channels="RGB"`` (the + default) for scikit-image and most other frameworks. - Notes: - When requesting ``"RGBA"``, the alpha channel may not be - meaningful on all platforms/backends. + When requesting ``"RGBA"`` or ``"BGRA"``, the alpha channel may + not represent meaningful transparency on all platforms/backends. + + TODO(jholveck): After #536 is resolved, add a note about the + buffer sharing semantics. """ channels = cast("Channels", channels.upper()) layout = cast("Layout", layout.upper()) @@ -244,28 +239,25 @@ def to_torch( ) -> torch.Tensor: """Convert the screenshot to a PyTorch tensor. - Args: - channels: The requested channel order. Must be - ``"BGRA"``, ``"BGR"``, ``"RGB"`` (default), or - ``"RGBA"``. - layout: The requested layout. Must be ``"CHW"`` (default) - or ``"HWC"``. - dtype: The requested dtype as a ``torch.dtype``. Defaults to - ``torch.float32``. + :param channels: The requested channel order. Must be + ``"BGRA"``, ``"BGR"``, ``"RGB"`` (default), or ``"RGBA"``. + :param layout: The requested layout. Must be ``"CHW"`` + (default) or ``"HWC"``. + :param dtype: The requested dtype as a ``torch.dtype``. + Defaults to ``torch.float32``. - Returns: - A PyTorch tensor. + Floating point dtypes are scaled to the ``[0, 1]`` range. - Notes: - The default layout is ``"CHW"`` because it is more - commonly used in PyTorch models. This is different than - in :py:meth:`to_numpy` or :py:meth:`to_tensorflow`, which - default to ``"HWC"``. + The default layout is ``"CHW"`` because it is more commonly used + in PyTorch models. This is different than in + :py:meth:`to_numpy` or :py:meth:`to_tensorflow`, which default + to ``"HWC"``. - When requesting ``"RGBA"``, the alpha channel may not be - meaningful on all platforms/backends. + When requesting ``"RGBA"`` or ``"BGRA"``, the alpha channel may + not represent meaningful transparency on all platforms/backends. - Floating point dtypes are scaled to the ``[0, 1]`` range. + TODO(jholveck): After #536 is resolved, add a note about the + buffer sharing semantics. """ import torch # noqa: PLC0415 @@ -274,8 +266,8 @@ def to_torch( if dtype is None: dtype = torch.float32 elif not isinstance(dtype, torch.dtype): - msg = "Dtype must be a torch dtype" - raise ValueError(msg) + msg = 'argument "dtype" must be a torch.dtype' + raise TypeError(msg) tensor = torch.from_numpy(frame) tensor = tensor.to(dtype=dtype) @@ -287,40 +279,40 @@ def to_tensorflow( self, channels: Channels = "RGB", layout: Layout = "HWC", - dtype: tf.DType | str = "float32", + dtype: tf.DType | np.dtype | int | str = "float32", ) -> tf.Tensor: """Convert the screenshot to a TensorFlow tensor. - Args: - channels: The requested channel order. Must be - ``"BGRA"``, ``"BGR"``, ``"RGB"`` (default), or - ``"RGBA"``. - layout: The requested layout. Must be ``"HWC"`` (default) - or ``"CHW"``. - dtype: The requested dtype. Can be a string like - ``"float32"`` (default) or a ``tf.DType``. + :param channels: The requested channel order. Must be + ``"BGRA"``, ``"BGR"``, ``"RGB"`` (default), or ``"RGBA"``. + :param layout: The requested layout. Must be ``"HWC"`` + (default) or ``"CHW"``. + :param dtype: The requested dtype. Can be a string like + ``"float32"`` (default), a :py:class:``tf.DType``, an int + representing a TensorFlow ``DataClass`` enum value, or a + ``np.dtype``. - Returns: - A TensorFlow tensor. + Floating point dtypes are scaled to the ``[0, 1]`` range. - Notes: - When requesting ``"RGBA"``, the alpha channel may not be - meaningful on all platforms/backends. + When requesting ``"RGBA"`` or ``"BGRA"``, the alpha channel may + not represent meaningful transparency on all platforms/backends. - Floating point dtypes are scaled to the ``[0, 1]`` range. + Currently, the returned :py:class:`tf.Tensor` does not share + memory with the :py:class:`ScreenShot`. This is expected to + change in the future. TODO(jholveck): After #536 is resolved, + add a note about the expected buffer sharing semantics. """ import tensorflow as tf # noqa: PLC0415 frame = self.to_numpy(channels=channels, layout=layout) - try: - tf_dtype = tf.as_dtype(dtype) - except (TypeError, ValueError) as exc: - msg = "Dtype must be a TensorFlow DType or valid string" - raise ValueError(msg) from exc + # TypeErrors from tf.as_dtype are passed up to the caller. + tf_dtype = tf.as_dtype(dtype) tensor = tf.convert_to_tensor(frame, dtype=tf_dtype) if tf_dtype.is_floating: + # TensorFlow's implicit dtype conversion rules are not trivial. We use an explicit dtype on both sides + # instead, by making a tf.constant. tensor = tensor / tf.constant(255.0, dtype=tf_dtype) return tensor From 1aa94301f73e8c13870b1055b7b40ef58ff852c5 Mon Sep 17 00:00:00 2001 From: Joel Holveck Date: Fri, 26 Jun 2026 03:32:42 +0000 Subject: [PATCH 06/20] Doc improvements The main point of this commit is to be explicit about the memory-sharing semantics: mostly, that we don't guarantee memory sharing one way or another. More to the point, that if they modify the pixel data in one of the exported objects, it may or may not affect the other exported objects. Memory sharing isn't expected to affect most users, since they are unlikely to access the pixel data with two different exported objects, but it's best to be explicit. To have a good place to put the memory-sharing discussion, I added a new section to the usage.rst file, listing the ways to access the pixel data. This is also where I documented the NumPy array interface protocol, which previously didn't show up in our Sphinx docs. (It might have been in the docs for 10.1, which didn't use autodoc.) I also put a discussion of the alpha channel in this new section: one thing I keep forgetting is to delete the alpha channel when using Matplotlib, and wondering why it's blank. TODO: We probably should document `grab` in the usage.rst file for the new section to make sense. In the process, I've also added intersphinx links to libraries like PIL and NumPy, so that our Sphinx docs will link to them. While checking that, I fixed a few other references (like referring to bytes with :py:type:, when it's indexed under :py:class:). --- docs/source/conf.py | 13 +++- docs/source/release-history/v11.0.0.md | 9 +-- docs/source/usage.rst | 98 +++++++++++++++++++++++++- src/mss/base.py | 32 +++++---- src/mss/screenshot.py | 39 ++++------ 5 files changed, 140 insertions(+), 51 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 11f3627c..3ef2b422 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -92,5 +92,14 @@ # ---------------------------------------------- -# Example configuration for intersphinx: refer to the Python standard library. -intersphinx_mapping = {"python": ("https://docs.python.org/3", None)} +intersphinx_mapping = { + "numpy": ("https://numpy.org/doc/stable/", None), + "pil": ("https://pillow.readthedocs.io/en/stable/", None), + "python": ("https://docs.python.org/3", None), + # TensorFlow doesn't have an official InterSphinx mapping, but there is a community one. + "tensorflow": ( + "https://www.tensorflow.org/api_docs/python", + "https://github.com/GPflow/tensorflow-intersphinx/raw/master/tf2_py_objects.inv", + ), + "torch": ("https://docs.pytorch.org/docs/stable/", None), +} diff --git a/docs/source/release-history/v11.0.0.md b/docs/source/release-history/v11.0.0.md index 6601a08b..055173b4 100644 --- a/docs/source/release-history/v11.0.0.md +++ b/docs/source/release-history/v11.0.0.md @@ -18,7 +18,7 @@ writing, but should be by the time we release 11.0!) The {py:attr}`mss.ScreenShot.raw` attribute has been removed. Use the {py:attr}`mss.ScreenShot.bgra` property instead. The {py:attr}`mss.ScreenShot.bgra` and {py:attr}`mss.ScreenShot.rgb` properties now will return read-only bytes-like -{py:type}`memoryview` objects, not necessarily {py:type}`bytes` or {py:type}`bytearray` objects. For practical use +{py:class}`memoryview` objects, not necessarily {py:class}`bytes` or {py:class}`bytearray` objects. For practical use cases, this should not be noticible. This change was allows faster access to screenshot data, with fewer memory copies. ### Python 3.9 EOL @@ -49,10 +49,3 @@ Support for additional operating systems is planned. ### General Improvements The MSS context object will now always surface inner exceptions, even if `__exit__` may also generate an exception during tear-down. - -### Documentation - -The documentation has received numerous small improvements. A few highlights: - -* The Pillow examples and demos using {py:meth}`PIL.Image.frombuffer` now explicitly specify the decoder arguments, - as recommended by Pillow. (#535) diff --git a/docs/source/usage.rst b/docs/source/usage.rst index 08acdf8c..a963243b 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -32,6 +32,103 @@ still available in 11.0, but are also deprecated:: # Microsoft Windows from mss.windows import MSS +Accessing Pixel Data +==================== + +.. attention:: + TODO(jholveck): We should have something documenting ``grab`` before this. + +Once you have the :py:class:`mss.ScreenShot` object, you'll want to use the pixel data. You can get this from the +:py:class:`mss.ScreenShot` object directly, using any of several methods: + +* :py:attr:`mss.ScreenShot.bgra`: (fastest) Direct access to the pixel data, as a :py:class:`memoryview` of + ``BGRABGRA...`` bytes. +* :py:attr:`mss.ScreenShot.rgb`: A :py:class:`memoryview` of ``RGBRGB...`` bytes. +* :py:attr:`mss.ScreenShot.pixels`: A 2d array (list of lists) of ``(R, G, B)`` tuples. +* :py:meth:`mss.ScreenShot.pixel`: Access ``(R, G, B)`` tuples of a particular x, y coordinate. + +Often, though, you'll export screenshot data to a different framework. You can often do this by passing the +:py:attr:`mss.ScreenShot.bgra` data to the framework's appropriate function. MSS also provides easy-to-use methods to +work with many popular frameworks: + +* :py:meth:`mss.ScreenShot.to_pil`: Creates a :py:class:`PIL.Image.Image` for use with the popular Python Imaging + Library, `Pillow `_. +* :py:meth:`mss.ScreenShot.to_numpy`: Creates a :py:class:`numpy.ndarray` for use with the high-speed NumPy scientific + computing library. This is compatible with most other Python frameworks that have image manipulation capabilities, + such as `scikit-image `_ and `OpenCV `_. +* :py:meth:`mss.ScreenShot.to_torch`: Creates a :py:class:`torch.Tensor` for use with + `PyTorch `_, a popular deep learning framework. +* :py:meth:`mss.ScreenShot.to_tensorflow`: Creates a :py:class:`tf.Tensor` for use with + `TensorFlow `_, another popular deep learning framework. + +NumPy Array Interface Protocol +------------------------------ + +Many libraries support the `NumPy array interface protocol +`_. This allows them to accept a +:py:class:`mss.ScreenShot` object directly to these libraries, without needing to convert it to a NumPy array first. +Some examples include the following libraries: + +* Many `SciPy `_ projects +* `CuPy `_, a GPU-accelerated NumPy-like library +* `JAX `_, a high-performance machine learning library +* `Pandas `_, a popular data analysis library +* `scikit-learn `_, a popular machine learning library +* `Matplotlib `_, a popular plotting library +* Some functions from `OpenCV `_, a popular computer vision library + +When using the NumPy array interface protocol, the returned object is in HWC (height, width, channels) format, with the +channels in BGRA order. + +Note that OpenCV uses RGB order, rather than the BGRA order used in this automatic conversion. You may prefer to +use the :py:meth:`mss.ScreenShot.to_numpy` method instead. + +Alpha Channel +------------- + +The alpha channel is used for transparency in images. However, it's also sometimes just used as a placeholder for an +unused channel. In the case of screenshots, the alpha channel is often not used for transparency, and may be filled +with zeros. If an image processing library interprets the alpha channel as transparency, this can make it think the +image is transparent. + +For instance, if you use Matplotlib to display a screenshot, you might see nothing at all. This happens if the OS has +filled the alpha channel with zeros (which is common on many platforms). + +The methods described above, such as :py:meth:`mss.ScreenShot.to_numpy`, can convert the pixel data to BGR (or RGB) +format, removing the alpha channel entirely. + +In other words, instead of ``plt.imshow(img)``, you can use ``plt.imshow(img.to_numpy(channels="RGB"))`` to display the +screenshot correctly. + +In the future, MSS may provide an indicator of whether the alpha channel is meaningful or not, as well as whether it is +premultiplied or straight alpha. For now, you should assume that the alpha channel is not meaningful, and either ignore +or remove it, unless you know that it's meaningful for your specific circumstances. + +Memory Sharing +-------------- + +There's a subtlety to be aware of in the following conditions: + +1. You are using any of the above methods (or properties, or the NumPy array interface protocol) to convert a + :py:class:`mss.ScreenShot` object to another format, *and* +2. You use two different methods, or the same method twice, on the same :py:class:`mss.ScreenShot` object, *and* +3. You modify the pixel data of the returned object (e.g., a NumPy array or a PIL image). + +When using any of the above methods, the returned object might (but does not always) share pixel memory with the +original :py:class:`mss.ScreenShot` object. This means that if you modify the returned object's pixels, it may also +modify the original :py:class:`mss.ScreenShot` object, or other objects that share the same memory. + +For instance, if you use :py:meth:`mss.ScreenShot.to_numpy` to create a NumPy array, then use +:py:meth:`mss.ScreenShot.to_pil` to create a PIL image, both objects may share the same memory. If you modify the +pixels of the NumPy array, it may also modify the pixels of the PIL image, and vice versa. + +Pixel memory is never guaranteed to be shared; it depends on many specifics. Whether memory is shared or not is an +implementation detail, and not part of the semantic versioning guarantees of MSS: it may change in future versions, +or even when a program is run in different environments. + +If you want to ensure that memory is not shared, you can make a copy of the returned object. For instance, if you +want to ensure that a NumPy array does not share memory with the original :py:class:`mss.ScreenShot` object, you can +use the :py:meth:`numpy.ndarray.copy` method to create a copy of the array. Intensive Use ============= @@ -141,7 +238,6 @@ There are three available backends. The legacy backend powered by :c:func:`XGetImage`. It is kept solely for systems where XCB libraries are unavailable and no new features are being added to it. - Command Line ============ diff --git a/src/mss/base.py b/src/mss/base.py index 28529c44..09c0a775 100644 --- a/src/mss/base.py +++ b/src/mss/base.py @@ -165,15 +165,15 @@ def _choose_impl(**kwargs: Any) -> MSSImplementation: class MSS: """Multiple ScreenShots class - :param backend: Backend selector, for platforms with multiple backends. + :param backend: Backend selector, for platforms with multiple + backends. :param compression_level: PNG compression level. - :param with_cursor: Include the mouse cursor in screenshots - (GNU/Linux only) - :type display: bool, optional (default False) - :param display: X11 display name (GNU/Linux only). - :type display: bytes | str, optional (default :envvar:`$DISPLAY`) - :param max_displays: Maximum number of displays to enumerate (macOS only). - :type max_displays: int, optional (default 32) + :param with_cursor: Include the mouse cursor in screenshots. + Optional, default False. (GNU/Linux only) + :param display: X11 display name. Optional; default + :envvar:`$DISPLAY`. (GNU/Linux only) + :param max_displays: Maximum number of displays to enumerate. + Optional, default 32. (macOS only). .. versionadded:: 8.0.0 ``compression_level``, ``display``, ``max_displays``, and @@ -295,7 +295,7 @@ def grab(self, monitor: Monitor | tuple[int, int, int, int], /) -> ScreenShot: """Retrieve screen pixels for a given monitor. Note: ``monitor`` can be a tuple like the one - :py:meth:`PIL.ImageGrab.grab` accepts: ``(left, top, right, bottom)`` + :py:func:`PIL.ImageGrab.grab` accepts: ``(left, top, right, bottom)`` :param monitor: The coordinates and size of the box to capture. See :meth:`monitors ` for object details. @@ -391,12 +391,14 @@ def save( ) -> Iterator[str]: """Grab a screenshot and save it to a file. - :param int mon: The monitor to screenshot (default=0). ``-1`` grabs all - monitors, ``0`` grabs each monitor, and ``N`` grabs monitor ``N``. - :param str output: The output filename. Keywords: ``{mon}``, ``{top}``, - ``{left}``, ``{width}``, ``{height}``, ``{date}``. - :param callable callback: Called before saving the screenshot; receives - the ``output`` argument. + :param int mon: The monitor to screenshot (default=0). ``-1`` + grabs all monitors, ``0`` grabs each monitor, and ``N`` + grabs monitor ``N``. + :param str output: The output filename. Keywords: ``{mon}``, + ``{top}``, ``{left}``, ``{width}``, ``{height}``, + ``{date}``. + :param typing.Callable callback: Called before saving the + screenshot; receives the ``output`` argument. :return: Created file(s). """ monitors = self.monitors diff --git a/src/mss/screenshot.py b/src/mss/screenshot.py index 983e32a5..df0a6e4a 100644 --- a/src/mss/screenshot.py +++ b/src/mss/screenshot.py @@ -12,12 +12,14 @@ from collections.abc import Iterator from typing import Any - import numpy as np + # We don't import numpy as np, since their InterSphinx reference isn't set up with that shortcut. + import numpy # noqa: ICN001 import PIL.Image import tensorflow as tf import torch from typing_extensions import Buffer +# Type checkers can see these, but they don't get into the Sphinx docs. I'm not sure if we should do this differently. Channels = Literal["BGRA", "BGR", "RGB", "RGBA"] Layout = Literal["HWC", "CHW"] @@ -98,10 +100,6 @@ def bgra(self) -> memoryview: BGRxBGRx... sequence. A specific pixel can be accessed as ``bgra[(y * width + x) * 4:(y * width + x) * 4 + 4].`` - The memoryview is read-only. PyTorch will issue a warning - when given a read-only buffer, but will still work. However, - actually modifying the data may cause undefined behavior. - While the name is ``bgra``, the alpha channel may not represent meaningful transparency on all platforms/backends. """ @@ -145,10 +143,6 @@ def rgb(self) -> memoryview: RGBRGB... sequence. A specific pixel can be accessed as ``rgb[(y * width + x) * 4:(y * width + x) * 4 + 4].`` - The memoryview is read-only. PyTorch will issue a warning - when given a read-only buffer, but will still work. However, - actually modifying the data may cause undefined behavior. - .. note:: This is a computed property. If possible, using the :py:attr:`bgra` property directly is usually more @@ -173,8 +167,7 @@ def to_pil(self, mode: str = "RGB") -> PIL.Image.Image: When requesting ``"RGBA"``, the alpha channel may not represent meaningful transparency on all platforms/backends. - TODO(jholveck): After #536 is resolved, add a note about the - buffer sharing semantics. + .. version-added:: 11.0.0 """ mode = mode.upper() if mode not in {"RGB", "RGBA"}: @@ -186,7 +179,7 @@ def to_pil(self, mode: str = "RGB") -> PIL.Image.Image: raw_mode = "BGRX" if mode == "RGB" else "BGRA" return Image.frombuffer(mode, self.size, self._raw, "raw", raw_mode, 0, 1) - def to_numpy(self, channels: Channels = "RGB", layout: Layout = "HWC") -> np.ndarray: + def to_numpy(self, channels: Channels = "RGB", layout: Layout = "HWC") -> numpy.ndarray: """Convert the screenshot to a NumPy array. :param channels: The requested channel order. Must be @@ -201,8 +194,7 @@ def to_numpy(self, channels: Channels = "RGB", layout: Layout = "HWC") -> np.nda When requesting ``"RGBA"`` or ``"BGRA"``, the alpha channel may not represent meaningful transparency on all platforms/backends. - TODO(jholveck): After #536 is resolved, add a note about the - buffer sharing semantics. + .. version-added:: 11.0.0 """ channels = cast("Channels", channels.upper()) layout = cast("Layout", layout.upper()) @@ -243,7 +235,7 @@ def to_torch( ``"BGRA"``, ``"BGR"``, ``"RGB"`` (default), or ``"RGBA"``. :param layout: The requested layout. Must be ``"CHW"`` (default) or ``"HWC"``. - :param dtype: The requested dtype as a ``torch.dtype``. + :param dtype: The requested dtype as a :py:class:`torch.dtype`. Defaults to ``torch.float32``. Floating point dtypes are scaled to the ``[0, 1]`` range. @@ -256,8 +248,7 @@ def to_torch( When requesting ``"RGBA"`` or ``"BGRA"``, the alpha channel may not represent meaningful transparency on all platforms/backends. - TODO(jholveck): After #536 is resolved, add a note about the - buffer sharing semantics. + .. version-added:: 11.0.0 """ import torch # noqa: PLC0415 @@ -279,7 +270,7 @@ def to_tensorflow( self, channels: Channels = "RGB", layout: Layout = "HWC", - dtype: tf.DType | np.dtype | int | str = "float32", + dtype: tf.dtypes.DType | numpy.dtype | int | str = "float32", ) -> tf.Tensor: """Convert the screenshot to a TensorFlow tensor. @@ -288,19 +279,17 @@ def to_tensorflow( :param layout: The requested layout. Must be ``"HWC"`` (default) or ``"CHW"``. :param dtype: The requested dtype. Can be a string like - ``"float32"`` (default), a :py:class:``tf.DType``, an int - representing a TensorFlow ``DataClass`` enum value, or a - ``np.dtype``. + ``"float32"`` (default), a :py:class:`tf.DType + `, an int representing a TensorFlow + ``DataClass`` enum value, or a :py:class:`np.dtype + `. Floating point dtypes are scaled to the ``[0, 1]`` range. When requesting ``"RGBA"`` or ``"BGRA"``, the alpha channel may not represent meaningful transparency on all platforms/backends. - Currently, the returned :py:class:`tf.Tensor` does not share - memory with the :py:class:`ScreenShot`. This is expected to - change in the future. TODO(jholveck): After #536 is resolved, - add a note about the expected buffer sharing semantics. + .. version-added:: 11.0.0 """ import tensorflow as tf # noqa: PLC0415 From 9afa3a24ec71bf09d56d5e8cb6b66d3948e637d5 Mon Sep 17 00:00:00 2001 From: Joel Holveck Date: Fri, 26 Jun 2026 06:21:01 +0000 Subject: [PATCH 07/20] Add some introductory documentation about grab --- docs/source/examples.rst | 2 ++ docs/source/usage.rst | 38 ++++++++++++++++++++++++++++++++------ 2 files changed, 34 insertions(+), 6 deletions(-) diff --git a/docs/source/examples.rst b/docs/source/examples.rst index 7f2c27ba..cf77b7a4 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -1,3 +1,5 @@ +.. _examples: + ======== Examples ======== diff --git a/docs/source/usage.rst b/docs/source/usage.rst index a963243b..a84170e2 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -32,14 +32,39 @@ still available in 11.0, but are also deprecated:: # Microsoft Windows from mss.windows import MSS +Capturing Screenshots +===================== + +If you simply need to capture one or more monitors to PNG files, the :ref:`examples` section has code ready for you to +copy and paste. + +If instead you want to use the pixel data yourself, you can do so easily, with the :py:meth:`mss.MSS.grab` method. + +You'll first need to decide whether you want to capture all the monitors, a single monitor, or a specific region of the +screen. The :py:class:`mss.MSS` object has a :py:attr:`mss.MSS.monitors` attribute that is a list of all the monitors, +starting from index 1, as well as the full virtual screen (all monitors combined) at index 0. + +Once you've decided what you want to capture, you can call :py:meth:`mss.MSS.grab` with the appropriate monitor or +region. This will return a :py:class:`mss.ScreenShot` object, which contains the pixel data and other information about +the screenshot. + +For instance, you can capture the first monitor and get a :py:class:`mss.ScreenShot` object like this:: + + with MSS() as sct: + sct_img = sct.grab(sct.monitors[1]) # Capture the first monitor + +Ok, now you've got the :py:class:`mss.ScreenShot` object. But what do +you do with it? + Accessing Pixel Data ==================== -.. attention:: - TODO(jholveck): We should have something documenting ``grab`` before this. +Once you have the :py:class:`mss.ScreenShot` object, you'll want to use the pixel data. There are several ways, +depending on what you want to do with it. This is a quick overview of the options, with more details in the linked +references. -Once you have the :py:class:`mss.ScreenShot` object, you'll want to use the pixel data. You can get this from the -:py:class:`mss.ScreenShot` object directly, using any of several methods: +If you want to examine individual pixels, you can get them from the :py:class:`mss.ScreenShot` object directly, using +any of several methods: * :py:attr:`mss.ScreenShot.bgra`: (fastest) Direct access to the pixel data, as a :py:class:`memoryview` of ``BGRABGRA...`` bytes. @@ -52,7 +77,8 @@ Often, though, you'll export screenshot data to a different framework. You can work with many popular frameworks: * :py:meth:`mss.ScreenShot.to_pil`: Creates a :py:class:`PIL.Image.Image` for use with the popular Python Imaging - Library, `Pillow `_. + Library, `Pillow `_. This provides a wide range of image manipulation capabilities, + including saving to many different formats. * :py:meth:`mss.ScreenShot.to_numpy`: Creates a :py:class:`numpy.ndarray` for use with the high-speed NumPy scientific computing library. This is compatible with most other Python frameworks that have image manipulation capabilities, such as `scikit-image `_ and `OpenCV `_. @@ -116,7 +142,7 @@ There's a subtlety to be aware of in the following conditions: When using any of the above methods, the returned object might (but does not always) share pixel memory with the original :py:class:`mss.ScreenShot` object. This means that if you modify the returned object's pixels, it may also -modify the original :py:class:`mss.ScreenShot` object, or other objects that share the same memory. +modify the pixels stored in the original :py:class:`mss.ScreenShot` object, or other objects that share the same memory. For instance, if you use :py:meth:`mss.ScreenShot.to_numpy` to create a NumPy array, then use :py:meth:`mss.ScreenShot.to_pil` to create a PIL image, both objects may share the same memory. If you modify the From 3a693f5f3d57d6ccbbca7b37c6eff9dada561475 Mon Sep 17 00:00:00 2001 From: Joel Holveck Date: Fri, 26 Jun 2026 08:41:13 +0000 Subject: [PATCH 08/20] Small improvements Doc fixes. Speed up converting to a NumPy array in RGB order. That sort of thing. --- src/mss/screenshot.py | 42 +++++++++++++++++++++++++++--------------- 1 file changed, 27 insertions(+), 15 deletions(-) diff --git a/src/mss/screenshot.py b/src/mss/screenshot.py index 143c7637..4788a5aa 100644 --- a/src/mss/screenshot.py +++ b/src/mss/screenshot.py @@ -45,11 +45,8 @@ def __init__(self, data: Buffer, monitor: Monitor, /, *, size: Size | None = Non #: NamedTuple of the screenshot size. self.size: Size = Size(monitor["width"], monitor["height"]) if size is None else size - # Buffer of the raw BGRA pixels, retrieved by the platform-specific implementations. This is kept read-write if - # it was originally so, in order for _merge to work, and so it can be used with ctypes. However, it should not - # be modified once __pixels or __rgb have been accessed, so that the cached values for __pixels and __rgb aren't - # potentially inconsistent if the user changes data. - self._raw: memoryview = memoryview(data) + # Buffer of the raw BGRA pixels, retrieved by the platform-specific implementations. + self._raw: memoryview[int] = memoryview(data) assert self._raw.nbytes == self.size.width * self.size.height * 4, ( # noqa: S101 "Data size does not match screenshot dimensions." ) @@ -70,10 +67,7 @@ def __array_interface__(self) -> dict[str, Any]: :py:class:`ScreenShot` instance directly to these libraries without needing to convert it first. - This is in HWC order, with 4 channels (BGRA). - - The array is read-write, for maximum compatibility. However, - actually modifying the data may cause undefined behavior. + This is in HWC order, with 4 channels in BGRA order. .. seealso:: @@ -94,7 +88,7 @@ def from_size(cls: type[ScreenShot], data: Buffer, width: int, height: int, /) - return cls(data, monitor) @property - def bgra(self) -> memoryview: + def bgra(self) -> memoryview[int]: """BGRx values from the BGRx raw pixels. The format is a memoryview object of bytes. These are in a @@ -116,7 +110,7 @@ def bgra(self) -> memoryview: return self._raw @property - def raw(self) -> memoryview: + def raw(self) -> memoryview[int]: """Deprecated alias for :py:attr:`bgra`. .. version-deprecated:: 10.2.0 @@ -231,6 +225,11 @@ def to_numpy(self, channels: Channels = "RGB", layout: Layout = "HWC") -> numpy. .. version-added:: 11.0.0 """ + # This is used as the channel and layout manipulation function for the other frameworks, too. This is on the + # assumption that NumPy is probably going to be the fastest on CPU, although I haven't tested this. It's also a + # prerequisite for the others, so we'd need to have a NumPy-only implementation anyway: the NumPy implementation + # shouldn't depend on PyTorch / TensorFlow, while the converse isn't true. + channels = cast("Channels", channels.upper()) layout = cast("Layout", layout.upper()) @@ -249,11 +248,17 @@ def to_numpy(self, channels: Channels = "RGB", layout: Layout = "HWC") -> numpy. elif channels == "BGR": data = frame[:, :, :3] elif channels == "RGB": - data = frame[:, :, [2, 1, 0]] + # Using [2,1,0] instead of 2::-1 would copy, rather than creating a view. + data = frame[:, :, 2::-1] else: # RGBA + # This can't be represented as a view, since the channels within a pixel are not ordered in a way that can + # be represented with a constant offset. In other words, the way that NumPy strides work, you can only make + # a view if you can express the desired element order in a x:y:z style range relative to the base array's + # order. data = frame[:, :, [2, 1, 0, 3]] if layout == "CHW": + # This will always create a view. (We're reordering the axes, not the elements.) data = np.transpose(data, (2, 0, 1)) return data @@ -271,7 +276,9 @@ def to_torch( :param layout: The requested layout. Must be ``"CHW"`` (default) or ``"HWC"``. :param dtype: The requested dtype as a :py:class:`torch.dtype`. - Defaults to ``torch.float32``. + Defaults to the current PyTorch default dtype, which is + usually ``torch.float32``; see + :py:func:`torch.get_default_dtype`. Floating point dtypes are scaled to the ``[0, 1]`` range. @@ -290,15 +297,19 @@ def to_torch( frame = self.to_numpy(channels=channels, layout=layout) if dtype is None: - dtype = torch.float32 + dtype = torch.get_default_dtype() elif not isinstance(dtype, torch.dtype): msg = 'argument "dtype" must be a torch.dtype' raise TypeError(msg) + # PyTorch doesn't support negative strides like NumPy does; we have to copy in that case. This happens if the + # user requested a RGB layout. (And no, we can't use [:,:,2::-1] in PyTorch either.) + if any(s < 0 for s in frame.strides): + frame = frame.copy() tensor = torch.from_numpy(frame) tensor = tensor.to(dtype=dtype) if dtype.is_floating_point: - tensor = tensor / 255.0 + tensor.div_(255.0) return tensor def to_tensorflow( @@ -333,6 +344,7 @@ def to_tensorflow( # TypeErrors from tf.as_dtype are passed up to the caller. tf_dtype = tf.as_dtype(dtype) + # TODO(jholveck): Make sure that we can use convert_to_tensor if the source has negative strides. tensor = tf.convert_to_tensor(frame, dtype=tf_dtype) if tf_dtype.is_floating: # TensorFlow's implicit dtype conversion rules are not trivial. We use an explicit dtype on both sides From e2c3e6700b0ee847eb13399fb3b3bba163c3f357 Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Sat, 27 Jun 2026 02:32:35 -0700 Subject: [PATCH 09/20] Update cat-detector to use to_torch While I was writing this, I realized something else. We may want to change the to_torch and to_tensorflow methods to accept device and stream arguments, and do the channel and layout shuffling on the GPU instead of the CPU. This wouldn't require us to natively support CUDA or anything like that; we'd just be using PyTorch / TensorFlow's existing CUDA support, which is easy. This demo still (as before) shows the GPU-focused way to do the channel and layout shuffling, presented as an alternative to using to_torch. A quick test shows about 107 ms / frame (to_torch) vs. about 62 ms / frame (GPU-focused routine), and that's including the CNN itself too. So there's a very significant benefit there. (That's not quite a fair comparison, since the GPU-focused routine is exactly what this program needs, but not so it'd be a big difference.) --- demos/cat-detector.py | 24 +++--------------------- 1 file changed, 3 insertions(+), 21 deletions(-) diff --git a/demos/cat-detector.py b/demos/cat-detector.py index cde7d833..40dc0e0b 100755 --- a/demos/cat-detector.py +++ b/demos/cat-detector.py @@ -278,27 +278,9 @@ def main() -> None: # Grab the screenshot. sct_img = sct.grab(monitor) - # We transfer the image from MSS to PyTorch via a Pillow Image. Faster approaches exist (see - # screenshot_to_tensor), but PIL is more readable. The bulk of the time in this program is spent doing - # the AI work, so we just use the most convenient mechanism. - img = Image.frombuffer("RGB", sct_img.size, sct_img.bgra, "raw", "BGRX", 0, 1) - - # We explicitly convert it to a tensor here, even though Torchvision can also convert it in the preprocess - # step. This is so that we send it to the GPU before we do the preprocessing: PIL Images are always on - # the CPU, and doing the preprocessing on the GPU is much faster. - # - # Most image APIs, including MSS, use an array layout of [height, width, channels]. In MSS, the - # ScreenShot.bgra data follows this convention, even though it's exposed as a flat bytes object. - # - # In contrast, most AI frameworks expect images in [channels, height, width] order. The pil_to_tensor - # helper performs this rearrangement for us. - img_tensor = torchvision.transforms.v2.functional.pil_to_tensor(img).to(device) - - # An alternative to using PIL is shown in screenshot_to_tensor. In one test, this saves about 20 ms per - # frame if using a GPU, and about 200 ms if using a CPU. This would replace the "img=" and "img_tensor=" - # lines above. - # - #img_tensor = screenshot_to_tensor(sct_img, device) + # Transfer the image from MSS to PyTorch. This does the channel shuffling and reordering on the CPU. For + # better performance, doing the shuffling on the GPU can be more efficient; see screenshot_to_tensor. + img_tensor = sct_img.to_torch().to(device) # Do the preprocessing stages that the trained model expects; see the comment where we define preprocess. # The traditional name for inputs to a neural net is "x", because AI programmers aren't terribly From 5bbbcc07917f8171c85dafdbfff8a2a2b6363baf Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Sat, 27 Jun 2026 02:54:06 -0700 Subject: [PATCH 10/20] Remove a TODO that I did --- src/mss/screenshot.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mss/screenshot.py b/src/mss/screenshot.py index 4788a5aa..b9f0da66 100644 --- a/src/mss/screenshot.py +++ b/src/mss/screenshot.py @@ -344,7 +344,7 @@ def to_tensorflow( # TypeErrors from tf.as_dtype are passed up to the caller. tf_dtype = tf.as_dtype(dtype) - # TODO(jholveck): Make sure that we can use convert_to_tensor if the source has negative strides. + # TensorFlow will always copy in convert_to_tensor. tensor = tf.convert_to_tensor(frame, dtype=tf_dtype) if tf_dtype.is_floating: # TensorFlow's implicit dtype conversion rules are not trivial. We use an explicit dtype on both sides From 1c84624fcc20cdae43f1d53c5c1fa2e24c74fdeb Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Sat, 27 Jun 2026 17:59:07 -0700 Subject: [PATCH 11/20] Add more testing for the third-party array exports. While improving performance earlier in this PR, I had introduced a case where to_torch would fail only for the RGB order. It also seemed likely that some cases might fail only if using particular combinations of channel order and layout, so we just test them all. --- .github/workflows/tests.yml | 26 ++++++++- src/tests/test_setup.py | 1 + src/tests/third_party/conftest.py | 57 +++++++++++++++++++ src/tests/third_party/test_numpy_method.py | 10 ++++ .../third_party/test_tensorflow_method.py | 15 +++++ src/tests/third_party/test_torch_method.py | 18 ++++++ 6 files changed, 125 insertions(+), 2 deletions(-) create mode 100644 src/tests/third_party/conftest.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 3baeb668..a10bce61 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -57,10 +57,13 @@ jobs: matrix: os: - emoji: 🐧 + name: linux runs-on: [ubuntu-latest] - emoji: 🍎 + name: macos runs-on: [macos-latest] - emoji: 🪟 + name: windows runs-on: [windows-latest] python: - name: CPython 3.10 @@ -84,11 +87,30 @@ jobs: run: | python -m pip install -U pip python -m pip install -e '.[dev,tests]' + + # We don't include PyTorch or TensorFlow in the dependencies. This is partly because they are large and not + # needed for most users. But it's mostly because the right way to install them depends a lot on the user's + # needs (such as GPU support) and OS. + + # The PyPI versions of PyTorch for Linux include CUDA for Linux, so we use an index with a CPU-only version there. + # For the other OSs, the PyPI versions are CPU-only. + - name: Install PyTorch (Linux CPU) + if: matrix.os.name == 'linux' + run: python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + - name: Install PyTorch (macOS/Windows) + if: matrix.os.name != 'linux' + run: python -m pip install torch + # TensorFlow stable (as of this June 2026) only supports up to Python 3.12; see + # https://www.tensorflow.org/install/pip#software_requirements + - name: Install TensorFlow (Python 3.10-3.12) + if: contains(fromJSON('["3.10", "3.11", "3.12"]'), matrix.python.runs-on) + run: python -m pip install tensorflow + - name: Tests (GNU/Linux) - if: matrix.os.emoji == '🐧' + if: matrix.os.name == 'linux' run: xvfb-run python -m pytest - name: Tests (macOS, Windows) - if: matrix.os.emoji != '🐧' + if: matrix.os.name != 'linux' run: python -m pytest automerge: diff --git a/src/tests/test_setup.py b/src/tests/test_setup.py index d65fc604..4160f210 100644 --- a/src/tests/test_setup.py +++ b/src/tests/test_setup.py @@ -113,6 +113,7 @@ def test_sdist() -> None: f"mss-{__version__}/src/tests/test_windows.py", f"mss-{__version__}/src/tests/test_xcb.py", f"mss-{__version__}/src/tests/third_party/__init__.py", + f"mss-{__version__}/src/tests/third_party/conftest.py", f"mss-{__version__}/src/tests/third_party/test_numpy.py", f"mss-{__version__}/src/tests/third_party/test_numpy_method.py", f"mss-{__version__}/src/tests/third_party/test_pil.py", diff --git a/src/tests/third_party/conftest.py b/src/tests/third_party/conftest.py new file mode 100644 index 00000000..4661f06a --- /dev/null +++ b/src/tests/third_party/conftest.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +import mss + +if TYPE_CHECKING: + # We have a separate import for just type hints, since we don't want to import numpy in the code itself (just + # importorskip), but type checkers don't always know about importorskip. + import numpy as np_typehints # noqa: ICN001 + +np = pytest.importorskip("numpy") + +TEST_SIZE = (320, 240) + + +def reordered_test_image(channels: str, layout: str) -> np_typehints.ndarray: + """Create a test image with particular channels and layout. + + This allows us to test all the paths of channels and layouts, to + make sure they all work. Restrictions on things like striding + require special cases for some paths, so we test all the paths. + """ + y = np.arange(TEST_SIZE[1], dtype=np.uint32)[:, None] + x = np.arange(TEST_SIZE[0], dtype=np.uint32)[None, :] + + rv = np.zeros((TEST_SIZE[1], TEST_SIZE[0], 0), dtype=np.uint32) + + for ch in channels: + if ch == "R": + charr = (31 * y + 1 * x + 17) & 0xFF + elif ch == "G": + charr = (7 * y + 17 * x + 53) & 0xFF + elif ch == "B": + charr = (13 * y + 29 * x + 101) & 0xFF + elif ch == "A": + charr = (19 * y + 11 * x + 149) & 0xFF + else: + msg = f'Unexpected channel "{ch}"' + raise ValueError(msg) + rv = np.dstack((rv, charr)) + rv = rv.astype(np.uint8) + + source_axes = ["HWC".index(ax) for ax in layout] + destination_axes = [0, 1, 2] + return np.moveaxis(rv, source_axes, destination_axes) + + +@pytest.fixture +def framework_test_image() -> mss.ScreenShot: + ndarray = reordered_test_image("BGRA", "HWC") + # We need the packed buffer to be R/W, hence the extra bytearray copy. + packed = bytearray(ndarray.tobytes()) + width, height = ndarray.shape[1], ndarray.shape[0] + return mss.ScreenShot(packed, {"left": 0, "top": 0, "width": width, "height": height}) diff --git a/src/tests/third_party/test_numpy_method.py b/src/tests/third_party/test_numpy_method.py index c8e67405..cb35a460 100644 --- a/src/tests/third_party/test_numpy_method.py +++ b/src/tests/third_party/test_numpy_method.py @@ -7,6 +7,7 @@ import pytest from mss import ScreenShot +from tests.third_party.conftest import reordered_test_image np = pytest.importorskip("numpy") @@ -55,3 +56,12 @@ def test_to_numpy_bad_layout() -> None: with pytest.raises(ValueError, match="Layout must be 'HWC' or 'CHW'"): shot.to_numpy(layout="NHWC") # type: ignore[arg-type] + + +@pytest.mark.parametrize("layout", ["HWC", "CHW"]) +@pytest.mark.parametrize("channels", ["BGRA", "BGR", "RGBA", "RGB"]) +def test_to_numpy_permutations(framework_test_image: ScreenShot, channels: str, layout: str) -> None: + """Test all permutations of channels and layouts.""" + result = framework_test_image.to_numpy(channels=channels, layout=layout) # type: ignore[arg-type] + target = reordered_test_image(channels=channels, layout=layout) + assert np.array_equal(result, target) diff --git a/src/tests/third_party/test_tensorflow_method.py b/src/tests/third_party/test_tensorflow_method.py index 6b2036d7..602d1184 100644 --- a/src/tests/third_party/test_tensorflow_method.py +++ b/src/tests/third_party/test_tensorflow_method.py @@ -7,6 +7,7 @@ import pytest from mss import ScreenShot +from tests.third_party.conftest import reordered_test_image np = pytest.importorskip("numpy") tf = pytest.importorskip("tensorflow") @@ -34,3 +35,17 @@ def test_to_tensorflow_dtype_string() -> None: tensor = shot.to_tensorflow(dtype="uint8") assert tensor.dtype == tf.uint8 assert (tensor.numpy()[0, 0] == [7, 8, 9]).all() + + +@pytest.mark.parametrize("layout", ["HWC", "CHW"]) +@pytest.mark.parametrize("channels", ["BGRA", "BGR", "RGBA", "RGB"]) +def test_to_tensorflow_permutations(framework_test_image: ScreenShot, channels: str, layout: str) -> None: + """Test all permutations of channels and layouts.""" + uint8_target = reordered_test_image(channels=channels, layout=layout) + bfloat16_target = uint8_target.astype(np.float32) / 255.0 + + uint8_result = framework_test_image.to_tensorflow(channels=channels, layout=layout, dtype="uint8") # type: ignore[arg-type] + assert np.array_equal(uint8_result.numpy(), uint8_target) + + bfloat16_result = framework_test_image.to_tensorflow(channels=channels, layout=layout, dtype="bfloat16") # type: ignore[arg-type] + assert np.allclose(bfloat16_result.numpy(), bfloat16_target, rtol=0, atol=1 / 512.0) diff --git a/src/tests/third_party/test_torch_method.py b/src/tests/third_party/test_torch_method.py index 7f7bba65..5846c68e 100644 --- a/src/tests/third_party/test_torch_method.py +++ b/src/tests/third_party/test_torch_method.py @@ -7,6 +7,7 @@ import pytest from mss import ScreenShot +from tests.third_party.conftest import reordered_test_image np = pytest.importorskip("numpy") torch = pytest.importorskip("torch") @@ -30,3 +31,20 @@ def test_to_torch_dtype_uint8() -> None: tensor = shot.to_torch(dtype=torch.uint8) assert tensor.dtype == torch.uint8 assert torch.equal(tensor[:, 0, 0], torch.tensor([7, 6, 5], dtype=torch.uint8)) + + +@pytest.mark.parametrize("layout", ["HWC", "CHW"]) +@pytest.mark.parametrize("channels", ["BGRA", "BGR", "RGBA", "RGB"]) +def test_to_torch_permutations(framework_test_image: ScreenShot, channels: str, layout: str) -> None: + """Test all permutations of channels and layouts.""" + uint8_target = reordered_test_image(channels=channels, layout=layout) + bfloat16_target = uint8_target.astype(np.float32) / 255.0 + + uint8_result = framework_test_image.to_torch(channels=channels, layout=layout, dtype=torch.uint8) # type: ignore[arg-type] + assert np.array_equal(uint8_result.numpy(), uint8_target) + + bfloat16_result = framework_test_image.to_torch(channels=channels, layout=layout, dtype=torch.bfloat16) # type: ignore[arg-type] + # We have to explicitly cast back to float32 for the comparison, because PyTorch won't directly convert bfloat16 to + # NumPy. + bfloat16_result = bfloat16_result.to(torch.float32) + assert np.allclose(bfloat16_result.numpy(), bfloat16_target, rtol=0, atol=1 / 512.0) From 106b20b23e55fc8605a5681db9408d2f2700f8a5 Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Sat, 27 Jun 2026 21:12:33 -0700 Subject: [PATCH 12/20] Add better support for non-CPU devices in to_torch and to_tensorflow. This adds a device= flag to to_torch, so the user can specify the destination device. (TensorFlow has more automatic device management, and we don't need to do anything explicit.) We also do all the channel shuffling to the GPU. For cases when channel shuffling or dtype conversion requires a copy, this is immensely faster than the alternative. I ran a rough test with the cat detector on my 4k monitor (for everything, including the CNN inference). On a GPU, there's a significant performance boost. However, this is now slower when run on the CPU. * Old code, with GPU: 99 ms/frame * New code, with GPU: 56 ms/frame * Old code, no GPU: 3164 ms/frame * New code, no GPU: 3219 ms/frame * New code, running to_torch(device="cpu") and and then transferring to the GPU for inference: 117 ms/frame It seems the new to_torch is slower when running on the CPU. I suspect (but this is a loose guess) that NumPy, which always did the channel shuffles in the old code, was faster at that. It might also be specific to the extra copy that I think this change introduces in the RGB CHW conversion path in this, due to PyTorch limitations. I don't think this is a significant problem: PyTorch and TensorFlow users are almost all going to be using GPUs. I will still look at ways to improve the CPU performance on PyTorch, though. --- demos/cat-detector.py | 25 +-- src/mss/screenshot.py | 174 +++++++++++++++------ src/tests/third_party/test_numpy_method.py | 4 +- src/tests/third_party/test_torch_method.py | 26 ++- 4 files changed, 150 insertions(+), 79 deletions(-) diff --git a/demos/cat-detector.py b/demos/cat-detector.py index 40dc0e0b..094664f0 100755 --- a/demos/cat-detector.py +++ b/demos/cat-detector.py @@ -128,26 +128,6 @@ MIN_AREA_FRAC = 0.001 -# This function is here for illustrative purposes: the demo doesn't currently call it, but there's a commented-out -# line in the main loop that shows how you might use it. -def screenshot_to_tensor(sct_img: mss.ScreenShot, device: str | torch.device) -> torch.Tensor: - """Convert an MSS ScreenShot to a CHW PyTorch tensor.""" - - # Get a 1d tensor of BGRA values. PyTorch will issue a warning at this step: the ScreenShot's bgra object is - # read-only, but PyTorch doesn't support read-only tensors. However, this is harmless in our case: we'll end up - # copying the data anyway. - img = torch.frombuffer(sct_img.bgra, dtype=torch.uint8) - # Bring everything to the desired device. This is still just a linear buffer of BGRA bytes. - img = img.to(device) - # The next two steps will all just create views of the original tensor, without copying the data. - img = img.view(sct_img.height, sct_img.width, 4) # Interpret as BGRA HWC - img = img.permute(2, 0, 1) # Permute the axes: BGRA CHW - # This final step will create a copy. Copying the data is required to reorder the channels. This also has the - # advantage of also making the tensor contiguous, for more efficient access. - img = img[[2, 1, 0], ...] # Reorder the channels: RGB CHW - return img - - def top_unique_labels(labels: torch.Tensor, scores: torch.Tensor) -> torch.Tensor: """Return the unique labels, ordered by descending score. @@ -278,9 +258,8 @@ def main() -> None: # Grab the screenshot. sct_img = sct.grab(monitor) - # Transfer the image from MSS to PyTorch. This does the channel shuffling and reordering on the CPU. For - # better performance, doing the shuffling on the GPU can be more efficient; see screenshot_to_tensor. - img_tensor = sct_img.to_torch().to(device) + # Transfer the image from MSS to PyTorch. + img_tensor = sct_img.to_torch(device=device) # Do the preprocessing stages that the trained model expects; see the comment where we define preprocess. # The traditional name for inputs to a neural net is "x", because AI programmers aren't terribly diff --git a/src/mss/screenshot.py b/src/mss/screenshot.py index b9f0da66..3510d847 100644 --- a/src/mss/screenshot.py +++ b/src/mss/screenshot.py @@ -208,14 +208,19 @@ def to_pil(self, mode: str = "RGB") -> PIL.Image.Image: raw_mode = "BGRX" if mode == "RGB" else "BGRA" return Image.frombuffer(mode, self.size, self._raw, "raw", raw_mode, 0, 1) - def to_numpy(self, channels: Channels = "RGB", layout: Layout = "HWC") -> numpy.ndarray: + def to_numpy( + self, channels: Channels = "RGB", layout: Layout = "HWC", dtype: numpy.dtype | type | None = None + ) -> numpy.ndarray: """Convert the screenshot to a NumPy array. :param channels: The requested channel order. Must be ``"BGRA"``, ``"BGR"``, ``"RGB"`` (default), or ``"RGBA"``. :param layout: The requested layout. Must be ``"HWC"`` (default) or ``"CHW"``. - :returns: A NumPy array of dtype ``uint8``. + :param dtype: The requested data type. The default is + ``np.uint8``. + + Floating point dtypes are scaled to the ``[0, 1]`` range. Use ``channels="BGR"`` for OpenCV, and ``channels="RGB"`` (the default) for scikit-image and most other frameworks. @@ -225,49 +230,53 @@ def to_numpy(self, channels: Channels = "RGB", layout: Layout = "HWC") -> numpy. .. version-added:: 11.0.0 """ - # This is used as the channel and layout manipulation function for the other frameworks, too. This is on the - # assumption that NumPy is probably going to be the fastest on CPU, although I haven't tested this. It's also a - # prerequisite for the others, so we'd need to have a NumPy-only implementation anyway: the NumPy implementation - # shouldn't depend on PyTorch / TensorFlow, while the converse isn't true. - channels = cast("Channels", channels.upper()) layout = cast("Layout", layout.upper()) - if channels not in {"BGRA", "BGR", "RGB", "RGBA"}: - msg = "Channels must be 'BGRA', 'BGR', 'RGB', or 'RGBA'" - raise ValueError(msg) - if layout not in {"HWC", "CHW"}: - msg = "Layout must be 'HWC' or 'CHW'" - raise ValueError(msg) - import numpy as np # noqa: PLC0415 - frame = np.frombuffer(self._raw, dtype=np.uint8).reshape((self.height, self.width, 4)) + rv = np.frombuffer(self._raw, dtype=np.uint8).reshape((self.height, self.width, 4)) + if channels == "BGRA": - data = frame + pass elif channels == "BGR": - data = frame[:, :, :3] + rv = rv[:, :, :3] elif channels == "RGB": # Using [2,1,0] instead of 2::-1 would copy, rather than creating a view. - data = frame[:, :, 2::-1] - else: # RGBA + rv = rv[:, :, 2::-1] + elif channels == "RGBA": # This can't be represented as a view, since the channels within a pixel are not ordered in a way that can # be represented with a constant offset. In other words, the way that NumPy strides work, you can only make # a view if you can express the desired element order in a x:y:z style range relative to the base array's # order. - data = frame[:, :, [2, 1, 0, 3]] + rv = rv[:, :, [2, 1, 0, 3]] + else: + msg = 'Channels must be "BGRA", "BGR", "RGB", or "RGBA"' + raise ValueError(msg) - if layout == "CHW": + if layout == "HWC": + pass + elif layout == "CHW": # This will always create a view. (We're reordering the axes, not the elements.) - data = np.transpose(data, (2, 0, 1)) + rv = np.transpose(rv, (2, 0, 1)) + else: + msg = 'Layout must be "HWC" or "CHW"' + raise ValueError(msg) + + dtype = np.uint8 if dtype is None else np.dtype(dtype) + if dtype != np.uint8: + rv = rv.astype(dtype) + if np.issubdtype(dtype, np.floating): + rv /= 255.0 - return data + return rv def to_torch( self, channels: Channels = "RGB", layout: Layout = "CHW", dtype: torch.dtype | None = None, + device: torch.device | str | None = None, ) -> torch.Tensor: """Convert the screenshot to a PyTorch tensor. @@ -279,6 +288,10 @@ def to_torch( Defaults to the current PyTorch default dtype, which is usually ``torch.float32``; see :py:func:`torch.get_default_dtype`. + :param device: The requested destination device, as a + :py:class:`torch.device` or string. Default is the current + default PyTorch device; see + :py:func:`torch.get_default_device`. Floating point dtypes are scaled to the ``[0, 1]`` range. @@ -292,31 +305,65 @@ def to_torch( .. version-added:: 11.0.0 """ - import torch # noqa: PLC0415 + channels = cast("Channels", channels.upper()) + layout = cast("Layout", layout.upper()) - frame = self.to_numpy(channels=channels, layout=layout) + import torch # noqa: PLC0415 if dtype is None: - dtype = torch.get_default_dtype() - elif not isinstance(dtype, torch.dtype): + torch_dtype = torch.get_default_dtype() + elif isinstance(dtype, torch.dtype): + torch_dtype = dtype + else: msg = 'argument "dtype" must be a torch.dtype' raise TypeError(msg) + torch_device = torch.get_default_device() if device is None else torch.device(device) + + # Build a new tensor from the raw bytes. This is a view, not a copy. The shape is HWC with 4 channels in BGRA + # order at this point. The dtype here tells PyTorch how to interpret the data (unlike TensorFlow's + # convert_to_tensor with a memoryview). + rv = torch.frombuffer(self._raw, dtype=torch.uint8) + rv = rv.reshape((self.height, self.width, 4)) + + # Move the data to the desired device. If no copy is needed, this is a no-op. PyTorch using CUDA can do this + # transfer without blocking; the other devices can't. (Well, they technically can, but then our subsequent ops + # may corrupt data unless we synchronize explicitly.) + rv = rv.to(device=torch_device, non_blocking=(torch_device.type == "cuda")) + + if channels == "BGRA": + pass + elif channels == "BGR": + rv = rv[:, :, :3] + elif channels == "RGB": + # We can't use 2::-1 with PyTorch, since it doesn't support negative strides. We always have to copy. + rv = rv[:, :, [2, 1, 0]] + elif channels == "RGBA": + # This can't be represented as a view; see the comment in the NumPy version. + rv = rv[:, :, [2, 1, 0, 3]] + else: + msg = 'Channels must be "BGRA", "BGR", "RGB", or "RGBA"' + raise ValueError(msg) + + if layout == "HWC": + pass + elif layout == "CHW": + rv = rv.movedim((2, 0, 1), (0, 1, 2)) + else: + msg = 'Layout must be "HWC" or "CHW"' + raise ValueError(msg) + + # Do the conversion last to save memory bandwidth during channel shuffles. + rv = rv.to(dtype=torch_dtype) + if torch_dtype.is_floating_point: + rv.div_(255.0) - # PyTorch doesn't support negative strides like NumPy does; we have to copy in that case. This happens if the - # user requested a RGB layout. (And no, we can't use [:,:,2::-1] in PyTorch either.) - if any(s < 0 for s in frame.strides): - frame = frame.copy() - tensor = torch.from_numpy(frame) - tensor = tensor.to(dtype=dtype) - if dtype.is_floating_point: - tensor.div_(255.0) - return tensor + return rv def to_tensorflow( self, channels: Channels = "RGB", layout: Layout = "HWC", - dtype: tf.dtypes.DType | numpy.dtype | int | str = "float32", + dtype: tf.dtypes.DType | numpy.dtype | str = "float32", ) -> tf.Tensor: """Convert the screenshot to a TensorFlow tensor. @@ -326,9 +373,9 @@ def to_tensorflow( (default) or ``"CHW"``. :param dtype: The requested dtype. Can be a string like ``"float32"`` (default), a :py:class:`tf.DType - `, an int representing a TensorFlow - ``DataClass`` enum value, or a :py:class:`np.dtype - `. + `, or a :py:class:`np.dtype `. + + Device and stream management is handled by TensorFlow. Floating point dtypes are scaled to the ``[0, 1]`` range. @@ -337,20 +384,49 @@ def to_tensorflow( .. version-added:: 11.0.0 """ - import tensorflow as tf # noqa: PLC0415 + channels = cast("Channels", channels.upper()) + layout = cast("Layout", layout.upper()) - frame = self.to_numpy(channels=channels, layout=layout) + import tensorflow as tf # noqa: PLC0415 # TypeErrors from tf.as_dtype are passed up to the caller. tf_dtype = tf.as_dtype(dtype) - # TensorFlow will always copy in convert_to_tensor. - tensor = tf.convert_to_tensor(frame, dtype=tf_dtype) - if tf_dtype.is_floating: - # TensorFlow's implicit dtype conversion rules are not trivial. We use an explicit dtype on both sides - # instead, by making a tf.constant. - tensor = tensor / tf.constant(255.0, dtype=tf_dtype) - return tensor + # We don't need to do anything explicit about device management in TensorFlow; it handles that for us. + # convert_to_tensor will always copy. We verified in __init__ that self._raw is a memoryview of unsigned bytes, + # so that's what TensorFlow will take as the source format; we just give an explicit dtype for clarity. + rv = tf.convert_to_tensor(self._raw, dtype=tf.uint8) + rv = tf.reshape(rv, (self.height, self.width, 4)) + + if channels == "BGRA": + pass + elif channels == "BGR": + rv = rv[:, :, :3] + elif channels == "RGB": + rv = tf.gather(rv, [2, 1, 0], axis=2) + elif channels == "RGBA": + rv = tf.gather(rv, [2, 1, 0, 3], axis=2) + else: + msg = 'Channels must be "BGRA", "BGR", "RGB", or "RGBA"' + raise ValueError(msg) + + if layout == "HWC": + pass + elif layout == "CHW": + rv = tf.transpose(rv, perm=(2, 0, 1)) + else: + msg = 'Layout must be "HWC" or "CHW"' + raise ValueError(msg) + + # Do the conversion last to save memory bandwidth during channel shuffles. + if tf_dtype != tf.uint8: + rv = tf.cast(rv, dtype=tf_dtype) + if tf_dtype.is_floating: + # TensorFlow's implicit dtype conversion rules are not trivial. We use an explicit dtype on both sides + # instead, by making a tf.constant. + rv = rv / tf.constant(255.0, dtype=tf_dtype) + + return rv @property def top(self) -> int: diff --git a/src/tests/third_party/test_numpy_method.py b/src/tests/third_party/test_numpy_method.py index cb35a460..4faf5c6e 100644 --- a/src/tests/third_party/test_numpy_method.py +++ b/src/tests/third_party/test_numpy_method.py @@ -46,7 +46,7 @@ def test_to_numpy_bad_channels() -> None: raw = bytearray([0, 0, 0, 0]) shot = ScreenShot.from_size(raw, 1, 1) - with pytest.raises(ValueError, match="Channels must be 'BGRA', 'BGR', 'RGB', or 'RGBA'"): + with pytest.raises(ValueError, match='Channels must be "BGRA", "BGR", "RGB", or "RGBA"'): shot.to_numpy(channels="gray") # type: ignore[arg-type] @@ -54,7 +54,7 @@ def test_to_numpy_bad_layout() -> None: raw = bytearray([0, 0, 0, 0]) shot = ScreenShot.from_size(raw, 1, 1) - with pytest.raises(ValueError, match="Layout must be 'HWC' or 'CHW'"): + with pytest.raises(ValueError, match='Layout must be "HWC" or "CHW"'): shot.to_numpy(layout="NHWC") # type: ignore[arg-type] diff --git a/src/tests/third_party/test_torch_method.py b/src/tests/third_party/test_torch_method.py index 5846c68e..634283d0 100644 --- a/src/tests/third_party/test_torch_method.py +++ b/src/tests/third_party/test_torch_method.py @@ -35,16 +35,32 @@ def test_to_torch_dtype_uint8() -> None: @pytest.mark.parametrize("layout", ["HWC", "CHW"]) @pytest.mark.parametrize("channels", ["BGRA", "BGR", "RGBA", "RGB"]) -def test_to_torch_permutations(framework_test_image: ScreenShot, channels: str, layout: str) -> None: +@pytest.mark.parametrize("cuda", [False, True]) +def test_to_torch_permutations(framework_test_image: ScreenShot, channels: str, layout: str, cuda: bool) -> None: """Test all permutations of channels and layouts.""" + if cuda and not torch.cuda.is_available: + # The CUDA versions won't be run in CI/CD, but it's still worth checking them on developers' machines if they + # happen to have PyTorch with CUDA support. + pytest.skip() + uint8_target = reordered_test_image(channels=channels, layout=layout) bfloat16_target = uint8_target.astype(np.float32) / 255.0 - uint8_result = framework_test_image.to_torch(channels=channels, layout=layout, dtype=torch.uint8) # type: ignore[arg-type] - assert np.array_equal(uint8_result.numpy(), uint8_target) + uint8_result = framework_test_image.to_torch( + channels=channels, # type: ignore[arg-type] + layout=layout, # type: ignore[arg-type] + dtype=torch.uint8, + device="cuda" if cuda else "cpu", + ) + assert np.array_equal(uint8_result.cpu().numpy(), uint8_target) - bfloat16_result = framework_test_image.to_torch(channels=channels, layout=layout, dtype=torch.bfloat16) # type: ignore[arg-type] + bfloat16_result = framework_test_image.to_torch( + channels=channels, # type: ignore[arg-type] + layout=layout, # type: ignore[arg-type] + dtype=torch.bfloat16, + device="cuda" if cuda else "cpu", + ) # We have to explicitly cast back to float32 for the comparison, because PyTorch won't directly convert bfloat16 to # NumPy. bfloat16_result = bfloat16_result.to(torch.float32) - assert np.allclose(bfloat16_result.numpy(), bfloat16_target, rtol=0, atol=1 / 512.0) + assert np.allclose(bfloat16_result.cpu().numpy(), bfloat16_target, rtol=0, atol=1 / 512.0) From 706ea74dba5b03285bb5d3d2a2a579ab6ed02a49 Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Sat, 27 Jun 2026 22:45:26 -0700 Subject: [PATCH 13/20] You actually have to call functions for them to work. :-/ --- src/tests/third_party/test_torch_method.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tests/third_party/test_torch_method.py b/src/tests/third_party/test_torch_method.py index 634283d0..e892b751 100644 --- a/src/tests/third_party/test_torch_method.py +++ b/src/tests/third_party/test_torch_method.py @@ -38,7 +38,7 @@ def test_to_torch_dtype_uint8() -> None: @pytest.mark.parametrize("cuda", [False, True]) def test_to_torch_permutations(framework_test_image: ScreenShot, channels: str, layout: str, cuda: bool) -> None: """Test all permutations of channels and layouts.""" - if cuda and not torch.cuda.is_available: + if cuda and not torch.cuda.is_available(): # The CUDA versions won't be run in CI/CD, but it's still worth checking them on developers' machines if they # happen to have PyTorch with CUDA support. pytest.skip() From 38c50b0a787a04d56cff0e2a8f7c00f9ad996a93 Mon Sep 17 00:00:00 2001 From: Joel Holveck Date: Sun, 28 Jun 2026 08:00:23 +0000 Subject: [PATCH 14/20] Minor fixes from AI review --- docs/source/examples.rst | 2 +- docs/source/examples/opencv_numpy.py | 5 +++-- docs/source/index.rst | 2 +- docs/source/release-history/v11.0.0.md | 2 +- docs/source/usage.rst | 4 ++-- src/tests/test_setup.py | 11 ++++++----- src/tests/third_party/array_frameworks/__init__.py | 0 .../third_party/{ => array_frameworks}/conftest.py | 1 + .../third_party/{ => array_frameworks}/test_numpy.py | 4 ++-- .../{ => array_frameworks}/test_numpy_method.py | 2 +- .../{ => array_frameworks}/test_tensorflow_method.py | 2 +- .../{ => array_frameworks}/test_torch_method.py | 4 ++-- 12 files changed, 21 insertions(+), 18 deletions(-) create mode 100644 src/tests/third_party/array_frameworks/__init__.py rename src/tests/third_party/{ => array_frameworks}/conftest.py (97%) rename src/tests/third_party/{ => array_frameworks}/test_numpy.py (83%) rename src/tests/third_party/{ => array_frameworks}/test_numpy_method.py (96%) rename src/tests/third_party/{ => array_frameworks}/test_tensorflow_method.py (95%) rename src/tests/third_party/{ => array_frameworks}/test_torch_method.py (94%) diff --git a/docs/source/examples.rst b/docs/source/examples.rst index cf77b7a4..e84a89f2 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -142,7 +142,7 @@ This is an example using :py:meth:`mss.ScreenShot.to_pil` and direct pixel edits .. versionadded:: 3.0.0 -OpenCV/Numpy +OpenCV/NumPy ============ See how fast you can record the screen. diff --git a/docs/source/examples/opencv_numpy.py b/docs/source/examples/opencv_numpy.py index 5ca3e952..820bf4a5 100644 --- a/docs/source/examples/opencv_numpy.py +++ b/docs/source/examples/opencv_numpy.py @@ -17,11 +17,12 @@ while "Screen capturing": last_time = time.time() - # Get raw pixels from the screen, save it to a NumPy array + # Get raw pixels from the screen, save it to a NumPy array. + # Note that OpenCV expects colors in BGR order. img = sct.grab(monitor).to_numpy(channels="BGR") # Display the picture - cv2.imshow("OpenCV/Numpy normal", img) + cv2.imshow("OpenCV/NumPy normal", img) # Display the picture in grayscale # cv2.imshow('OpenCV/Numpy grayscale', diff --git a/docs/source/index.rst b/docs/source/index.rst index df5130aa..4cfc3c81 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -23,7 +23,7 @@ An ultra fast cross-platform multiple screenshots module in pure python using ct - **Python 3.10+**, :pep:`8` compliant, no dependency, thread-safe; - very basic, it will grab one screenshot by monitor or a screenshot of all monitors and save it to a PNG file; - but you can use PIL and benefit from all its formats (or add yours directly); - - integrate well with Numpy and OpenCV; + - integrate well with NumPy, OpenCV, PyTorch, and TensorFlow; - it could be easily embedded into games and other software which require fast and platform optimized methods to grab screenshots (like AI, Computer Vision); - get the `source code on GitHub `_; - learn with a `bunch of examples `_; diff --git a/docs/source/release-history/v11.0.0.md b/docs/source/release-history/v11.0.0.md index 9fa58139..ef64efb9 100644 --- a/docs/source/release-history/v11.0.0.md +++ b/docs/source/release-history/v11.0.0.md @@ -20,7 +20,7 @@ The {py:attr}`mss.ScreenShot.raw` attribute has been deprecated, and will soon b The {py:attr}`mss.ScreenShot.bgra` and {py:attr}`mss.ScreenShot.rgb` properties now will return bytes-like {py:class}`memoryview` objects, rather than {py:class}`bytes` or {py:class}`bytearray` objects. For practical use -cases, this should not be noticible. This change allows faster access to screenshot data, with fewer memory copies. +cases, this should not be noticeable. This change allows faster access to screenshot data, with fewer memory copies. ### Python 3.9 EOL diff --git a/docs/source/usage.rst b/docs/source/usage.rst index a84170e2..d52370fb 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -106,8 +106,8 @@ Some examples include the following libraries: When using the NumPy array interface protocol, the returned object is in HWC (height, width, channels) format, with the channels in BGRA order. -Note that OpenCV uses RGB order, rather than the BGRA order used in this automatic conversion. You may prefer to -use the :py:meth:`mss.ScreenShot.to_numpy` method instead. +Note that most libraries do not expect the alpha channel to be present, or expect an order other than the BGRA order +used in this automatic conversion. You may prefer to use the :py:meth:`mss.ScreenShot.to_numpy` method instead. Alpha Channel ------------- diff --git a/src/tests/test_setup.py b/src/tests/test_setup.py index 4160f210..f77f6275 100644 --- a/src/tests/test_setup.py +++ b/src/tests/test_setup.py @@ -113,13 +113,14 @@ def test_sdist() -> None: f"mss-{__version__}/src/tests/test_windows.py", f"mss-{__version__}/src/tests/test_xcb.py", f"mss-{__version__}/src/tests/third_party/__init__.py", - f"mss-{__version__}/src/tests/third_party/conftest.py", - f"mss-{__version__}/src/tests/third_party/test_numpy.py", - f"mss-{__version__}/src/tests/third_party/test_numpy_method.py", + f"mss-{__version__}/src/tests/third_party/array_frameworks/__init__.py", + f"mss-{__version__}/src/tests/third_party/array_frameworks/conftest.py", + f"mss-{__version__}/src/tests/third_party/array_frameworks/test_numpy.py", + f"mss-{__version__}/src/tests/third_party/array_frameworks/test_numpy_method.py", + f"mss-{__version__}/src/tests/third_party/array_frameworks/test_tensorflow_method.py", + f"mss-{__version__}/src/tests/third_party/array_frameworks/test_torch_method.py", f"mss-{__version__}/src/tests/third_party/test_pil.py", f"mss-{__version__}/src/tests/third_party/test_pil_method.py", - f"mss-{__version__}/src/tests/third_party/test_tensorflow_method.py", - f"mss-{__version__}/src/tests/third_party/test_torch_method.py", f"mss-{__version__}/src/tests/thread_helpers.py", f"mss-{__version__}/src/xcbproto/README.md", f"mss-{__version__}/src/xcbproto/gen_xcb_to_py.py", diff --git a/src/tests/third_party/array_frameworks/__init__.py b/src/tests/third_party/array_frameworks/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/tests/third_party/conftest.py b/src/tests/third_party/array_frameworks/conftest.py similarity index 97% rename from src/tests/third_party/conftest.py rename to src/tests/third_party/array_frameworks/conftest.py index 4661f06a..565bc0f6 100644 --- a/src/tests/third_party/conftest.py +++ b/src/tests/third_party/array_frameworks/conftest.py @@ -11,6 +11,7 @@ # importorskip), but type checkers don't always know about importorskip. import numpy as np_typehints # noqa: ICN001 +# This entire tree is skipped if NumPy isn't installed. np = pytest.importorskip("numpy") TEST_SIZE = (320, 240) diff --git a/src/tests/third_party/test_numpy.py b/src/tests/third_party/array_frameworks/test_numpy.py similarity index 83% rename from src/tests/third_party/test_numpy.py rename to src/tests/third_party/array_frameworks/test_numpy.py index c06e6180..f216cf3a 100644 --- a/src/tests/third_party/test_numpy.py +++ b/src/tests/third_party/array_frameworks/test_numpy.py @@ -6,9 +6,9 @@ import pytest -from mss import MSS +import numpy as np -np = pytest.importorskip("numpy", reason="Numpy module not available.") +from mss import MSS def test_numpy(mss_impl: Callable[..., MSS]) -> None: diff --git a/src/tests/third_party/test_numpy_method.py b/src/tests/third_party/array_frameworks/test_numpy_method.py similarity index 96% rename from src/tests/third_party/test_numpy_method.py rename to src/tests/third_party/array_frameworks/test_numpy_method.py index 4faf5c6e..8acab8be 100644 --- a/src/tests/third_party/test_numpy_method.py +++ b/src/tests/third_party/array_frameworks/test_numpy_method.py @@ -7,7 +7,7 @@ import pytest from mss import ScreenShot -from tests.third_party.conftest import reordered_test_image +from tests.third_party.array_frameworks.conftest import reordered_test_image np = pytest.importorskip("numpy") diff --git a/src/tests/third_party/test_tensorflow_method.py b/src/tests/third_party/array_frameworks/test_tensorflow_method.py similarity index 95% rename from src/tests/third_party/test_tensorflow_method.py rename to src/tests/third_party/array_frameworks/test_tensorflow_method.py index 602d1184..99463ac7 100644 --- a/src/tests/third_party/test_tensorflow_method.py +++ b/src/tests/third_party/array_frameworks/test_tensorflow_method.py @@ -7,7 +7,7 @@ import pytest from mss import ScreenShot -from tests.third_party.conftest import reordered_test_image +from tests.third_party.array_frameworks.conftest import reordered_test_image np = pytest.importorskip("numpy") tf = pytest.importorskip("tensorflow") diff --git a/src/tests/third_party/test_torch_method.py b/src/tests/third_party/array_frameworks/test_torch_method.py similarity index 94% rename from src/tests/third_party/test_torch_method.py rename to src/tests/third_party/array_frameworks/test_torch_method.py index e892b751..f5008fbc 100644 --- a/src/tests/third_party/test_torch_method.py +++ b/src/tests/third_party/array_frameworks/test_torch_method.py @@ -7,7 +7,7 @@ import pytest from mss import ScreenShot -from tests.third_party.conftest import reordered_test_image +from tests.third_party.array_frameworks.conftest import reordered_test_image np = pytest.importorskip("numpy") torch = pytest.importorskip("torch") @@ -19,7 +19,7 @@ def test_to_torch_default() -> None: tensor = shot.to_torch() assert tuple(tensor.shape) == (3, 1, 1) - assert tensor.dtype == torch.float32 + assert tensor.dtype == torch.get_default_dtype() expected = torch.tensor([3 / 255.0, 2 / 255.0, 1 / 255.0], dtype=torch.float32) assert torch.allclose(tensor[:, 0, 0], expected) From 672e0f004d3e41b89b6404a46aa3e7a1ae0cfe8d Mon Sep 17 00:00:00 2001 From: Joel Holveck Date: Sun, 28 Jun 2026 09:59:05 +0000 Subject: [PATCH 15/20] Improve to_torch performance on the CPU The PyTorch implementation of the channel shuffling operations on the CPU is significantly slower (like by 5x or so) than the NumPy implementation. (It's still much faster on the GPU, were the PyTorch devs spend their focus, and also have much faster RAM.) So, if the user requests a CPU tensor, we now do the channel shuffling work in NumPy. --- src/mss/screenshot.py | 15 ++++++++++++++- .../third_party/array_frameworks/test_numpy.py | 2 -- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/mss/screenshot.py b/src/mss/screenshot.py index 3510d847..87cce3bf 100644 --- a/src/mss/screenshot.py +++ b/src/mss/screenshot.py @@ -271,7 +271,7 @@ def to_numpy( return rv - def to_torch( + def to_torch( # noqa: PLR0912 self, channels: Channels = "RGB", layout: Layout = "CHW", @@ -319,6 +319,19 @@ def to_torch( raise TypeError(msg) torch_device = torch.get_default_device() if device is None else torch.device(device) + if torch_device.type == "cpu": + # NumPy handles the necessary CPU operations significantly more efficiently than PyTorch, so we defer to it. + ndarray = self.to_numpy(channels=channels, layout=layout) + # to_numpy can return tensors with negative strides, which PyTorch doesn't support. + if any(s < 0 for s in ndarray.strides): + ndarray = ndarray.copy() + rv = torch.from_numpy(ndarray) + # We do the dtype conversion ourselves because PyTorch has dtypes that NumPy doesn't, like bfloat16. + rv = rv.to(dtype=torch_dtype, device=torch_device) + if torch_dtype.is_floating_point: + rv.div_(255.0) + return rv + # Build a new tensor from the raw bytes. This is a view, not a copy. The shape is HWC with 4 channels in BGRA # order at this point. The dtype here tells PyTorch how to interpret the data (unlike TensorFlow's # convert_to_tensor with a memoryview). diff --git a/src/tests/third_party/array_frameworks/test_numpy.py b/src/tests/third_party/array_frameworks/test_numpy.py index f216cf3a..0bf39413 100644 --- a/src/tests/third_party/array_frameworks/test_numpy.py +++ b/src/tests/third_party/array_frameworks/test_numpy.py @@ -4,8 +4,6 @@ from collections.abc import Callable -import pytest - import numpy as np from mss import MSS From a914e4add8086642dee6d55f1efa08508a181c1a Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Sun, 28 Jun 2026 16:15:23 -0700 Subject: [PATCH 16/20] Don't use a non-blocking copy. It's fragile, harder to manage lifetimes, and doesn't win us measurable gains. --- src/mss/screenshot.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mss/screenshot.py b/src/mss/screenshot.py index 87cce3bf..10e82a25 100644 --- a/src/mss/screenshot.py +++ b/src/mss/screenshot.py @@ -338,10 +338,10 @@ def to_torch( # noqa: PLR0912 rv = torch.frombuffer(self._raw, dtype=torch.uint8) rv = rv.reshape((self.height, self.width, 4)) - # Move the data to the desired device. If no copy is needed, this is a no-op. PyTorch using CUDA can do this - # transfer without blocking; the other devices can't. (Well, they technically can, but then our subsequent ops - # may corrupt data unless we synchronize explicitly.) - rv = rv.to(device=torch_device, non_blocking=(torch_device.type == "cuda")) + # Move the data to the desired device. If no copy is needed, this is a no-op. + # We don't use a non-blocking copy because it can be fragile, hard to manage lifetimes, and doesn't win us + # measurable gains. + rv = rv.to(device=torch_device) if channels == "BGRA": pass From 43d5b076933dfe93f2347dc12d2aaba6fbfe02b1 Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Mon, 29 Jun 2026 17:52:25 -0700 Subject: [PATCH 17/20] Speedups for to_tensorflow TensorFlow is designed around graphs of ops for a DNN, not eager channel shuffling and similar operations. It's much faster to do the channel shuffles in NumPy, although we still do the dtype conversion in TensorFlow. --- src/mss/screenshot.py | 39 ++++++------------- .../test_tensorflow_method.py | 2 + .../array_frameworks/test_torch_method.py | 8 ++-- 3 files changed, 18 insertions(+), 31 deletions(-) diff --git a/src/mss/screenshot.py b/src/mss/screenshot.py index 10e82a25..723cd7a0 100644 --- a/src/mss/screenshot.py +++ b/src/mss/screenshot.py @@ -405,39 +405,22 @@ def to_tensorflow( # TypeErrors from tf.as_dtype are passed up to the caller. tf_dtype = tf.as_dtype(dtype) - # We don't need to do anything explicit about device management in TensorFlow; it handles that for us. - # convert_to_tensor will always copy. We verified in __init__ that self._raw is a memoryview of unsigned bytes, - # so that's what TensorFlow will take as the source format; we just give an explicit dtype for clarity. - rv = tf.convert_to_tensor(self._raw, dtype=tf.uint8) - rv = tf.reshape(rv, (self.height, self.width, 4)) - - if channels == "BGRA": - pass - elif channels == "BGR": - rv = rv[:, :, :3] - elif channels == "RGB": - rv = tf.gather(rv, [2, 1, 0], axis=2) - elif channels == "RGBA": - rv = tf.gather(rv, [2, 1, 0, 3], axis=2) - else: - msg = 'Channels must be "BGRA", "BGR", "RGB", or "RGBA"' - raise ValueError(msg) + # Whether we're sending to a CPU or GPU, the channels and layout conversion are much faster in NumPy. I suspect + # this is because doing this in TensorFlow eager mode materializes the tensor at each operation, but that's just + # a guess, and putting the channel shuffles in a tf.function didn't help. It's still best (especially on GPU) + # to do the dtype conversion in TensorFlow. + ndarray = self.to_numpy(channels=channels, layout=layout) - if layout == "HWC": - pass - elif layout == "CHW": - rv = tf.transpose(rv, perm=(2, 0, 1)) - else: - msg = 'Layout must be "HWC" or "CHW"' - raise ValueError(msg) + # We don't need to do anything explicit about device management in TensorFlow; it handles that for us. + # convert_to_tensor will always copy. + rv = tf.convert_to_tensor(ndarray) - # Do the conversion last to save memory bandwidth during channel shuffles. if tf_dtype != tf.uint8: + # I have no idea why, but doing the cast here instead of in convert_to_tensor is significantly faster, + # especially on the GPU. rv = tf.cast(rv, dtype=tf_dtype) if tf_dtype.is_floating: - # TensorFlow's implicit dtype conversion rules are not trivial. We use an explicit dtype on both sides - # instead, by making a tf.constant. - rv = rv / tf.constant(255.0, dtype=tf_dtype) + rv /= 255.0 return rv diff --git a/src/tests/third_party/array_frameworks/test_tensorflow_method.py b/src/tests/third_party/array_frameworks/test_tensorflow_method.py index 99463ac7..3dd8c7df 100644 --- a/src/tests/third_party/array_frameworks/test_tensorflow_method.py +++ b/src/tests/third_party/array_frameworks/test_tensorflow_method.py @@ -45,7 +45,9 @@ def test_to_tensorflow_permutations(framework_test_image: ScreenShot, channels: bfloat16_target = uint8_target.astype(np.float32) / 255.0 uint8_result = framework_test_image.to_tensorflow(channels=channels, layout=layout, dtype="uint8") # type: ignore[arg-type] + assert uint8_result.dtype == tf.uint8 assert np.array_equal(uint8_result.numpy(), uint8_target) bfloat16_result = framework_test_image.to_tensorflow(channels=channels, layout=layout, dtype="bfloat16") # type: ignore[arg-type] + assert bfloat16_result.dtype == tf.bfloat16 assert np.allclose(bfloat16_result.numpy(), bfloat16_target, rtol=0, atol=1 / 512.0) diff --git a/src/tests/third_party/array_frameworks/test_torch_method.py b/src/tests/third_party/array_frameworks/test_torch_method.py index f5008fbc..7226b1e2 100644 --- a/src/tests/third_party/array_frameworks/test_torch_method.py +++ b/src/tests/third_party/array_frameworks/test_torch_method.py @@ -41,7 +41,7 @@ def test_to_torch_permutations(framework_test_image: ScreenShot, channels: str, if cuda and not torch.cuda.is_available(): # The CUDA versions won't be run in CI/CD, but it's still worth checking them on developers' machines if they # happen to have PyTorch with CUDA support. - pytest.skip() + pytest.skip("CUDA is not available") uint8_target = reordered_test_image(channels=channels, layout=layout) bfloat16_target = uint8_target.astype(np.float32) / 255.0 @@ -52,6 +52,7 @@ def test_to_torch_permutations(framework_test_image: ScreenShot, channels: str, dtype=torch.uint8, device="cuda" if cuda else "cpu", ) + assert uint8_result.dtype == torch.uint8 assert np.array_equal(uint8_result.cpu().numpy(), uint8_target) bfloat16_result = framework_test_image.to_torch( @@ -60,7 +61,8 @@ def test_to_torch_permutations(framework_test_image: ScreenShot, channels: str, dtype=torch.bfloat16, device="cuda" if cuda else "cpu", ) + assert bfloat16_result.dtype == torch.bfloat16 # We have to explicitly cast back to float32 for the comparison, because PyTorch won't directly convert bfloat16 to # NumPy. - bfloat16_result = bfloat16_result.to(torch.float32) - assert np.allclose(bfloat16_result.cpu().numpy(), bfloat16_target, rtol=0, atol=1 / 512.0) + bfloat16_result_f32 = bfloat16_result.to(torch.float32) + assert np.allclose(bfloat16_result_f32.cpu().numpy(), bfloat16_target, rtol=0, atol=1 / 512.0) From 13e1366c1838b4581588a3dad40ff13df688b18b Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Mon, 29 Jun 2026 23:25:23 -0700 Subject: [PATCH 18/20] Add release notes --- docs/source/release-history/v11.0.0.md | 10 ++++++++++ docs/source/usage.rst | 2 ++ 2 files changed, 12 insertions(+) diff --git a/docs/source/release-history/v11.0.0.md b/docs/source/release-history/v11.0.0.md index ef64efb9..c39c5e6e 100644 --- a/docs/source/release-history/v11.0.0.md +++ b/docs/source/release-history/v11.0.0.md @@ -47,6 +47,16 @@ processing time decreased from 22.64 ms to 18.59 ms per frame (approximately 18% Support for additional operating systems is planned. +### Export Methods + +There are now convenient, easy-to-use methods to export screenshots to several popular formats: the popular PIL image +library, the NumPy array format for scientific computing, and the PyTorch and TensorFlow deep learning frameworks. + +Exporting to these formats was always possible from MSS, but the best way to do it hasn’t always been obvious. By +providing these methods, MSS gives you tested, high-speed ways to transfer the data to these popular formats. + +See the documentation section {ref}`accessing_pixel_data` for details. + ### General Improvements The MSS context object will now always surface inner exceptions, even if `__exit__` may also generate an exception during tear-down. diff --git a/docs/source/usage.rst b/docs/source/usage.rst index d52370fb..3ef6cbe1 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -56,6 +56,8 @@ For instance, you can capture the first monitor and get a :py:class:`mss.ScreenS Ok, now you've got the :py:class:`mss.ScreenShot` object. But what do you do with it? +.. _accessing_pixel_data: + Accessing Pixel Data ==================== From 8f0060cb44d2e2828f30e80b77dc8f5fce6def30 Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Tue, 30 Jun 2026 01:25:28 -0700 Subject: [PATCH 19/20] Document primary_monitor in the usage section As suggested by @BoboTiG Also a couple of other tiny improvements in that area I noticed. --- docs/source/usage.rst | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/docs/source/usage.rst b/docs/source/usage.rst index 3ef6cbe1..d0e1fa4c 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -2,6 +2,10 @@ Usage ===== +.. role:: python(code) + :language: python + :class: highlight + Import ====== @@ -41,8 +45,17 @@ copy and paste. If instead you want to use the pixel data yourself, you can do so easily, with the :py:meth:`mss.MSS.grab` method. You'll first need to decide whether you want to capture all the monitors, a single monitor, or a specific region of the -screen. The :py:class:`mss.MSS` object has a :py:attr:`mss.MSS.monitors` attribute that is a list of all the monitors, -starting from index 1, as well as the full virtual screen (all monitors combined) at index 0. +screen. + +For capturing one or more monitors, the :py:class:`mss.MSS` object has a :py:attr:`mss.MSS.monitors` attribute that is a +list of all the monitors, starting from index 1, as well as the full virtual screen (all monitors combined) at index 0. +The primary monitor, the one that holds the taskbar or similar system UI, is available as +:py:attr:`mss.MSS.primary_monitor`. + +For capturing a specific region, you can pass :py:meth:`mss.MSS.grab` a dictionary with the keys ``top``, ``left``, +``width``, and ``height``. For instance, to capture a 100x100 pixel region starting at the top-left corner of the +screen, you could use :python:`{"top": 0, "left": 0, "width": 100, "height": 100}`. You can also use a PIL-style box, +which is a 4-tuple of ``(left, top, right, bottom)``. Once you've decided what you want to capture, you can call :py:meth:`mss.MSS.grab` with the appropriate monitor or region. This will return a :py:class:`mss.ScreenShot` object, which contains the pixel data and other information about @@ -51,7 +64,7 @@ the screenshot. For instance, you can capture the first monitor and get a :py:class:`mss.ScreenShot` object like this:: with MSS() as sct: - sct_img = sct.grab(sct.monitors[1]) # Capture the first monitor + sct_img = sct.grab(sct.primary_monitor) Ok, now you've got the :py:class:`mss.ScreenShot` object. But what do you do with it? @@ -106,10 +119,11 @@ Some examples include the following libraries: * Some functions from `OpenCV `_, a popular computer vision library When using the NumPy array interface protocol, the returned object is in HWC (height, width, channels) format, with the -channels in BGRA order. +channels in BGRA order, and with a data type of :py:attr:`numpy.uint8`. Note that most libraries do not expect the alpha channel to be present, or expect an order other than the BGRA order -used in this automatic conversion. You may prefer to use the :py:meth:`mss.ScreenShot.to_numpy` method instead. +used in this automatic conversion. You may prefer to use the :py:meth:`mss.ScreenShot.to_numpy` method instead, since +it can return the pixel data in most common layouts, orders, and data types. Alpha Channel ------------- From a2e1e8b1c47460d8a68587cce98dd29cc71c24f6 Mon Sep 17 00:00:00 2001 From: Joel Ray Holveck Date: Tue, 30 Jun 2026 01:54:27 -0700 Subject: [PATCH 20/20] Add a 10-minute timeout to tests The default is six hours for the whole job. We've been seeing PyTest freezing during some macOS runs. Not sure what's up, but at least we can not wait for a six hour timeout. Most tests seem to take between two and three minutes, so this seems like quite a generous timeout. --- .github/workflows/tests.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index a10bce61..936cdd80 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -107,9 +107,11 @@ jobs: run: python -m pip install tensorflow - name: Tests (GNU/Linux) + timeout-minutes: 10 if: matrix.os.name == 'linux' run: xvfb-run python -m pytest - name: Tests (macOS, Windows) + timeout-minutes: 10 if: matrix.os.name != 'linux' run: python -m pytest