fix for test?

Author: ianardee

mindee/image/extracted_image.py

@@ -20,42 +20,30 @@ class ExtractedImage:
     """Generic class for image extraction."""
 
     buffer: BinaryIO
+    filename: str
     _page_id: int
     """Id of the page the image was extracted from."""
     _element_id: int
     """Id of the element on a given page."""
-    filename: str
-    """Name of the file the image was extracted from."""
 
     def __init__(
         self,
         img_byte_stream: BinaryIO,
-        orig_filename: str,
-        orig_extension: str,
+        filename: str,
         page_id: int,
         element_id: int,
     ) -> None:
         """
         Initialize the ExtractedImage with a buffer and an internal file name.
 
         :param img_byte_stream: The raw image bytes.
-        :param orig_filename: Name of the file the image was extracted from.
+        :param filename: Name of the file.
         :param page_id: ID of the page the element was found on.
         :param element_id: ID of the element in a page.
         """
         self.buffer = img_byte_stream
-        self.filename = orig_filename
-
-        if orig_extension.lower().endswith("pdf"):
-            extension = "jpg"
-        else:
-            extension = orig_extension.lower()
         self.buffer.seek(0)
-        pg_number = str(page_id).zfill(3)
-        elem_number = str(element_id).zfill(3)
-        self.internal_file_name = (
-            f"{orig_filename}_page{pg_number}-{elem_number}.{extension}"
-        )
+        self.filename = filename
         self._page_id = page_id
         self._element_id = 0 if element_id is None else element_id
 
@@ -88,7 +76,7 @@ def as_input_source(self) -> BytesInput:
         :returns: A BufferInput source.
         """
         self.buffer.seek(0)
-        return BytesInput(self.buffer.read(), self.internal_file_name)
+        return BytesInput(self.buffer.read(), self.filename)
 
     @property
     def page_id(self):

mindee/image/extracted_images.py

@@ -1,5 +1,12 @@
+from pathlib import Path
+
 from mindee.image.extracted_image import ExtractedImage
 
 
 class ExtractedImages(list[ExtractedImage]):
     """List of extracted images."""
+
+    def save_all_to_disk(self, output_path: Path | str) -> None:
+        """Save all extracted images to disk."""
+        for image in self:
+            image.save_to_file(output_path)

mindee/image/image_extractor.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import io
+from pathlib import Path
 from typing import Any, BinaryIO
 
 from mindee.dependencies import requires_pypdfium2
@@ -28,7 +29,7 @@
 
 @requires_pillow
 @requires_pypdfium2
-def attach_image_as_new_file(  # type: ignore
+def _attach_image_as_new_file(  # type: ignore
     input_buffer: BinaryIO,
 ) -> pdfium.PdfDocument:
     """
@@ -86,11 +87,11 @@ def extract_image_from_polygon(
             int(min_max_y.max * height),
         )
     )
-    return save_image_to_buffer(cropped_image, file_format)
+    return _save_image_to_buffer(cropped_image, file_format)
 
 
 @requires_pillow
-def save_image_to_buffer(image: Image.Image, file_format: str) -> BinaryIO:
+def _save_image_to_buffer(image: Image.Image, file_format: str) -> BinaryIO:
     """
     Saves an image as a buffer.
 
@@ -144,7 +145,8 @@ def extract_multiple_images_from_source(
     :param polygons: List of coordinates to pull the elements from.
     :return: List of byte arrays representing the extracted elements.
     """
-    page = load_pdf_doc(input_source).get_page(page_id)
+    stem = Path(input_source.filename).stem
+    page = _load_pdf_doc(input_source).get_page(page_id)
     page_content = page.render().to_pil()
     width, height = page.get_size()
 
