perf: Free image buffers sooner during figure extraction

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.27.23
+
+### Enhancements
+
+- **Release image extraction buffers after each crop.** Image and PDF figure extraction closes owned decoded pages, crops, and JPEG buffers promptly, and copies image uploads to its existing temporary file in bounded chunks. JPEG payloads, saved files, metadata, and caller-owned streams retain their existing behavior.
+
 ## 0.27.19-dev0
 
 ### Maintenance

scripts/performance/benchmark_image_extraction_memory.py

@@ -0,0 +1,98 @@
+"""Measure peak memory of PIL image extraction and hash its output for cross-checkout comparison.
+
+Run with PYTHONPATH pointing at the desired library checkout. The deterministic
+100-megapixel BMP is written a row at a time, so fixture generation is bounded.
+"""
+
+import argparse
+import hashlib
+import json
+import random
+import resource
+import struct
+import sys
+import tempfile
+import time
+from pathlib import Path
+
+from unstructured.documents.coordinates import PixelSpace
+from unstructured.documents.elements import ElementMetadata, Image
+from unstructured.partition.pdf_image.pdf_image_utils import save_elements
+
+
+def _positive_int(value: str) -> int:
+    """Accept only positive dimensions and extraction counts."""
+    result = int(value)
+    if result <= 0:
+        raise argparse.ArgumentTypeError("must be greater than zero")
+    return result
+
+
+def main() -> None:
+    """Extract full-page figures from one high-resolution raster upload."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--side", type=_positive_int, default=10000)
+    parser.add_argument("--count", type=_positive_int, default=2)
+    args = parser.parse_args()
+    side = args.side
+    stride = (side * 3 + 3) // 4 * 4
+    size = stride * side
+    with tempfile.TemporaryDirectory() as directory:
+        path = Path(directory) / "scan.bmp"
+        rng = random.Random(901)
+        with path.open("wb") as output:
+            output.write(struct.pack("<2sIHHI", b"BM", 54 + size, 0, 0, 54))
+            output.write(struct.pack("<IiiHHIIiiII", 40, side, side, 1, 24, 0, size, 0, 0, 0, 0))
+            for _ in range(side):
+                output.write(rng.randbytes(side * 3))
+                output.write(b"\0" * (stride - side * 3))
+        elements = [
+            Image(
+                text=f"figure {index}",
+                element_id=str(index),
+                coordinates=((0, 0), (0, side), (side, side), (side, 0)),
+                coordinate_system=PixelSpace(width=side, height=side),
+                metadata=ElementMetadata(page_number=1),
+            )
+            for index in range(args.count)
+        ]
+        start = time.perf_counter()
+        with path.open("rb") as source:
+            save_elements(
+                elements,
+                1,
+                "Image",
+                200,
+                file=source,
+                is_image=True,
+                extract_image_block_to_payload=True,
+            )
+            assert not source.closed
+            assert source.tell() == path.stat().st_size
+        seconds = time.perf_counter() - start
+        # Hash each field without constructing another complete output JSON.
+        digest = hashlib.sha256()
+        for element in elements:
+            metadata = element.to_dict()
+            payload = metadata["metadata"].pop("image_base64")
+            digest.update(json.dumps(metadata, sort_keys=True).encode())
+            for offset in range(0, len(payload), 65536):
+                digest.update(payload[offset : offset + 65536].encode())
+        peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+        print(
+            json.dumps(
+                {
+                    "side": side,
+                    "count": len(elements),
+                    "seconds": seconds,
+                    "peak_rss_kib": peak // 1024 if sys.platform == "darwin" else peak,
+                    "output_sha256": digest.hexdigest(),
+                    "payload_bytes_each": len(elements[0].metadata.image_base64 or ""),
+                }
+            ),
+            flush=True,
+        )
+
+
+if __name__ == "__main__":
+    main()

test_unstructured/partition/pdf_image/test_extracted_image_memory.py

@@ -0,0 +1,148 @@
+"""Extraction byte parity, bounded upload reads, and owned image lifetime regressions."""
+
+import base64
+import io
+
+import pytest
+from PIL import Image as PILImage
+
+from unstructured.documents.coordinates import PixelSpace
+from unstructured.documents.elements import ElementMetadata, Image
+from unstructured.partition.pdf_image import pdf_image_utils
+
+
+def _closed(image):
+    try:
+        image.getpixel((0, 0))
+    except ValueError:
+        return True
+    return False
+
+
+def _element():
+    return Image(
+        text="figure",
+        coordinates=((-3, 2), (-3, 35), (45, 35), (45, 2)),
+        coordinate_system=PixelSpace(width=50, height=40),
+        metadata=ElementMetadata(page_number=7),
+    )
+
+
+@pytest.mark.parametrize("mode", ["RGB", "RGBA", "L"])
+@pytest.mark.parametrize("payload", [True, False])
+def test_extraction_preserves_native_jpeg_bytes_and_metadata(tmp_path, mode, payload):
+    source = tmp_path / "source.png"
+    with PILImage.new(mode, (50, 40), 123) as image:
+        image.save(source)
+    with PILImage.open(source) as image:
+        crop = image.crop((-3, 2, 45, 35))
+        if mode == "RGBA":
+            crop = crop.convert("RGB")
+        with crop, io.BytesIO() as expected:
+            crop.save(expected, format="JPEG")
+            expected_bytes = expected.getvalue()
+    element = _element()
+    pdf_image_utils.save_elements(
+        [element],
+        7,
+        "Image",
+        200,
+        filename=str(source),
+        is_image=True,
+        extract_image_block_to_payload=payload,
+        output_dir_path=str(tmp_path),
+    )
+    assert element.text == "figure"
+    assert element.metadata.page_number == 7
+    if payload:
+        assert element.metadata.image_base64 == base64.b64encode(expected_bytes).decode()
+        assert element.metadata.image_mime_type == "image/jpeg"
+        assert element.metadata.image_path is None
+    else:
+        path = tmp_path / "figure-7-1.jpg"
+        assert path.read_bytes() == expected_bytes
+        assert element.metadata.image_path == str(path)
+        assert element.metadata.image_base64 is None
+
+
+def test_image_upload_reads_are_bounded_and_leave_borrowed_stream_open(tmp_path):
+    class BoundedStream(io.BytesIO):
+        def read(self, size=-1):
+            assert 0 < size <= 1024 * 1024
+            return super().read(size)
+
+    with PILImage.new("RGB", (50, 40)) as image, io.BytesIO() as encoded:
+        image.save(encoded, format="PNG")
+        raw = encoded.getvalue()
+    stream = BoundedStream(raw)
+    stream.seek(3)
+    element = _element()
+    pdf_image_utils.save_elements(
+        [element],
+        7,
+        "Image",
+        200,
+        file=stream,
+        is_image=True,
+        extract_image_block_to_payload=True,
+    )
+    assert not stream.closed
+    assert stream.tell() == len(raw)
+    stream.seek(0)
+    assert stream.read(1024 * 1024) == raw
+
+
+@pytest.mark.parametrize("fail_encode", [False, True])
+def test_extraction_releases_page_before_encoding_and_crops_before_next_decode(
+    tmp_path,
+    monkeypatch,
+    fail_encode,
+):
+    source = tmp_path / "source.png"
+    with PILImage.new("RGBA", (50, 40)) as image:
+        image.save(source)
+    opened = []
+    crops = []
+    real_open = PILImage.open
+    real_crop = PILImage.Image.crop
+    real_save = PILImage.Image.save
+    real_convert = PILImage.Image.convert
+
+    def tracked_open(*args, **kwargs):
+        assert all(_closed(image) for image in crops)
+        image = real_open(*args, **kwargs)
+        opened.append(image)
+        return image
+
+    def tracked_crop(image, *args, **kwargs):
+        crop = real_crop(image, *args, **kwargs)
+        crops.append(crop)
+        return crop
+
+    def tracked_convert(image, *args, **kwargs):
+        converted = real_convert(image, *args, **kwargs)
+        crops.append(converted)
+        return converted
+
+    def tracked_save(image, *args, **kwargs):
+        assert all(_closed(page) for page in opened)
+        assert _closed(crops[-2])  # RGBA crop is released after RGB conversion.
+        if fail_encode:
+            raise OSError("failed JPEG encoder")
+        return real_save(image, *args, **kwargs)
+
+    monkeypatch.setattr(PILImage, "open", tracked_open)
+    monkeypatch.setattr(PILImage.Image, "crop", tracked_crop)
+    monkeypatch.setattr(PILImage.Image, "convert", tracked_convert)
+    monkeypatch.setattr(PILImage.Image, "save", tracked_save)
+    pdf_image_utils.save_elements(
+        [_element(), _element()],
+        7,
+        "Image",
+        200,
+        filename=str(source),
+        is_image=True,
+        extract_image_block_to_payload=True,
+    )
+    assert len(opened) == 2
+    assert all(_closed(image) for image in opened + crops)

unstructured/__version__.py

@@ -1 +1 @@
-__version__ = "0.27.19-dev0"  # pragma: no cover
+__version__ = "0.27.23"  # pragma: no cover

unstructured/partition/pdf_image/pdf_image_utils.py

@@ -3,6 +3,7 @@
 import base64
 import os
 import re
+import shutil
 import tempfile
 import unicodedata
 from copy import deepcopy
@@ -150,15 +151,13 @@ def save_elements(
             if file is None:
                 image_paths = [filename]
             else:
-                if isinstance(file, bytes):
-                    file_data = file
-                else:
-                    file.seek(0)
-                    file_data = file.read()
-
                 tmp_file_path = os.path.join(temp_dir, "tmp_file")
                 with open(tmp_file_path, "wb") as tmp_file:
-                    tmp_file.write(file_data)
+                    if isinstance(file, bytes):
+                        tmp_file.write(file)
+                    else:
+                        file.seek(0)
+                        shutil.copyfileobj(file, tmp_file, length=1024 * 1024)
                 image_paths = [tmp_file_path]
         else:
             _image_paths = convert_pdf_to_image(
@@ -206,30 +205,38 @@ def save_elements(
             figure_number += 1
             try:
                 image_path = image_paths[page_index]
+                # The crop owns its pixels; release the decoded page before encoding it.
+                # Image.close() frees the pixel data, which the context manager does not.
                 image = Image.open(image_path)
-                cropped_image = image.crop(padded_bbox)
-
-                # PNG images with transparency need to be converted before saving
-                if cropped_image.mode == "RGBA":
-                    cropped_image = cropped_image.convert("RGB")
-
-                if extract_image_block_to_payload:
-                    buffered = BytesIO()
-                    cropped_image.save(buffered, format="JPEG")
-                    img_base64 = base64.b64encode(buffered.getvalue())
-                    img_base64_str = img_base64.decode()
-                    el.metadata.image_base64 = img_base64_str
-                    el.metadata.image_mime_type = "image/jpeg"
-                else:
-                    basename = "table" if el.category == ElementType.TABLE else "figure"
-                    assert output_dir_path
-                    output_f_path = os.path.join(
-                        output_dir_path,
-                        f"{basename}-{metadata_page_number}-{figure_number}.jpg",
-                    )
-                    write_image(cropped_image, output_f_path)
-                    # add image path to element metadata
-                    el.metadata.image_path = output_f_path
+                try:
+                    cropped_image = image.crop(padded_bbox)
+                finally:
+                    image.close()
+                try:
+                    # PNG images with transparency need conversion before JPEG encoding.
+                    if cropped_image.mode == "RGBA":
+                        converted_image = cropped_image.convert("RGB")
+                        cropped_image.close()
+                        cropped_image = converted_image
+
+                    if extract_image_block_to_payload:
+                        with BytesIO() as buffered:
+                            cropped_image.save(buffered, format="JPEG")
+                            el.metadata.image_base64 = base64.b64encode(
+                                buffered.getvalue()
+                            ).decode()
+                        el.metadata.image_mime_type = "image/jpeg"
+                    else:
+                        basename = "table" if el.category == ElementType.TABLE else "figure"
+                        assert output_dir_path
+                        output_f_path = os.path.join(
+                            output_dir_path,
+                            f"{basename}-{metadata_page_number}-{figure_number}.jpg",
+                        )
+                        write_image(cropped_image, output_f_path)
+                        el.metadata.image_path = output_f_path
+                finally:
+                    cropped_image.close()
             except (ValueError, IOError):
                 logger.warning("Image Extraction Error: Skipping the failed image", exc_info=True)