fix: release image backend resources after frame extraction (#3134)

Signed-off-by: LarytheLord <llawlietbagsum@gmail.com>

Author: LarytheLord

docling/backend/image_backend.py

@@ -1,7 +1,7 @@
 import logging
 from io import BytesIO
 from pathlib import Path
-from typing import Iterable, List, Optional, Union
+from typing import Iterable, List, Union
 
 from docling_core.types.doc import BoundingBox, CoordOrigin
 from docling_core.types.doc.page import (
@@ -24,7 +24,7 @@
 
 class _ImagePageBackend(PdfPageBackend):
     def __init__(self, image: Image.Image):
-        self._image: Optional[Image.Image] = image
+        self._image: Image.Image | None = image
         self.valid: bool = self._image is not None
 
     def is_valid(self) -> bool:
@@ -85,7 +85,7 @@ def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]:
         yield full_page_bbox
 
     def get_page_image(
-        self, scale: float = 1, cropbox: Optional[BoundingBox] = None
+        self, scale: float = 1, cropbox: BoundingBox | None = None
     ) -> Image.Image:
         assert self._image is not None
         img = self._image
@@ -147,20 +147,22 @@ def __init__(
         # Load frames eagerly for thread-safety across pages
         self._frames: List[Image.Image] = []
         try:
-            img = Image.open(self.path_or_stream)  # type: ignore[arg-type]
-
-            # Handle multi-frame and single-frame images
-            # - multiframe formats: TIFF, GIF, ICO
-            # - singleframe formats: JPEG (.jpg, .jpeg), PNG (.png), BMP, WEBP (unless animated), HEIC
-            frame_count = getattr(img, "n_frames", 1)
-
-            if frame_count > 1:
-                for i in range(frame_count):
-                    img.seek(i)
-                    self._frames.append(img.copy().convert("RGB"))
-            else:
-                self._frames.append(img.convert("RGB"))
+            with Image.open(self.path_or_stream) as img:  # type: ignore[arg-type]
+                # Handle multi-frame and single-frame images
+                # - multiframe formats: TIFF, GIF, ICO
+                # - singleframe formats: JPEG (.jpg, .jpeg), PNG (.png), BMP, WEBP (unless animated), HEIC
+                frame_count = getattr(img, "n_frames", 1)
+
+                if frame_count > 1:
+                    for i in range(frame_count):
+                        img.seek(i)
+                        self._frames.append(img.copy().convert("RGB"))
+                else:
+                    self._frames.append(img.convert("RGB"))
         except Exception as e:
+            for frame in self._frames:
+                frame.close()
+            self._frames = []
             raise RuntimeError(f"Could not load image for document {self.file}") from e
 
     def is_valid(self) -> bool:
@@ -184,5 +186,7 @@ def supports_pagination(cls) -> bool:
         return True
 
     def unload(self):
-        super().unload()
+        for frame in self._frames:
+            frame.close()
         self._frames = []
+        super().unload()

tests/test_backend_image_native.py

@@ -1,13 +1,17 @@
 from io import BytesIO
-from pathlib import Path
+from unittest.mock import MagicMock
 
 import pytest
 from docling_core.types.doc import BoundingBox, CoordOrigin
 from PIL import Image
 
 from docling.backend.image_backend import ImageDocumentBackend, _ImagePageBackend
 from docling.datamodel.base_models import DocumentStream, InputFormat
-from docling.datamodel.document import InputDocument, _DocumentConversionInput
+from docling.datamodel.document import (
+    InputDocument,
+    _DocumentConversionInput,
+    _DummyBackend,
+)
 from docling.document_converter import DocumentConverter, ImageFormatOption
 from docling.document_extractor import DocumentExtractor
 
@@ -216,3 +220,69 @@ def test_multipage_access():
         size = page_backend.get_size()
         assert size.width == 64
         assert size.height == 64
+
+
+def test_source_image_is_closed_after_backend_init(tmp_path, monkeypatch):
+    image_path = tmp_path / "test.png"
+    Image.new("RGB", (32, 32), (10, 20, 30)).save(image_path)
+
+    opened_images = []
+    original_open = Image.open
+
+    class TrackingImage:
+        def __init__(self, image):
+            self._image = image
+            self.closed = False
+
+        def __getattr__(self, attr):
+            return getattr(self._image, attr)
+
+        def close(self):
+            self.closed = True
+            return self._image.close()
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc, tb):
+            self.close()
+            return False
+
+    def tracking_open(*args, **kwargs):
+        tracked_image = TrackingImage(original_open(*args, **kwargs))
+        opened_images.append(tracked_image)
+        return tracked_image
+
+    input_doc = InputDocument(
+        path_or_stream=image_path,
+        format=InputFormat.IMAGE,
+        backend=_DummyBackend,
+        filename=image_path.name,
+    )
+
+    monkeypatch.setattr("docling.backend.image_backend.Image.open", tracking_open)
+    backend = ImageDocumentBackend(
+        in_doc=input_doc,
+        path_or_stream=image_path,
+    )
+
+    assert len(opened_images) == 1
+    assert opened_images[0].closed is True
+    backend.unload()
+
+
+def test_unload_closes_cached_frames():
+    stream = _make_multipage_tiff_stream(num_pages=3, size=(32, 32))
+    doc_backend = _get_backend_from_stream(stream)
+
+    tracked_closers = []
+    for frame in doc_backend._frames:
+        closer = MagicMock(wraps=frame.close)
+        frame.close = closer
+        tracked_closers.append(closer)
+
+    doc_backend.unload()
+
+    assert doc_backend._frames == []
+    for closer in tracked_closers:
+        closer.assert_called_once()