fix lint, typeignore a lot of pypdfium stuff

Author: sebastianMindee

mindee/image_extraction/common/__init__.py

@@ -1,2 +1,6 @@
 from mindee.image_extraction.common.extracted_image import ExtractedImage
-from mindee.image_extraction.common.image_extractor import extract_from_page, attach_bitmap_as_new_page, get_image_size
+from mindee.image_extraction.common.image_extractor import (
+    attach_bitmap_as_new_page,
+    extract_from_page,
+    get_image_size,
+)

mindee/image_extraction/common/extracted_image.py

@@ -7,6 +7,8 @@
 
 
 class ExtractedImage:
+    """Generic class for image extraction."""
+
     def __init__(self, buffer: bytes, file_name: str):
         """
         Initialize the ExtractedImage with a buffer and an internal file name.
@@ -26,13 +28,13 @@ def save_to_file(self, output_path: str):
         """
         try:
             resolved_path = Path(output_path).resolve()
-            with open(resolved_path, 'wb') as f:
-                f.write(self.buffer.read())
-            logger.info(f"File saved successfully to {resolved_path}.")
-        except TypeError:
-            raise MindeeError("Invalid path/filename provided.")
-        except Exception as e:
-            raise e
+            with open(resolved_path, "wb") as file:
+                file.write(self.buffer.read())
+                logger.info("File saved successfully to %s.", resolved_path)
+        except TypeError as exc:
+            raise MindeeError("Invalid path/filename provided.") from exc
+        except Exception as exc:
+            raise MindeeError(f"Could not save file {Path(output_path).name}.") from exc
 
     def as_source(self) -> FileInput:
         """

mindee/image_extraction/common/image_extractor.py

@@ -1,12 +1,11 @@
 import io
-from typing import List, BinaryIO, Tuple
+import struct
+from typing import BinaryIO, List, Tuple
 
 import pypdfium2 as pdfium
 
 from mindee.error import MimeTypeError
-from mindee.geometry import get_min_max_x, get_min_max_y, Polygon
-
-import struct
+from mindee.geometry import Polygon, get_min_max_x, get_min_max_y
 
 
 def get_image_size(data: BinaryIO) -> Tuple[int, int]:
@@ -20,48 +19,53 @@ def get_image_size(data: BinaryIO) -> Tuple[int, int]:
     signature = data.read(8)
 
     # Check for PNG signature
-    if signature[:8] == b'\x89PNG\r\n\x1a\n':
-        # PNG file
+    if signature[:8] == b"\x89PNG\r\n\x1a\n":
         data.seek(16)
-        width, height = struct.unpack('>II', data.read(8))
+        width, height = struct.unpack(">II", data.read(8))
         return width, height
 
-    # Check for JPEG SOI marker
-    elif signature[:2] == b'\xff\xd8':
+    # Check for JPEG SOI marker (also works for jpga)
+    if signature[:2] == b"\xff\xd8":
         data.seek(2)
         while True:
-            marker, = struct.unpack('>H', data.read(2))
-            if marker == 0xFFC0 or marker == 0xFFC2:  # SOF0 or SOF2
+            (marker,) = struct.unpack(">H", data.read(2))
+            if marker in (0xFFC0, 0xFFC2):  # SOF0 or SOF2
                 data.seek(3, 1)  # Skip length and precision
-                height, width = struct.unpack('>HH', data.read(4))
+                height, width = struct.unpack(">HH", data.read(4))
                 return width, height
-            else:
-                length, = struct.unpack('>H', data.read(2))
-                data.seek(length - 2, 1)
+            (length,) = struct.unpack(">H", data.read(2))
+            data.seek(length - 2, 1)
     data.close()
     raise MimeTypeError("Size could not be retrieved for file.")
 
 
-def attach_bitmap_as_new_page(pdf_doc: pdfium.PdfDocument, bitmap: pdfium.PdfBitmap, new_width: float,
-                              new_height: float) -> pdfium.PdfDocument:
+def attach_bitmap_as_new_page(  # type: ignore
+    pdf_doc: pdfium.PdfDocument,
+    bitmap: pdfium.PdfBitmap,
+    new_width: float,
+    new_height: float,
+) -> pdfium.PdfDocument:
     """
     Attaches a created PdfBitmap object as a new page in a PdfDocument object.
 
     :param pdf_doc: The PdfDocument to which the new page will be added.
     :param bitmap: The PdfBitmap object to be added as a new page.
     :param new_width: The width of the new page.
     :param new_height: The height of the new page.
+    :return: A PdfDocument handle.
     """
     # Create a new page in the PdfDocument
     new_page = pdf_doc.new_page(new_width, new_height)
 
+    pdf_obj = pdfium.PdfImage.new(pdf_doc)
+    pdf_obj.set_bitmap(bitmap)
     # Create a device context to render the bitmap onto the new page
-    new_page.insert_obj(bitmap.buffer)
+    new_page.insert_obj(pdf_obj)
 
     return pdf_doc
 
 
-def extract_from_page(pdf_page: pdfium.PdfPage, polygons: List[Polygon]):
+def extract_from_page(pdf_page: pdfium.PdfPage, polygons: List[Polygon]):  # type: ignore
     """
     Extracts elements from a page based on a list of bounding boxes.
 
