rehaul all to accomodate newer changes in bernie

Author: sebO-S

mindee/image/extracted_image.py

@@ -69,7 +69,7 @@ def write_to_file(self, output_path: Path | str, file_format: str | None = None)
             if file_format is None:
                 out_file_path.write_bytes(self.buffer.read())
             else:
-                image = bernard_image.decode(self.buffer)
+                image = bernard_image.Image(self.buffer.read())  # type: ignore[call-arg]
                 image.save(out_file_path, format=file_format)
             logger.info("File saved successfully to '%s'.", out_file_path)
         except Exception as e:

mindee/image/image_compressor.py

@@ -2,160 +2,51 @@
 
 import io
 from pathlib import Path
-from typing import BinaryIO
 
 from mindee.dependencies import requires_bernard_ledit
 from mindee.dependencies.checkers import BERNARD_LEDIT_AVAILABLE
 from mindee.geometry.point import Point
-from mindee.geometry.polygon import Polygon, get_min_max_x, get_min_max_y
+from mindee.geometry.polygon import Polygon
 from mindee.image.extracted_image import ExtractedImage
 from mindee.input.local_input_source import LocalInputSource
 
 if BERNARD_LEDIT_AVAILABLE:
     # pylint: disable=import-error
     import bernard_ledit.image as bernard_image
-    import bernard_ledit.pdf as bernard_pdf
 else:
-    bernard_pdf = None  # type: ignore[assignment]  # pylint: disable=invalid-name
     bernard_image = None  # type: ignore[assignment]  # pylint: disable=invalid-name
 
 
-@requires_bernard_ledit
-def _attach_image_as_new_file(
-    input_buffer: BinaryIO,
-) -> bernard_pdf.PdfDocument:
-    """
-    Attaches an image as a new page in a PdfDocument object.
-
-    :param input_buffer: Input buffer.
-    :return: A PdfDocument handle.
-    """
-    input_buffer.seek(0)
-    pdf = bernard_pdf.PdfDocument.new()
-    pdf.append_jpeg_page(input_buffer.read())
-    return pdf
-
-
-@requires_bernard_ledit
-def extract_image_from_polygon(
-    page_content: bernard_image.Image,
-    polygon: list[Point],
-    width: float,
-    height: float,
-    file_format: str,
-    quality: int = 70,
-) -> BinaryIO:
-    """
-    Crops the image from the given polygon.
-
-    :param page_content: Contents of the page as a Bernard L'Édit Image object.
-    :param polygon: Polygon coordinates for the image.
-    :param width: Width of the generated image.
-    :param height: Height of the generated image.
-    :param file_format: Format for the generated file.
-    :param quality: JPEG quality for the output (default 70 — sufficient for OCR).
-    :return: A generated image as a buffer.
-    """
-    min_max_x = get_min_max_x(polygon)
-    min_max_y = get_min_max_y(polygon)
-    cropped_image = page_content.crop(
-        left=int(min_max_x.min * width),
-        top=int(min_max_y.min * height),
-        right=int(min_max_x.max * width),
-        bottom=int(min_max_y.max * height),
-    )
-    return _save_image_to_buffer(cropped_image, file_format, quality)
-
-
-@requires_bernard_ledit
-def _save_image_to_buffer(
-    image: bernard_image.Image, file_format: str, quality: int = 70
-) -> BinaryIO:
-    """
-    Saves an image as a buffer.
-
-    :param image: Bernard L'Édit wrapper for the image.
-    :param file_format: Format to save the file as.
-    :param quality: JPEG quality (default 70 — sufficient for OCR).
-    :return: A valid buffer.
-    """
-    buffer = io.BytesIO()
-    image.save(buffer, format=file_format, quality=quality)
-    buffer.seek(0)
-    return buffer
-
-
-def get_file_extension(file_format: str):
-    """
-    Extract the correct file extension.
-
-    :param file_format: Format of the file.
-    :return: A valid file extension.
-    """
-    return file_format.lower() if file_format != "JPEG" else "jpg"
-
-
 @requires_bernard_ledit
 def extract_multiple_images_from_source(
     input_source: LocalInputSource,
     page_id: int,
     polygons: list[Polygon | list[Point]],
-    quality: int = 70,
 ) -> list[ExtractedImage]:
     """
-    Extracts elements from a page based on a list of bounding boxes.
-
-    :param input_source: Local Input source to extract elements from.
-    :param page_id: id of the page to extract from.
-    :param polygons: List of coordinates to pull the elements from.
-    :param quality: JPEG quality for extracted images (default 70).
-    :return: List of byte arrays representing the extracted elements.
+    Extract multiple images from a source.
+    :param input_source: Source to extract images from.
+    :param page_id: Page ID.
+    :param polygons: List of polygons to extract.
+    :return: List of extracted images.
     """
     stem = Path(input_source.filename).stem
