Add convenience export methods to ScreenShot#14
Closed
jholveck wants to merge 18 commits into
Closed
Conversation
Currently, this shares memory with the .raw bytearray in some cases. We should consider whether we want to change that to always copy, or just wait until we have the full shared buffer support. This implementation always does the array manipulation on the CPU. For PyTorch or TensorFlow, the user is likely to want to do the array manipulation on the GPU for performance. This is very easy to implement.
Python 3.9 doesn't define TypeAlias, and it's not required in this case, so we can just remove it.
The reverted change was from a temporary test, but I accidentally included it in the commit. It's not relevant to the issue at hand.
The docstrings weren't all in Sphinx format, and could use a little clarification in places. TensorFlow itself allows NumPy dtypes and DataType enums, so expand to_tensorflow to accept those too.
The main point of this commit is to be explicit about the memory-sharing semantics: mostly, that we don't guarantee memory sharing one way or another. More to the point, that if they modify the pixel data in one of the exported objects, it may or may not affect the other exported objects. Memory sharing isn't expected to affect most users, since they are unlikely to access the pixel data with two different exported objects, but it's best to be explicit. To have a good place to put the memory-sharing discussion, I added a new section to the usage.rst file, listing the ways to access the pixel data. This is also where I documented the NumPy array interface protocol, which previously didn't show up in our Sphinx docs. (It might have been in the docs for 10.1, which didn't use autodoc.) I also put a discussion of the alpha channel in this new section: one thing I keep forgetting is to delete the alpha channel when using Matplotlib, and wondering why it's blank. TODO: We probably should document `grab` in the usage.rst file for the new section to make sense. In the process, I've also added intersphinx links to libraries like PIL and NumPy, so that our Sphinx docs will link to them. While checking that, I fixed a few other references (like referring to bytes with :py:type:, when it's indexed under :py:class:).
Doc fixes. Speed up converting to a NumPy array in RGB order. That sort of thing.
While I was writing this, I realized something else. We may want to change the to_torch and to_tensorflow methods to accept device and stream arguments, and do the channel and layout shuffling on the GPU instead of the CPU. This wouldn't require us to natively support CUDA or anything like that; we'd just be using PyTorch / TensorFlow's existing CUDA support, which is easy. This demo still (as before) shows the GPU-focused way to do the channel and layout shuffling, presented as an alternative to using to_torch. A quick test shows about 107 ms / frame (to_torch) vs. about 62 ms / frame (GPU-focused routine), and that's including the CNN itself too. So there's a very significant benefit there. (That's not quite a fair comparison, since the GPU-focused routine is exactly what this program needs, but not so it'd be a big difference.)
While improving performance earlier in this PR, I had introduced a case where to_torch would fail only for the RGB order. It also seemed likely that some cases might fail only if using particular combinations of channel order and layout, so we just test them all.
This adds a device= flag to to_torch, so the user can specify the destination device. (TensorFlow has more automatic device management, and we don't need to do anything explicit.) We also do all the channel shuffling to the GPU. For cases when channel shuffling or dtype conversion requires a copy, this is immensely faster than the alternative. I ran a rough test with the cat detector on my 4k monitor (for everything, including the CNN inference). On a GPU, there's a significant performance boost. However, this is now slower when run on the CPU. * Old code, with GPU: 99 ms/frame * New code, with GPU: 56 ms/frame * Old code, no GPU: 3164 ms/frame * New code, no GPU: 3219 ms/frame * New code, running to_torch(device="cpu") and and then transferring to the GPU for inference: 117 ms/frame It seems the new to_torch is slower when running on the CPU. I suspect (but this is a loose guess) that NumPy, which always did the channel shuffles in the old code, was faster at that. It might also be specific to the extra copy that I think this change introduces in the RGB CHW conversion path in this, due to PyTorch limitations. I don't think this is a significant problem: PyTorch and TensorFlow users are almost all going to be using GPUs. I will still look at ways to improve the CPU performance on PyTorch, though.
📝 WalkthroughWalkthrough
ChangesScreenShot Framework Conversion API
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Owner
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (2)
src/tests/third_party/test_numpy_method.py (1)
61-67: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd one
dtypecase forto_numpy().This suite never exercises the new cast/scale branch in
src/mss/screenshot.pyLines 266-270, so regressions indtype=np.float32normalization would slip through even though the Torch and TensorFlow variants are covered.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/tests/third_party/test_numpy_method.py` around lines 61 - 67, Add a test case in test_to_numpy_permutations or a nearby dedicated test that calls ScreenShot.to_numpy with dtype=np.float32 so the cast/scale branch in screenshot conversion is exercised. Use the existing test helpers like framework_test_image and reordered_test_image, and assert the float32 output matches the expected normalized values for one representative channels/layout combination. Keep the current permutation coverage intact and extend it with one dtype assertion to catch regressions in to_numpy.docs/source/conf.py (1)
99-103: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winPin the TensorFlow inventory to an immutable source.
This mapping now follows a community
raw/masterURL, so doc builds can change or break when that branch moves even if this repo does not. Please pin it to a specific commit/tag or vendor the.invfile into the docs tree.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/source/conf.py` around lines 99 - 103, The TensorFlow InterSphinx mapping in conf.py is using a mutable raw/master inventory URL, which can make docs builds unstable. Update the tensorflow entry in the intersphinx mapping to use an immutable source such as a pinned commit/tag URL, or vendor the tf2_py_objects.inv file into the docs tree and reference that local copy.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/source/examples/pil_pixels.py`:
- Line 14: The example in the pil_pixels flow is mixing two conversion paths:
img is already created via sct_img.to_pil("RGB"), but the later manual
rebuilding from sct_img.bgra duplicates the conversion and makes the example
unclear. Update the example to use only one path by either keeping the manual
Image.new plus putdata approach throughout or removing the manual pixel
reconstruction and relying on the existing img variable; check the pil_pixels
example flow around the sct_img.to_pil and Image.new/putdata steps.
In `@docs/source/release-history/v11.0.0.md`:
- Around line 22-23: Fix the typo in the release note text by updating the word
“noticible” to “noticeable” in the v11.0.0 documentation entry. The change is in
the release-history markdown content describing memoryview screenshot data, so
just correct that misspelling in the relevant sentence.
In `@docs/source/usage.rst`:
- Around line 109-110: The OpenCV channel-order note is incorrect in the usage
docs and conflicts with the example and surrounding text. Update the wording in
the relevant docs paragraph to refer to BGR/BGRA instead of RGB, and if
appropriate point readers to mss.ScreenShot.to_numpy with channels="BGR" for
OpenCV use.
In `@src/tests/third_party/conftest.py`:
- Line 14: The conftest-level numpy import guard is skipping the entire
src/tests/third_party subtree, including PIL-only tests that should still run
without NumPy. Remove the pytest.importorskip("numpy") from conftest.py and move
the guard into the specific NumPy-dependent helpers or test modules that
actually require np, so PIL-only files like test_pil.py and test_pil_method.py
remain discoverable.
In `@src/tests/third_party/test_torch_method.py`:
- Around line 20-24: The test in shot.to_torch() is hardcoding torch.float32
even though the method defaults to torch.get_default_dtype() when dtype is
omitted. Update the assertion to compare against the current default dtype, or
make the call explicit by passing dtype=torch.float32, so the expectation stays
aligned with the behavior of to_torch().
---
Nitpick comments:
In `@docs/source/conf.py`:
- Around line 99-103: The TensorFlow InterSphinx mapping in conf.py is using a
mutable raw/master inventory URL, which can make docs builds unstable. Update
the tensorflow entry in the intersphinx mapping to use an immutable source such
as a pinned commit/tag URL, or vendor the tf2_py_objects.inv file into the docs
tree and reference that local copy.
In `@src/tests/third_party/test_numpy_method.py`:
- Around line 61-67: Add a test case in test_to_numpy_permutations or a nearby
dedicated test that calls ScreenShot.to_numpy with dtype=np.float32 so the
cast/scale branch in screenshot conversion is exercised. Use the existing test
helpers like framework_test_image and reordered_test_image, and assert the
float32 output matches the expected normalized values for one representative
channels/layout combination. Keep the current permutation coverage intact and
extend it with one dtype assertion to catch regressions in to_numpy.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 7720f49c-2ffa-4462-bfbd-876a9682742a
📒 Files selected for processing (21)
.github/workflows/tests.ymldemos/cat-detector.pydemos/tinytv-stream-simple.pydemos/tinytv-stream.pydemos/video-capture-simple.pydemos/video-capture.pydocs/source/conf.pydocs/source/examples.rstdocs/source/examples/opencv_numpy.pydocs/source/examples/pil.pydocs/source/examples/pil_pixels.pydocs/source/release-history/v11.0.0.mddocs/source/usage.rstsrc/mss/base.pysrc/mss/screenshot.pysrc/tests/test_setup.pysrc/tests/third_party/conftest.pysrc/tests/third_party/test_numpy_method.pysrc/tests/third_party/test_pil_method.pysrc/tests/third_party/test_tensorflow_method.pysrc/tests/third_party/test_torch_method.py
|
|
||
| # Create an Image | ||
| img = Image.new("RGB", sct_img.size) | ||
| img = sct_img.to_pil("RGB") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Keep this example on one conversion path.
img is already populated here, but the example still starts rebuilding RGB tuples from sct_img.bgra at Line 17. That makes the script do the conversion twice and muddies the point of a “pixels” example. Either keep the manual Image.new(...)+putdata(...) flow, or drop the manual pixel reconstruction below.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/source/examples/pil_pixels.py` at line 14, The example in the pil_pixels
flow is mixing two conversion paths: img is already created via
sct_img.to_pil("RGB"), but the later manual rebuilding from sct_img.bgra
duplicates the conversion and makes the example unclear. Update the example to
use only one path by either keeping the manual Image.new plus putdata approach
throughout or removing the manual pixel reconstruction and relying on the
existing img variable; check the pil_pixels example flow around the
sct_img.to_pil and Image.new/putdata steps.
Comment on lines
+22
to
+23
| {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. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix the typo in this release note.
noticible should be noticeable.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/source/release-history/v11.0.0.md` around lines 22 - 23, Fix the typo in
the release note text by updating the word “noticible” to “noticeable” in the
v11.0.0 documentation entry. The change is in the release-history markdown
content describing memoryview screenshot data, so just correct that misspelling
in the relevant sentence.
Comment on lines
+109
to
+110
| 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. |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== docs/source/usage.rst (around lines 90-125) ==\n'
sed -n '90,125p' docs/source/usage.rst | cat -n
printf '\n== docs/source/examples.rst (around lines 140-165) ==\n'
sed -n '140,165p' docs/source/examples.rst | cat -n
printf '\n== Search for to_numpy and channels docs ==\n'
rg -n "to_numpy|channels=\"BGR\"|OpenCV|BGRA|RGB order|BGR order" docs/source mss -SRepository: jholveck/python-mss
Length of output: 6047
Fix the OpenCV channel-order note.
docs/source/usage.rst says OpenCV uses RGB order, but the rest of the docs and the OpenCV example use BGR/BGRA. Update this to BGR/BGRA, or point readers to to_numpy(channels="BGR") for OpenCV.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/source/usage.rst` around lines 109 - 110, The OpenCV channel-order note
is incorrect in the usage docs and conflicts with the example and surrounding
text. Update the wording in the relevant docs paragraph to refer to BGR/BGRA
instead of RGB, and if appropriate point readers to mss.ScreenShot.to_numpy with
channels="BGR" for OpenCV use.
| # importorskip), but type checkers don't always know about importorskip. | ||
| import numpy as np_typehints # noqa: ICN001 | ||
|
|
||
| np = pytest.importorskip("numpy") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
fd -t f . src/tests/third_party | sort
echo
rg -n "importorskip\\(|reordered_test_image|framework_test_image" src/tests/third_partyRepository: jholveck/python-mss
Length of output: 3405
Don't gate src/tests/third_party on NumPy. pytest.importorskip("numpy") in conftest.py skips collection for the whole subtree, so PIL-only tests like test_pil.py and test_pil_method.py disappear in NumPy-free environments. Move the guard into the NumPy-dependent helpers/modules instead.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/tests/third_party/conftest.py` at line 14, The conftest-level numpy
import guard is skipping the entire src/tests/third_party subtree, including
PIL-only tests that should still run without NumPy. Remove the
pytest.importorskip("numpy") from conftest.py and move the guard into the
specific NumPy-dependent helpers or test modules that actually require np, so
PIL-only files like test_pil.py and test_pil_method.py remain discoverable.
Comment on lines
+20
to
+24
| 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) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n "set_default_dtype|get_default_dtype" src/tests src || trueRepository: jholveck/python-mss
Length of output: 309
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== screenshot implementation =="
sed -n '270,325p' src/mss/screenshot.py
echo
echo "== torch method test =="
sed -n '1,120p' src/tests/third_party/test_torch_method.pyRepository: jholveck/python-mss
Length of output: 5031
Assert against the current default dtype here. to_torch() uses torch.get_default_dtype() when dtype is omitted, so hardcoding torch.float32 makes this test fail if PyTorch’s global default is changed elsewhere. Use torch.get_default_dtype() in the expectation, or pass dtype=torch.float32 explicitly.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/tests/third_party/test_torch_method.py` around lines 20 - 24, The test in
shot.to_torch() is hardcoding torch.float32 even though the method defaults to
torch.get_default_dtype() when dtype is omitted. Update the assertion to compare
against the current default dtype, or make the call explicit by passing
dtype=torch.float32, so the expectation stays aligned with the behavior of
to_torch().
There was a problem hiding this comment.
Pull request overview
This PR adds new convenience export methods on mss.ScreenShot to convert screenshots efficiently into common image/ML framework formats (Pillow, NumPy, PyTorch, TensorFlow), along with updated documentation, examples/demos, CI changes, and new third-party tests to validate channel/layout/dtype permutations.
Changes:
- Add
ScreenShot.to_pil(),to_numpy(),to_torch(), andto_tensorflow()with channel/layout/dtype options and related typing. - Add third-party test suite for the new export methods (skipping when optional deps aren’t installed) and include them in sdist checks.
- Update docs/examples/demos and CI to exercise the new pathways and improve framework references.
Reviewed changes
Copilot reviewed 21 out of 21 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
src/mss/screenshot.py |
Adds new export methods for PIL/NumPy/Torch/TensorFlow and related types/docs. |
src/mss/base.py |
Docstring/reference formatting updates for public API docs. |
docs/source/usage.rst |
Adds usage documentation for capture/pixel access/export/memory sharing. |
docs/source/release-history/v11.0.0.md |
Minor documentation formatting/wording adjustments in release notes. |
docs/source/examples/pil.py |
Updates example to use ScreenShot.to_pil(). |
docs/source/examples/pil_pixels.py |
Updates pixel-edit example to use the new PIL conversion path. |
docs/source/examples/opencv_numpy.py |
Updates OpenCV example to use ScreenShot.to_numpy(channels="BGR"). |
docs/source/examples.rst |
Updates examples index text and adds guidance around to_numpy() channel choice. |
docs/source/conf.py |
Extends Sphinx intersphinx mappings (numpy/pillow/tf/torch). |
demos/video-capture.py |
Uses to_numpy() in the demo and updates explanatory comments. |
demos/video-capture-simple.py |
Uses to_pil() in the simple video capture demo and trims whitespace. |
demos/tinytv-stream.py |
Uses to_pil() in the TinyTV stream demo. |
demos/tinytv-stream-simple.py |
Uses to_pil() in the TinyTV simple demo. |
demos/cat-detector.py |
Switches the demo to use to_torch() and removes an older helper path. |
.github/workflows/tests.yml |
Installs optional ML deps in CI and adjusts OS selection logic. |
src/tests/test_setup.py |
Ensures new third-party tests are included in sdist contents. |
src/tests/third_party/conftest.py |
Adds shared fixtures/helpers for framework export tests. |
src/tests/third_party/test_numpy_method.py |
Adds tests for ScreenShot.to_numpy() (channels/layout/error cases). |
src/tests/third_party/test_pil_method.py |
Adds tests for ScreenShot.to_pil() (RGB/RGBA and validation). |
src/tests/third_party/test_tensorflow_method.py |
Adds tests for ScreenShot.to_tensorflow() (dtype/layout/channel permutations). |
src/tests/third_party/test_torch_method.py |
Adds tests for ScreenShot.to_torch() (dtype/layout/channel permutations + optional CUDA). |
Comments suppressed due to low confidence (1)
docs/source/examples/pil_pixels.py:18
- This example now creates the image via
ScreenShot.to_pil("RGB")and then immediately re-populates it viaputdata(). That makes the example slower and more confusing (it no longer demonstrates building an image from pixel data). Either keepImage.new(...)+putdata(...), or keepto_pil()and drop theputdata()path.
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)
img.putdata(list(pixels))
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+97
to
+102
| - 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 |
| 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) |
Comment on lines
+109
to
+110
| 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. |
| 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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add convenience methods to mss.ScreenShot to export to PIL, NumPy, PyTorch, and TensorFlow. These are high-performance methods that can be used by end users for most common export tasks.
This also removes the burden from the user of figuring out the best way to, for instance, convert from BGRA to RGB: by putting the fastest techniques in these methods, we let users focus on their applications, and even get benefits as we make changes to things like buffer management that allow faster exports.
Extensive documentation is added, since these may end up being the most common way for end users to access pixel data.
This is also the first time we explicitly add any GPU support to MSS. This is a pretty small thing in MSS’s terms; we still are copying the framebuffer to CPU memory. But we’re now enabling fast, easy export to the common DL frameworks, which is where we’re seeing a rise in screen image usage lately.
These methods are tested in CI/CD with CPU only. On developer machines, they’re tested if the dev has installed PyTorch or TensorFlow (NumPy is already a prerequisite), including the CUDA paths if that’s available.
Summary by CodeRabbit
New Features
Documentation
Bug Fixes