diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 3baeb668..936cdd80 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,32 @@ 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 == '🐧' + timeout-minutes: 10 + if: matrix.os.name == 'linux' run: xvfb-run python -m pytest - name: Tests (macOS, Windows) - if: matrix.os.emoji != '🐧' + timeout-minutes: 10 + if: matrix.os.name != 'linux' run: python -m pytest automerge: diff --git a/demos/cat-detector.py b/demos/cat-detector.py index cde7d833..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,27 +258,8 @@ 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. + 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/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..e84a89f2 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,20 +135,23 @@ 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- .. versionadded:: 3.0.0 -OpenCV/Numpy +OpenCV/NumPy ============ 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..820bf4a5 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,16 @@ 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. + # 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', - # 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/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 d4eef6b5..c39c5e6e 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 noticeable. This change allows faster access to screenshot data, with fewer memory copies. ### Python 3.9 EOL @@ -48,13 +47,16 @@ processing time decreased from 22.64 ms to 18.59 ms per frame (approximately 18% Support for additional operating systems is planned. -### General Improvements +### Export Methods -The MSS context object will now always surface inner exceptions, even if `__exit__` may also generate an exception during tear-down. +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. -### Documentation +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. -The documentation has received numerous small improvements. A few highlights: +See the documentation section {ref}`accessing_pixel_data` for details. -* The Pillow examples and demos using {py:meth}`PIL.Image.frombuffer` now explicitly specify the decoder arguments, - as recommended by Pillow. (#535) +### 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 08acdf8c..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 ====== @@ -32,6 +36,141 @@ 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. + +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 +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.primary_monitor) + +Ok, now you've got the :py:class:`mss.ScreenShot` object. But what do +you do with it? + +.. _accessing_pixel_data: + +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, 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, since +it can return the pixel data in most common layouts, orders, and data types. + +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 +280,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..723cd7a0 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,243 @@ 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", 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"``. + :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. + + When requesting ``"RGBA"`` or ``"BGRA"``, the alpha channel may + not represent meaningful transparency on all platforms/backends. + + .. version-added:: 11.0.0 + """ + channels = cast("Channels", channels.upper()) + layout = cast("Layout", layout.upper()) + + import numpy as np # noqa: PLC0415 + + rv = np.frombuffer(self._raw, dtype=np.uint8).reshape((self.height, self.width, 4)) + + if channels == "BGRA": + pass + elif channels == "BGR": + rv = rv[:, :, :3] + elif channels == "RGB": + # Using [2,1,0] instead of 2::-1 would copy, rather than creating a view. + 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. + 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": + # This will always create a view. (We're reordering the axes, not the elements.) + 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 rv + + def to_torch( # noqa: PLR0912 + 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. + + :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`. + :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. + + 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 + """ + channels = cast("Channels", channels.upper()) + layout = cast("Layout", layout.upper()) + + import torch # noqa: PLC0415 + + if dtype is None: + 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) + + 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). + 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. + # 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 + 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) + + return rv + + def to_tensorflow( + self, + channels: Channels = "RGB", + layout: Layout = "HWC", + dtype: tf.dtypes.DType | numpy.dtype | 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 + `, or a :py:class:`np.dtype `. + + Device and stream management is handled by TensorFlow. + + 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 + """ + channels = cast("Channels", channels.upper()) + layout = cast("Layout", layout.upper()) + + import tensorflow as tf # noqa: PLC0415 + + # TypeErrors from tf.as_dtype are passed up to the caller. + tf_dtype = tf.as_dtype(dtype) + + # 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) + + # 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) + + 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: + rv /= 255.0 + + return rv + @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..f77f6275 100644 --- a/src/tests/test_setup.py +++ b/src/tests/test_setup.py @@ -113,8 +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/test_numpy.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/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/array_frameworks/conftest.py b/src/tests/third_party/array_frameworks/conftest.py new file mode 100644 index 00000000..565bc0f6 --- /dev/null +++ b/src/tests/third_party/array_frameworks/conftest.py @@ -0,0 +1,58 @@ +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 + +# This entire tree is skipped if NumPy isn't installed. +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.py b/src/tests/third_party/array_frameworks/test_numpy.py similarity index 80% rename from src/tests/third_party/test_numpy.py rename to src/tests/third_party/array_frameworks/test_numpy.py index c06e6180..0bf39413 100644 --- a/src/tests/third_party/test_numpy.py +++ b/src/tests/third_party/array_frameworks/test_numpy.py @@ -4,12 +4,10 @@ from collections.abc import Callable -import pytest +import numpy as np from mss import MSS -np = pytest.importorskip("numpy", reason="Numpy module not available.") - def test_numpy(mss_impl: Callable[..., MSS]) -> None: box = {"top": 0, "left": 0, "width": 10, "height": 10} diff --git a/src/tests/third_party/array_frameworks/test_numpy_method.py b/src/tests/third_party/array_frameworks/test_numpy_method.py new file mode 100644 index 00000000..8acab8be --- /dev/null +++ b/src/tests/third_party/array_frameworks/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.array_frameworks.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/array_frameworks/test_tensorflow_method.py b/src/tests/third_party/array_frameworks/test_tensorflow_method.py new file mode 100644 index 00000000..3dd8c7df --- /dev/null +++ b/src/tests/third_party/array_frameworks/test_tensorflow_method.py @@ -0,0 +1,53 @@ +"""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.array_frameworks.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 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 new file mode 100644 index 00000000..7226b1e2 --- /dev/null +++ b/src/tests/third_party/array_frameworks/test_torch_method.py @@ -0,0 +1,68 @@ +"""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.array_frameworks.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.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) + + +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"]) +@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("CUDA is not available") + + 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, # type: ignore[arg-type] + layout=layout, # type: ignore[arg-type] + 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( + channels=channels, # type: ignore[arg-type] + layout=layout, # type: ignore[arg-type] + 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_f32 = bfloat16_result.to(torch.float32) + assert np.allclose(bfloat16_result_f32.cpu().numpy(), bfloat16_target, rtol=0, atol=1 / 512.0) 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")