diff --git a/Tests/test_image_encode.py b/Tests/test_image_encode.py new file mode 100644 index 00000000000..5ef9a3b853a --- /dev/null +++ b/Tests/test_image_encode.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import pytest + +from PIL import Image + +TYPE_CHECKING = False +if TYPE_CHECKING: + from typing import Any + + +@pytest.mark.parametrize( + "mode", + ("1", "L", "P", "LA", "I", "F", "I;16", "RGB", "RGBA", "CMYK"), +) +def test_raw_encoder_optimal_bufsize(mode: str) -> None: + # Test the raw encoder does know the exact buffer size. + im = Image.new(mode, (300, 200)) + encoder = Image._getencoder(mode, "raw", mode) + encoder.setimage(im.im, (0, 0) + im.size) + assert encoder.optimal_bufsize == len(im.tobytes()) + + +def test_raw_encoder_optimal_bufsize_stride() -> None: + im = Image.new("RGB", (100, 50)) + encoder = Image._getencoder("RGB", "raw", ("RGB", 400)) + encoder.setimage(im.im, (0, 0) + im.size) + assert encoder.optimal_bufsize == 400 * 50 + encoder.encode(400) # The encoder's internals change here + # ... but the result for this should remain the same. + assert encoder.optimal_bufsize == 400 * 50 + + +def test_encoder_optimal_bufsize_unknown() -> None: + im = Image.new("1", (64, 64)) + + # the size is not known before the image has been assigned + encoder = Image._getencoder("1", "raw", "1") + assert encoder.optimal_bufsize == 0 + + # nor for encoders whose output size depends on the pixel data + encoder = Image._getencoder("1", "xbm", "1") + encoder.setimage(im.im, (0, 0) + im.size) + assert encoder.optimal_bufsize == 0 + + +class EncoderSpy: + def __init__(self, encoder: Any, *, bufsizes: list[int]) -> None: + self._encoder = encoder + self.bufsizes = bufsizes + + def __getattr__(self, name: str) -> Any: + return getattr(self._encoder, name) + + def encode(self, bufsize: int) -> tuple[int, int, bytes]: + self.bufsizes.append(bufsize) + return self._encoder.encode(bufsize) + + +@pytest.mark.parametrize("mode", ("L", "RGB")) +@pytest.mark.parametrize( + "size", + ( + (20000, 1), + (1, 20000), + (3, 3), + (500, 500), + ), +) +def test_tobytes_exact_buffer_and_single_pass( + mode: str, + size: tuple[int, int], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """ + Test that the raw encoder gets asked exactly the correct size, + and ends up doing its work in a single pass. + """ + bufsizes: list[int] = [] + get_encoder = Image._getencoder + monkeypatch.setattr( + Image, + "_getencoder", + lambda *args: EncoderSpy(get_encoder(*args), bufsizes=bufsizes), + ) + + im = Image.new(mode, size) + n_bands = len(im.getbands()) + expected_size = size[0] * size[1] * n_bands + assert len(im.tobytes()) == expected_size + assert bufsizes == [expected_size] diff --git a/docs/releasenotes/13.0.0.rst b/docs/releasenotes/13.0.0.rst index d6d0ff6c40e..019e851b860 100644 --- a/docs/releasenotes/13.0.0.rst +++ b/docs/releasenotes/13.0.0.rst @@ -104,6 +104,33 @@ Two new filters are available for :py:meth:`~PIL.Image.Image.resize` and Other changes ============= +Reduced memory use in Image.tobytes() in raw mode +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Encoders may now report the "optimal" buffer size that they will need to encode the +entire image into their output format. + +In general, :py:meth:`~PIL.Image.Image.tobytes` collects encoded data in +``ImageFile.MAXBLOCK`` sized chunks and join them into the output byte buffer +afterwards. For raw uncompressed output, this would need, momentarily, roughly +twice the memory of the returned data. The raw encoder now reports the buffer size +that encodes the whole image in a single call, so no copy of the result is made. + +Encoders that cannot know their output size in advance, such as the compressed +encoders, are unaffected and keep using :py:data:`~PIL.ImageFile.MAXBLOCK`; +for many applications, it may be beneficial to adjust that value for improved +performance and reduced peak memory use. + +This also reduces peak memory use in other operations in Pillow that use this +path for encoding, including but not limited to: + +* AVIF and SGI encoding +* BLP JPEG decoding +* Image equivalence comparisons +* NumPy array interface support +* Pickling of images +* Qt interop + Python 3.15 ^^^^^^^^^^^ diff --git a/src/PIL/Image.py b/src/PIL/Image.py index c50cb777981..cd92b76aec3 100644 --- a/src/PIL/Image.py +++ b/src/PIL/Image.py @@ -886,7 +886,24 @@ def tobytes(self, encoder_name: str = "raw", *args: Any) -> bytes: from . import ImageFile - bufsize = max(ImageFile.MAXBLOCK, self.size[0] * 4) # see RawEncode.c + # Figure out an optimal buffer size. A buffer size that's too small + # to hold all of a given image's encoded bytes will result in `output` + # having a bunch of small chunks, and returning the actual bytes for + # them will then require a second contiguous allocation for the result, + # leading to double-plus-then-some peak memory allocation. + # + # CPython optimizes the case `b"".join([x])` for a simple `x` to be + # a reference increase of `x` and a return instead of any extra allocations; + # this has been the case since approximately 2004 (commit 05eba1fdc80). + # + # * e.optimal_bufsize will be nonzero for encoders that can compute it. + # * The heuristic `self.size[0] * 4` for the buffer size matches the widest + # possible row of pixels in a RGBA image; conveniently, this is also slightly + # larger than what is required for a single row in XbmEncode.c, + # which in itself is unable to cope with a too-small buffer. + # * MAXBLOCK can be tuned by the user; for the aforementioned reasons, + # we need to use the larger of that and the heuristic. + bufsize = e.optimal_bufsize or max(ImageFile.MAXBLOCK, self.size[0] * 4) output = [] while True: diff --git a/src/PIL/ImageFile.py b/src/PIL/ImageFile.py index 14038e4732d..f8e8466d4c4 100644 --- a/src/PIL/ImageFile.py +++ b/src/PIL/ImageFile.py @@ -892,6 +892,14 @@ class PyEncoder(PyCodec): def pushes_fd(self) -> bool: return self._pushes_fd + @property + def optimal_bufsize(self) -> int: + """ + The buffer size that would encode the whole tile in a single + :meth:`encode` call, or 0 if that is not known in advance. + """ + return 0 + def encode(self, bufsize: int) -> tuple[int, int, bytes]: """ Override to perform the encoding process. diff --git a/src/encode.c b/src/encode.c index bffc32938e8..d5b7f4eb917 100644 --- a/src/encode.c +++ b/src/encode.c @@ -43,6 +43,7 @@ typedef struct { Imaging im, ImagingCodecState state, UINT8 *buffer, int bytes ); int (*cleanup)(ImagingCodecState state); + Py_ssize_t (*optimal_bufsize)(ImagingCodecState state); struct ImagingCodecStateInstance state; Imaging im; PyObject *lock; @@ -83,8 +84,9 @@ PyImaging_EncoderNew(int contextsize) { /* Initialize encoder context */ encoder->state.context = context; - /* Most encoders don't need this */ + /* Most encoders don't need these */ encoder->cleanup = NULL; + encoder->optimal_bufsize = NULL; /* Target image */ encoder->lock = NULL; @@ -337,12 +339,31 @@ static struct PyMethodDef methods[] = { {NULL, NULL} /* sentinel */ }; +/* The buffer size that would let this encoder consume the whole tile in a + single call, or 0 if that size cannot be known up front (basically all + non-raw formats). Only meaningful once setimage() has been called. */ +static PyObject * +_get_optimal_bufsize(ImagingEncoderObject *encoder, void *closure) { + Py_ssize_t bufsize = 0; + + if (encoder->optimal_bufsize && encoder->im) { + bufsize = encoder->optimal_bufsize(&encoder->state); + } + + return PyLong_FromSsize_t(bufsize); +} + static struct PyGetSetDef getseters[] = { {"pushes_fd", (getter)_get_pushes_fd, NULL, "True if this decoder expects to push directly to self.fd", NULL}, + {"optimal_bufsize", + (getter)_get_optimal_bufsize, + NULL, + "Buffer size that encodes the entire tile in one call, or 0 if unknown", + NULL}, {NULL, NULL, NULL, NULL, NULL} /* sentinel */ }; @@ -497,6 +518,40 @@ PyImaging_PcxEncoderNew(PyObject *self, PyObject *args) { /* RAW */ /* -------------------------------------------------------------------- */ +static Py_ssize_t +_raw_optimal_bufsize(ImagingCodecState state) { + // ImagingRawEncode writes one row at a time into a buffer of at least `bytes`. + // Before the first `encode()` call, `bytes` is the packed row size, and `count` + // is the stride (or zero). After the first call, the two get swapped... + // The larger is the row size. + Py_ssize_t row = MAX(state->bytes, state->count); + Py_ssize_t bufsize = row * state->ysize; + // Cap the size to the whole number of rows that fits an `int`. + // TODO: this needs to be changed when encoders' buffer sizes become `ssize_t`, + // like decoders did in ca1cf5925. + if (bufsize > INT_MAX) { + bufsize = INT_MAX - (INT_MAX % row); + } + return bufsize; +} + +/** + * Instantiate a raw, uncompressed encoder. + * + * Python arguments: + * - mode (str): The mode name for finding a pixel packer. + * - rawmode (str): The rawmode name for finding a pixel packer. + * - stride (int, optional): The number of bytes each row occupies in the output stream, + * >= the packed pixel size. Each row is explicitly + * zero-padded to this size. If unset, the true packed size + * is used. + * - ystep (int, optional): If set to a negative value, the encoder will write rows in + * reverse order. The only effective values are -1 and +1; + * this does not have the encoder skip rows. + * @param self Unused. + * @param args Python arguments, see above. + * @return A Python object representing an encoder. + */ PyObject * PyImaging_RawEncoderNew(PyObject *self, PyObject *args) { ImagingEncoderObject *encoder; @@ -523,6 +578,7 @@ PyImaging_RawEncoderNew(PyObject *self, PyObject *args) { } encoder->encode = ImagingRawEncode; + encoder->optimal_bufsize = _raw_optimal_bufsize; encoder->state.ystep = ystep; encoder->state.count = stride;