@@ -159,18 +161,16 @@ def extract_multiple_images_from_source(
         extracted_elements.append(
             ExtractedImage(
                 image_data,
-                input_source.filename,
-                file_extension,
+                f"{stem}_page-{(page_id + 1):03d}-item-{(element_id + 1):03d}.{file_extension}",
                 page_id,
                 element_id,
             )
         )
-
     return extracted_elements
 
 
 @requires_pypdfium2
-def load_pdf_doc(input_file: LocalInputSource) -> pdfium.PdfDocument:  # type: ignore
+def _load_pdf_doc(input_file: LocalInputSource) -> pdfium.PdfDocument:  # type: ignore
     """
     Loads a PDF document from a local input source.
 
@@ -181,4 +181,4 @@ def load_pdf_doc(input_file: LocalInputSource) -> pdfium.PdfDocument:  # type: i
         input_file.file_object.seek(0)
         return pdfium.PdfDocument(input_file.file_object.read())
 
-    return attach_image_as_new_file(input_file.file_object)
+    return _attach_image_as_new_file(input_file.file_object)

mindee/pdf/extracted_pdf.py

@@ -20,10 +20,14 @@ class ExtractedPDF:
 
     buffer: BinaryIO
     filename: str
+    _page_indexes: tuple[int, int]
 
-    def __init__(self, pdf_byte_stream: BinaryIO, filename: str):
+    def __init__(
+        self, pdf_byte_stream: BinaryIO, filename: str, page_indexes: tuple[int, int]
+    ):
         self.buffer = pdf_byte_stream
         self.filename = filename
+        self._page_indexes = page_indexes
 
     @requires_pypdfium2
     def get_page_count(self) -> int:
@@ -40,12 +44,12 @@ def save_to_file(self, output_path: Path | str):
         """
         Writes the contents of the current PDF object to a file.
 
-        :param output_path: Path of the destination file. If
-         not extension is provided, pdf will be appended by default.
+        :param output_path: Path of the destination file.
+        If no extension is provided, '.pdf' will be appended by default.
         """
-        out_path = Path(output_path)
-        if out_path.resolve().is_dir():
-            raise MindeeError("Provided path is not a file.")
+        out_path = Path(output_path) / self.filename
+        if not out_path.resolve().is_dir():
+            raise MindeeError("Provided path is not a directory.")
         if not output_path or not out_path.parent.exists():
             raise MindeeError("Invalid save path provided {}.")
         if out_path.suffix.lower() != "pdf":
@@ -58,3 +62,8 @@ def as_input_source(self) -> BytesInput:
         """Returns the current PDF object as a usable BytesInput source."""
         self.buffer.seek(0)
         return BytesInput(self.buffer.read(), self.filename)
+
+    @property
+    def page_indexes(self) -> tuple[int, int]:
+        """This PDF was extracted from this page range of the original PDF."""
+        return self._page_indexes

mindee/pdf/extracted_pdfs.py

@@ -1,5 +1,13 @@
+from pathlib import Path
+
 from mindee.pdf.extracted_pdf import ExtractedPDF
 
 
 class ExtractedPDFs(list[ExtractedPDF]):
     """List of extracted PDFs."""
+
+    def save_all_to_disk(self, output_path: Path | str) -> None:
+        """Save all extracted images to disk."""
+
+        for image in self:
+            image.save_to_file(output_path)

mindee/pdf/pdf_extractor.py

@@ -68,7 +68,7 @@ def extract_sub_documents(
         """
         Extract the sub-documents from the main pdf, based on the given list of page indexes.
 
-        :param page_indexes: List of list of numbers, representing page indexes.
+        :param page_indexes: 2D list of numbers, representing page indexes.
         :return: A list of created PDFS.
         """
         extracted_pdfs: list[ExtractedPDF] = []
@@ -80,10 +80,12 @@ def extract_sub_documents(
             for page_index in page_index_elem:
                 if page_index > self.get_page_count():
                     raise MindeeError(f"Index {page_index} is out of range.")
-            formatted_max_index = f"{page_index_elem[len(page_index_elem) - 1] + 1:03d}"
-            field_filename = f"{stem}_{(page_index_elem[0] + 1):03d}-{formatted_max_index}{extension}"
+            first_page = page_index_elem[0]
+            last_page = page_index_elem[len(page_index_elem) - 1]
             extracted_pdf = ExtractedPDF(
-                self.cut_pages(page_index_elem), field_filename
+                self.cut_pages(page_index_elem),
+                f"{stem}_{(first_page + 1):03d}-{(last_page + 1):03d}{extension}",
+                (first_page, last_page),
             )
             extracted_pdfs.append(extracted_pdf)
         return extracted_pdfs

tests/v2/file_operations/test_crop_operation.py

@@ -13,56 +13,51 @@
 Image = pytest.importorskip("PIL.Image")
 
 
-@pytest.fixture
-def crops_single_page_path():
-    return V2_PRODUCT_DATA_DIR / "crop" / "default_sample.jpg"
-
-
-@pytest.fixture
-def crops_multi_page_path():
-    return V2_PRODUCT_DATA_DIR / "crop" / "multipage_sample.pdf"
-
-
-@pytest.fixture
-def crops_single_page_json_path():
-    return V2_PRODUCT_DATA_DIR / "crop" / "crop_single.json"
-
-
-@pytest.fixture
-def crops_multi_page_json_path():
-    return V2_PRODUCT_DATA_DIR / "crop" / "crop_multiple.json"
-
-
 @pytest.mark.pillow
 @pytest.mark.pypdfium2
-def test_single_page_crop_split(crops_single_page_path, crops_single_page_json_path):
-    input_sample = PathInput(crops_single_page_path)
-    with open(crops_single_page_json_path, "rb") as f:
+def test_single_page_crop():
+    input_sample = PathInput(V2_PRODUCT_DATA_DIR / "crop" / "default_sample.jpg")
+    with open(V2_PRODUCT_DATA_DIR / "crop" / "default_sample.json", "rb") as f:
         response = CropResponse(json.load(f))
     extracted_crops = response.inference.result.extract_from_input_source(input_sample)
-    assert len(extracted_crops) == 1
+    assert len(extracted_crops) == 2
+
+    crop0 = extracted_crops[0]
+    assert crop0.page_id == 0
+    assert crop0.element_id == 0
+    assert crop0.filename == "default_sample_page-001-item-001.jpg"
+    assert Image.open(crop0.buffer).size == (1057, 2071)
 
-    assert extracted_crops[0].page_id == 0
-    assert extracted_crops[0].element_id == 0
-    image_buffer_0 = Image.open(extracted_crops[0].buffer)
-    assert image_buffer_0.size == (2823, 1571)
+    crop1 = extracted_crops[1]
+    assert crop1.page_id == 0
+    assert crop1.element_id == 1
+    assert crop1.filename == "default_sample_page-001-item-002.jpg"
+    assert Image.open(crop1.buffer).size == (1298, 1869)
 
 
 @pytest.mark.pillow
 @pytest.mark.pypdfium2
-def test_multi_page_receipt_crop(crops_multi_page_path, crops_multi_page_json_path):
-    input_sample = PathInput(crops_multi_page_path)
-    with open(crops_multi_page_json_path, "rb") as f:
+def test_multi_page_crop():
+    input_sample = PathInput(V2_PRODUCT_DATA_DIR / "crop" / "multipage_sample.pdf")
+    with open(V2_PRODUCT_DATA_DIR / "crop" / "multipage_sample.json", "rb") as f:
         response = CropResponse(json.load(f))
     extracted_crops = response.inference.result.extract_from_input_source(input_sample)
-    assert len(extracted_crops) == 2
-
-    assert extracted_crops[0].page_id == 0
-    assert extracted_crops[0].element_id == 0
-    image_buffer_0 = Image.open(extracted_crops[0].buffer)
-    assert image_buffer_0.size == (156, 758)
-
-    assert extracted_crops[1].page_id == 0
-    assert extracted_crops[1].element_id == 1
-    image_buffer_1 = Image.open(extracted_crops[1].buffer)
-    assert image_buffer_1.size == (187, 690)
+    assert len(extracted_crops) == 5
+
+    crop0 = extracted_crops[0]
+    assert crop0.page_id == 0
+    assert crop0.element_id == 0
+    assert crop0.filename == "multipage_sample_page-001-item-001.jpg"
+    assert Image.open(crop0.buffer).size == (200, 553)
+
+    crop1 = extracted_crops[1]
+    assert crop1.page_id == 0
+    assert crop1.element_id == 1
+    assert crop1.filename == "multipage_sample_page-001-item-002.jpg"
+    assert Image.open(crop1.buffer).size == (203, 333)
+
+    crop4 = extracted_crops[4]
+    assert crop4.page_id == 1
+    assert crop4.element_id == 1
+    assert crop4.filename == "multipage_sample_page-002-item-002.jpg"
+    assert Image.open(crop4.buffer).size == (197, 520)

tests/v2/file_operations/test_crop_operation_integration.py

@@ -15,16 +15,17 @@
 from tests.utils import OUTPUT_DIR, V2_PRODUCT_DATA_DIR, cleanup_output_files
 
 
-@pytest.fixture
-def crop_sample():
-    return V2_PRODUCT_DATA_DIR / "crop" / "default_sample.jpg"
-
-
 def check_findoc_return(findoc_response: ExtractionResponse):
     assert len(findoc_response.inference.model.id) > 0
     assert findoc_response.inference.result.fields.get("total_amount").value > 0
 
 
+output_files = [
+    "default_sample_page-001-item-001.jpg",
+    "default_sample_page-001-item-001.jpg",
+]
+
+
 @pytest.mark.pillow
 @pytest.mark.pypdfium2
 @pytest.mark.integration
@@ -38,30 +39,30 @@ def test_image_should_extract_crops():
     )
     assert len(response.inference.result.crops) == 2
 
-    extracted_images = extract_multiple_crops(
+    extracted_crops = extract_multiple_crops(
         crop_input, response.inference.result.crops
     )
 
-    assert len(extracted_images) == 2
-    assert extracted_images[0].filename == "default_sample.jpg_page1-0.jpg"
-    assert extracted_images[1].filename == "default_sample.jpg_page1-1.jpg"
+    assert len(extracted_crops) == 2
+    assert extracted_crops[0].filename == output_files[0]
+    assert extracted_crops[1].filename == output_files[1]
 
     invoice_0 = client.enqueue_and_get_result(
         ExtractionResponse,
-        extracted_images[0].as_input_source(),
+        extracted_crops[0].as_input_source(),
         ExtractionParameters(
             getenv("MINDEE_V2_SE_TESTS_FINDOC_MODEL_ID"), close_file=False
         ),
     )
     check_findoc_return(invoice_0)
-    extracted_images.save_all_to_disk(OUTPUT_DIR)
-    crop1size = os.path.getsize(OUTPUT_DIR / "crop_001.jpg")
-    crop2size = os.path.getsize(OUTPUT_DIR / "crop_002.jpg")
-    assert 180000 <= crop1size <= 199685
-    assert 190000 <= crop2size <= 199433
+    extracted_crops.save_all_to_disk(OUTPUT_DIR)
+    crop0_size = os.path.getsize(OUTPUT_DIR / output_files[0])
+    crop1_size = os.path.getsize(OUTPUT_DIR / output_files[1])
+    assert 180000 <= crop0_size <= 199685
+    assert 190000 <= crop1_size <= 199433
 
 
 @pytest.fixture(scope="module", autouse=True)
 def cleanup():
     yield
-    cleanup_output_files(["crop_001.jpg", "crop_002.jpg"])
+    cleanup_output_files()