@@ -79,17 +83,21 @@ def extract_from_page(pdf_page: pdfium.PdfPage, polygons: List[Polygon]):
         min_max_x = get_min_max_x(polygon)
         min_max_y = get_min_max_y(polygon)
 
-        new_width = width * (min_max_x.max - min_max_x.min)
-        new_height = height * (min_max_y.max - min_max_y.min)
-
-        left = min_max_x.min * width
-        right = min_max_x.max * width
-        top = height - (min_max_y.min * height)
-        bottom = height - (min_max_y.max * height)
-
-        cropped_page: pdfium.PdfBitmap = pdf_page.render(crop=(left, bottom, right, top))
-
-        temp_pdf = attach_bitmap_as_new_page(temp_pdf, cropped_page, new_width, new_height)
+        left = min_max_x.min
+        right = min_max_x.max
+        top = (height - (min_max_y.min * height)) / height
+        bottom = (height - (min_max_y.max * height)) / height
+
+        cropped_page: pdfium.PdfBitmap = pdf_page.render(  # type: ignore
+            crop=(left, bottom, right, top)
+        )
+
+        temp_pdf = attach_bitmap_as_new_page(
+            temp_pdf,
+            cropped_page,
+            width * (min_max_x.max - min_max_x.min),
+            height * (min_max_y.max - min_max_y.min),
+        )
 
         temp_file = io.BytesIO()
         temp_pdf.save(temp_file)

mindee/image_extraction/multi_receipts_extractor/__init__.py

@@ -1,2 +1,6 @@
-from mindee.image_extraction.multi_receipts_extractor.mult_receipts_extractor import extract_receipts_from_page
-from mindee.image_extraction.multi_receipts_extractor.extracted_mult_receipt_image import ExtractedMultiReceiptImage
+from mindee.image_extraction.multi_receipts_extractor.extracted_multi_receipt_image import (
+    ExtractedMultiReceiptImage,
+)
+from mindee.image_extraction.multi_receipts_extractor.mult_receipts_extractor import (
+    extract_receipts_from_page,
+)

mindee/image_extraction/multi_receipts_extractor/extracted_mult_receipt_image.py renamed to mindee/image_extraction/multi_receipts_extractor/extracted_multi_receipt_image.py

@@ -2,8 +2,10 @@
 
 
 class ExtractedMultiReceiptImage(ExtractedImage):
+    """Wrapper class for extracted multiple-receipts images."""
+
     _receipt_id: int
-    page_id:  int
+    _page_id: int
 
     def __init__(self, buffer, receipt_id: int, page_id: int):
         super().__init__(buffer, f"receipt_p{page_id}_{receipt_id}.pdf")
@@ -12,8 +14,18 @@ def __init__(self, buffer, receipt_id: int, page_id: int):
 
     @property
     def receipt_id(self):
+        """
+        ID of the receipt on a given page.
+
+        :return:
+        """
         return self._receipt_id
 
     @property
     def page_id(self):
-        return self.page_id
+        """
+        ID of the page the receipt was found on.
+
+        :return:
+        """
+        return self._page_id

mindee/image_extraction/multi_receipts_extractor/mult_receipts_extractor.py

@@ -1,18 +1,30 @@
-from typing import List
+from typing import List, Union
 
 import pypdfium2 as pdfium
 
-from mindee.error import MimeTypeError
-from mindee.geometry import Polygon
-from mindee.image_extraction.common.image_extractor import extract_from_page, attach_bitmap_as_new_page, get_image_size
-from mindee.image_extraction.multi_receipts_extractor import ExtractedMultiReceiptImage
+from mindee.error import MimeTypeError, MindeeError
+from mindee.geometry.point import Point
+from mindee.geometry.polygon import Polygon
+from mindee.geometry.quadrilateral import Quadrilateral
+from mindee.image_extraction.common.image_extractor import (
+    attach_bitmap_as_new_page,
+    extract_from_page,
+    get_image_size,
+)
+from mindee.image_extraction.multi_receipts_extractor.extracted_multi_receipt_image import (
+    ExtractedMultiReceiptImage,
+)
 from mindee.input import LocalInputSource
+from mindee.product import MultiReceiptsDetectorV1
 
 
-def extract_receipts_from_page(pdf_page: pdfium.PdfPage, bounding_boxes: List[Polygon], page_id: int) \
-        -> List[ExtractedMultiReceiptImage]:
+def extract_receipts_from_page(  # type: ignore
+    pdf_page: pdfium.PdfPage,
+    bounding_boxes: List[Union[List[Point], Polygon, Quadrilateral]],
+    page_id: int,
+) -> List[ExtractedMultiReceiptImage]:
     """
-    Given a page and a set of coordinates, extracts & assigns individual receipts to an ExtractedMultiReceiptImage
+    Given a page and a set of coordinates, extracts & assigns individual receipts to an ExtractedMultiReceiptImage\
     object.
 
     :param pdf_page: PDF Page to extract from.
@@ -21,26 +33,76 @@ def extract_receipts_from_page(pdf_page: pdfium.PdfPage, bounding_boxes: List[Po
     pages.
     :return: A list of ExtractedMultiReceiptImage.
     """