-    doc = _load_pdf_doc(input_source)
-    page_content = doc.rasterize_page(page_id, 100)
-    width = doc.get_page(page_id).get_size().width
-    height = doc.get_page(page_id).get_size().height
-
-    if input_source.is_pdf():
-        file_format = "JPEG"
-    else:
-        input_source.file_object.seek(0)
-        file_format = bernard_image.guess_format(input_source.file_object)
-
-    extracted_elements = []
-    decoded_page = bernard_image.decode(page_content)
-    for element_id, polygon in enumerate(polygons):
-        image_data = extract_image_from_polygon(
-            decoded_page,
-            polygon,
-            width,
-            height,
-            file_format,
-            quality,
-        )
-        extracted_elements.append(
+    input_source.file_object.seek(0)
+    pts = [[(p.x, p.y) for p in poly] for poly in polygons]
+    images = bernard_image.extract_polygons(input_source.file_object, page_id, pts)  # type: ignore[attr-defined]
+
+    extracted = []
+    for element_id, img in enumerate(images):
+        buf = io.BytesIO()
+        img.save(buf)
+        fmt = str(img.format).lower() if img.format else "jpeg"
+        fmt = fmt if fmt != "jpeg" else "jpg"
+        extracted.append(
             ExtractedImage(
-                image_data,
-                f"{stem}_page-{(page_id + 1):03d}-item-{(element_id + 1):03d}."
-                f"{get_file_extension(file_format)}",
+                buf,
+                f"{stem}_page-{page_id + 1:03d}-item-{element_id + 1:03d}.{fmt}",
                 page_id,
                 element_id,
             )
         )
-    return extracted_elements
-
-
-@requires_bernard_ledit
-def _load_pdf_doc(input_file: LocalInputSource) -> bernard_pdf.PdfDocument:
-    """
-    Loads a PDF document from a local input source.
-
-    :param input_file: Local input.
-    :return: A valid PdfDocument handle.
-    """
-    if input_file.is_pdf():
-        input_file.file_object.seek(0)
-        return bernard_pdf.PdfDocument(input_file.file_object.read())
-
-    return _attach_image_as_new_file(input_file.file_object)
+    return extracted

mindee/image/image_extractor.py

@@ -10,17 +10,15 @@
 from mindee.dependencies.checkers import BERNARD_LEDIT_AVAILABLE
 from mindee.error.mimetype_error import MimeTypeError
 from mindee.error.mindee_error import MindeeError, MindeeSourceError
-from mindee.image.image_compressor import compress_image
 from mindee.input.page_options import KEEP_ONLY, REMOVE, PageOptions
 from mindee.logger import logger
-from mindee.pdf.pdf_compressor import compress_pdf
 from mindee.pdf.pdf_utils import pdf_has_source_text
 
 if BERNARD_LEDIT_AVAILABLE:
     # pylint: disable=import-error
-    import bernard_ledit.pdf as bernard_pdf
+    import bernard_ledit as bernard
 else:
-    bernard_pdf = None  # type: ignore[assignment]  # pylint: disable=invalid-name
+    bernard = None  # type: ignore[assignment]  # pylint: disable=invalid-name
 
 mimetypes.add_type("image/heic", ".heic")
 mimetypes.add_type("image/heif", ".heif")
@@ -56,9 +54,9 @@ def __init__(self) -> None:
             # Some broken (yet fixable) PDFs can cause Bernard to crash on open.
             if BERNARD_LEDIT_AVAILABLE:
                 try:
-                    pdf = bernard_pdf.PdfDocument(self.file_object)
+                    pdf = bernard.pdf.PdfDocument(self.file_object)
                     self.page_count = len(pdf)
-                except bernard_pdf.PdfiumError as e:
+                except bernard.pdf.PdfiumError as e:
                     logger.warning(
                         "Could not open PDF file: %s due to %s", self.filename, e
                     )
@@ -139,7 +137,7 @@ def apply_page_options(self, page_options: PageOptions) -> None:
             page_options.page_indexes,
         )
         self.file_object.seek(0)
