diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 3baeb668..5f565dc1 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,24 @@ jobs: run: | python -m pip install -U pip python -m pip install -e '.[dev,tests]' + # 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 writing, 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/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 diff --git a/demos/tinytv-stream-simple.py b/demos/tinytv-stream-simple.py index c10129b8..1fd3986c 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.frombuffer("RGB", screenshot.size, screenshot.bgra, "raw", "BGRX", 0, 1) + 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 c2781fcf..0bda3c1d 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.frombuffer("RGB", sct_img.size, sct_img.bgra, "raw", "BGRX", 0, 1) + 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 0e33b58f..22c5c0df 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.frombuffer("RGB", screenshot.size, screenshot.bgra, "raw", "BGRX", 0, 1) + 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 cc89ca99..66f068cc 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/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/examples.rst b/docs/source/examples.rst index 0085caeb..cf77b7a4 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -1,3 +1,5 @@ +.. _examples: + ======== Examples ======== @@ -122,8 +124,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 `frombuffer() `_: +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 +135,7 @@ This is an example using `frombuffer() `_: +This is an example using :py:meth:`mss.ScreenShot.to_pil` and direct pixel edits: .. literalinclude:: examples/pil_pixels.py :lines: 7- @@ -147,6 +149,9 @@ 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.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 1eb51fd5..5ca3e952 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 f888a4e3..f9b90247 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.frombuffer("RGB", sct_img.size, sct_img.bgra, "raw", "BGRX", 0, 1) - # The same, but less efficient: - # img = Image.frombuffer('RGB', sct_img.size, sct_img.rgb, 'raw', 'RGB', 0, 1) + 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 4398192d..57f50952 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,7 +11,7 @@ sct_img = sct.grab(sct.monitors[1]) # Create an Image - img = Image.new("RGB", sct_img.size) + img = sct_img.to_pil("RGB") # Best solution: create a list(tuple(R, G, B), ...) for putdata() pixels = zip(sct_img.bgra[2::4], sct_img.bgra[1::4], sct_img.bgra[::4], strict=False) diff --git a/docs/source/release-history/v11.0.0.md b/docs/source/release-history/v11.0.0.md index d4eef6b5..9fa58139 100644 --- a/docs/source/release-history/v11.0.0.md +++ b/docs/source/release-history/v11.0.0.md @@ -19,9 +19,8 @@ The {py:attr}`mss.ScreenShot.raw` attribute has been deprecated, and will soon b {py:attr}`mss.ScreenShot.bgra` property instead. The {py:attr}`mss.ScreenShot.bgra` and {py:attr}`mss.ScreenShot.rgb` properties now will return bytes-like -{py:type}`memoryview` objects, not necessarily {py:type}`bytes` or {py:type}`bytearray` objects. For practical use -cases, this should not be noticible. This change allows faster access to screenshot data, with fewer memory -copies. +{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. ### Python 3.9 EOL @@ -51,10 +50,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..a84170e2 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -32,6 +32,129 @@ 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 +==================== + +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. + +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. +* :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 `_. 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 `_. +* :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 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 +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 +264,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 d5081173..b9f0da66 100644 --- a/src/mss/screenshot.py +++ b/src/mss/screenshot.py @@ -4,7 +4,7 @@ from __future__ import annotations import warnings -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Literal, cast from mss.exception import ScreenShotError from mss.models import Monitor, Pixel, Pixels, Pos, Size @@ -13,8 +13,17 @@ from collections.abc import Iterator from typing import Any + # 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"] + class ScreenShot: """Screenshot object. @@ -36,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." ) @@ -58,13 +64,10 @@ def __array_interface__(self) -> dict[str, Any]: :py:func:`tf.convert_to_tensor`), JAX (via :py:func:`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). - - 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:: @@ -85,17 +88,13 @@ 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 BGRxBGRx... sequence. A specific pixel can be accessed as ``bgra[(y * width + x) * 4:(y * width + x) * 4 + 4].`` - The memoryview is read-write, for compatibility with ctypes' - :py:func:`from_buffer`. However, actually modifying the data - may cause undefined behavior. - .. note:: While the name is ``bgra``, the alpha channel may or may not be valid. @@ -111,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 @@ -162,10 +161,6 @@ def rgb(self) -> memoryview: RGBRGB... sequence. A specific pixel can be accessed as ``rgb[(y * width + x) * 3:(y * width + x) * 3 + 3].`` - The memoryview is read-write, for compatibility with ctypes' - :py:func:`from_buffer`. 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 efficient. @@ -192,6 +187,171 @@ def rgb(self) -> memoryview: return self.__rgb + def to_pil(self, mode: str = "RGB") -> PIL.Image.Image: + """Convert the screenshot to a Pillow image. + + :param mode: The requested image mode. Must be ``"RGB"`` + (default) or ``"RGBA"``. + + When requesting ``"RGBA"``, the alpha channel may not represent + meaningful transparency on all platforms/backends. + + .. version-added:: 11.0.0 + """ + 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") -> 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``. + + Use ``channels="BGR"`` for OpenCV, and ``channels="RGB"`` (the + default) for scikit-image and most other frameworks. + + When requesting ``"RGBA"`` or ``"BGRA"``, the alpha channel may + not represent meaningful transparency on all platforms/backends. + + .. 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)) + if channels == "BGRA": + data = frame + elif channels == "BGR": + data = frame[:, :, :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 + # 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 + + def to_torch( + self, + channels: Channels = "RGB", + layout: Layout = "CHW", + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + """Convert the screenshot to a PyTorch tensor. + + :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 :py:class:`torch.dtype`. + 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. + + 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"`` or ``"BGRA"``, the alpha channel may + not represent meaningful transparency on all platforms/backends. + + .. version-added:: 11.0.0 + """ + import torch # noqa: PLC0415 + + frame = self.to_numpy(channels=channels, layout=layout) + + if dtype is None: + 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.div_(255.0) + return tensor + + def to_tensorflow( + self, + channels: Channels = "RGB", + layout: Layout = "HWC", + dtype: tf.dtypes.DType | numpy.dtype | int | str = "float32", + ) -> tf.Tensor: + """Convert the screenshot to a TensorFlow tensor. + + :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 :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. + + .. version-added:: 11.0.0 + """ + import tensorflow as tf # noqa: PLC0415 + + frame = self.to_numpy(channels=channels, layout=layout) + + # 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 + @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 18e24f08..4160f210 100644 --- a/src/tests/test_setup.py +++ b/src/tests/test_setup.py @@ -113,8 +113,13 @@ 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", + 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/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 new file mode 100644 index 00000000..cb35a460 --- /dev/null +++ b/src/tests/third_party/test_numpy_method.py @@ -0,0 +1,67 @@ +"""This is part of the MSS Python's module. +Source: https://github.com/BoboTiG/python-mss. +""" + +from __future__ import annotations + +import pytest + +from mss import ScreenShot +from tests.third_party.conftest import reordered_test_image + +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] + + +@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_pil_method.py b/src/tests/third_party/test_pil_method.py new file mode 100644 index 00000000..04042dd3 --- /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 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..602d1184 --- /dev/null +++ b/src/tests/third_party/test_tensorflow_method.py @@ -0,0 +1,51 @@ +"""This is part of the MSS Python's module. +Source: https://github.com/BoboTiG/python-mss. +""" + +from __future__ import annotations + +import pytest + +from mss import ScreenShot +from tests.third_party.conftest import reordered_test_image + +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() + + +@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 new file mode 100644 index 00000000..5846c68e --- /dev/null +++ b/src/tests/third_party/test_torch_method.py @@ -0,0 +1,50 @@ +"""This is part of the MSS Python's module. +Source: https://github.com/BoboTiG/python-mss. +""" + +from __future__ import annotations + +import pytest + +from mss import ScreenShot +from tests.third_party.conftest import reordered_test_image + +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)) + + +@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)