-    extracted_receipts_raw = extract_from_page(pdf_page, bounding_boxes)
+    extracted_receipts_raw = extract_from_page(pdf_page, bounding_boxes)  # type: ignore
     extracted_receipts = []
-    for i in range(len(extracted_receipts_raw)):
-        extracted_receipts.append(ExtractedMultiReceiptImage(extracted_receipts_raw[i], page_id, i))
+    for i, extracted_receipt_raw in enumerate(extracted_receipts_raw):
+        extracted_receipts.append(
+            ExtractedMultiReceiptImage(extracted_receipt_raw, i, page_id)
+        )
     return extracted_receipts
 
 
-def load_pdf_doc(input_file: LocalInputSource) -> pdfium.PdfDocument:
+def load_pdf_doc(input_file: LocalInputSource) -> pdfium.PdfDocument:  # type: ignore
     """
     Loads a PDF document from a local input source.
 
     :param input_file: Local input.
     :return: A valid PdfDocument handle.
     """
-    if input_file.file_mimetype not in ["image/jpeg", "image/jpg", "image/png", "application/pdf"]:
-        raise MimeTypeError(f"Unsupported file type '{input_file.file_mimetype}'. Currently supported types are '.png',"
-                            f" '.jpg' and '.pdf'.")
+    if input_file.file_mimetype not in [
+        "image/jpeg",
+        "image/jpg",
+        "image/png",
+        "application/pdf",
+    ]:
+        raise MimeTypeError(
+            f"Unsupported file type '{input_file.file_mimetype}'. Currently supported types are '.png',"
+            f" '.jpg' and '.pdf'."
+        )
     if input_file.is_pdf():
-        pdf_document = pdfium.PdfDocument(input_file.file_object)
-    else:
-        pdf_document = pdfium.PdfDocument.new()
+        return pdfium.PdfDocument(input_file.file_object)
+    pdf_document = pdfium.PdfDocument.new()
+    height, width = get_image_size(input_file.file_object)
+    pdf_bitmap = pdfium.PdfBitmap.new_native(width, height, 4)
+    pdf_bitmap = pdfium.PdfBitmap(
+        raw=pdf_bitmap,
+        buffer=input_file.file_object,
+        height=height,
+        width=width,
+        needs_free=True,
+        rev_byteorder=False,
+        format=4,
+        stride=4,
+    )
+    # Bitmap format 4 should equate to RGBA, assumed to be equivalent to:
+    # https://docs.rs/pdfium-render/latest/pdfium_render/bitmap/enum.PdfBitmapFormat.html
 
-    return attach_bitmap_as_new_page(pdf_document, input_file.file_object, get_image_size(input_file.file_object))
+    return attach_bitmap_as_new_page(pdf_document, pdf_bitmap, height, width)
+
+
+def extract_receipts(
+    input_file: LocalInputSource, inference: MultiReceiptsDetectorV1
+) -> List[ExtractedMultiReceiptImage]:
+    """
+    Extracts individual receipts from multi-receipts documents.
+
+    :param input_file: File to extract sub-receipts from.
+    :param inference: Results of the inference.
+    :return: Individual extracted receipts as an array of ExtractedMultiReceiptImage.
+    """
+    images: List[ExtractedMultiReceiptImage] = []
+    if not inference.prediction.receipts:
+        raise MindeeError(
+            "No possible receipts candidates found for MultiReceipts extraction."
+        )
+    pdf_doc = load_pdf_doc(input_file)
+    for page_id in range(len(pdf_doc)):
+        receipt_positions = [
+            receipt.bounding_box
+            for receipt in inference.pages[page_id].prediction.receipts
+        ]
+        extracted_receipts = extract_receipts_from_page(
+            pdf_doc.get_page(page_id), receipt_positions, page_id  # type: ignore
+        )
+        images.extend(extracted_receipts)
+    return images

mindee/parsing/standard/locale.py

@@ -27,7 +27,11 @@ def __init__(
         :param reconstructed: Bool for reconstructed object (not extracted in the API)
         :param page_id: Page number for multi-page document
         """
-        value_key = "value" if ("value" in raw_prediction and raw_prediction["value"]) else "language"
+        value_key = (
+            "value"
+            if ("value" in raw_prediction and raw_prediction["value"])
+            else "language"
+        )
 
         super().__init__(
             raw_prediction,

tests/image_extraction/test_image_extractor.py

@@ -11,9 +11,12 @@
 def jpg_file_path():
     return FILE_TYPES_DIR / "receipt.jpg"
 
+
 @pytest.fixture
 def txt_file_path():
     return FILE_TYPES_DIR / "receipt.txt"
+
+
 @pytest.fixture
 def png_file_path():
     return FILE_TYPES_DIR / "receipt.png"