-        pdf = bernard_pdf.PdfDocument(self.file_object)
+        pdf = bernard.pdf.PdfDocument(self.file_object)
         self.page_count = len(pdf)
         pdf.close()
 
@@ -187,8 +185,8 @@ def merge_pdf_pages(self, page_numbers: set) -> None:
         :return: None
         """
         self.file_object.seek(0)
-        new_pdf = bernard_pdf.PdfDocument.new()
-        pdf = bernard_pdf.PdfDocument(self.file_object)
+        new_pdf = bernard.pdf.PdfDocument.new()
+        pdf = bernard.pdf.PdfDocument(self.file_object)
         new_pdf.import_pages(pdf, list(page_numbers))
         self.file_object.close()
         bytes_io = io.BytesIO()
@@ -207,7 +205,7 @@ def is_pdf_empty(self) -> bool:
         :return: ``True`` if the PDF is empty
         """
         self.file_object.seek(0)
-        with bernard_pdf.PdfDocument(self.file_object) as pdf:
+        with bernard.pdf.PdfDocument(self.file_object) as pdf:
             return pdf.has_no_content()
         raise MindeeSourceError(f"PDF couldn't be accessed: {self.filename}")
 
@@ -248,8 +246,6 @@ def compress(
         quality: int = 85,
         max_width: int | None = None,
         max_height: int | None = None,
-        force_source_text: bool = False,
-        disable_source_text: bool = True,
     ) -> None:
         """
         Compresses the file object, either as a PDF or an image.
@@ -258,20 +254,17 @@ def compress(
             For PDFs, this affects image quality within the PDF.
         :param max_width: Maximum width for image resizing. Ignored for PDFs.
         :param max_height: Maximum height for image resizing. Ignored for PDFs.
-        :param force_source_text: For PDFs, whether to force compression even if source text is present.
-        :param disable_source_text: For PDFs, whether to disable source text during compression.
         """
-        new_file_bytes: bytes
+
+        self.file_object.seek(0)
+        data = self.file_object.read()
         if self.is_pdf():
-            new_file_bytes = compress_pdf(
-                self.file_object, quality, force_source_text, disable_source_text
-            )
+            compressed = bernard.pdf.compress(data, quality)
         else:
-            new_file_bytes = compress_image(
-                self.file_object, quality, max_width, max_height
-            )
+            compressed = bernard.image.compress(data, quality, max_width, max_height)
 
-        self.file_object = io.BytesIO(new_file_bytes)
+        self.file_object = io.BytesIO(compressed)
+        self.file_object.seek(0)
 
     def __enter__(self):
         """Allows the class to be used as a context manager."""

mindee/input/local_input_source.py

@@ -1,11 +1,9 @@
-from mindee.pdf.pdf_compressor import compress_pdf
 from mindee.pdf.pdf_utils import (
     lerp,
     pdf_has_source_text,
 )
 
 __all__ = [
-    "compress_pdf",
     "lerp",
     "pdf_has_source_text",
 ]