From 3a72872476e469dc13de15d9323577196b579f75 Mon Sep 17 00:00:00 2001 From: sebastianMindee Date: Mon, 17 Jun 2024 14:03:15 +0200 Subject: [PATCH 1/7] :recycle: switch default image save format to PNG for image extractor --- mindee/image_extraction/common/image_extractor.py | 6 +++--- tests/image_extraction/test_image_extractor.py | 3 +++ 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/mindee/image_extraction/common/image_extractor.py b/mindee/image_extraction/common/image_extractor.py index f703f95e..ed6edaa2 100644 --- a/mindee/image_extraction/common/image_extractor.py +++ b/mindee/image_extraction/common/image_extractor.py @@ -15,7 +15,7 @@ def attach_image_as_new_file( # type: ignore """ Attaches an image as a new page in a PdfDocument object. - :param input_buffer: Input buffer. Only supports JPEG. + :param input_buffer: Input buffer. :return: A PdfDocument handle. """ # Create a new page in the PdfDocument @@ -70,13 +70,13 @@ def extract_multiple_images_from_source( ) ) buffer = io.BytesIO() - pillow_page.save(buffer, format="JPEG") + pillow_page.save(buffer, format="PNG") buffer.seek(0) extracted_elements.append( ExtractedImage( BytesInput( buffer.read(), - f"{input_source.filename}_p{page_id}_e{element_id}.jpg", + f"{input_source.filename}_p{page_id}_e{element_id}.png", ), page_id, element_id, diff --git a/tests/image_extraction/test_image_extractor.py b/tests/image_extraction/test_image_extractor.py index 8e0e0404..c7771e10 100644 --- a/tests/image_extraction/test_image_extractor.py +++ b/tests/image_extraction/test_image_extractor.py @@ -35,6 +35,9 @@ def test_barcode_image_extraction(barcode_path, barcode_json_path): assert len(extracted_barcodes_1d) == 1 assert len(extracted_barcodes_2d) == 2 + assert extracted_barcodes_1d[0].as_source().filename.endswith("png") assert Image.open(extracted_barcodes_1d[0].buffer).size == (353, 200) assert Image.open(extracted_barcodes_2d[0].buffer).size == (214, 216) + assert extracted_barcodes_2d[0].as_source().filename.endswith("png") + assert extracted_barcodes_2d[1].as_source().filename.endswith("png") assert Image.open(extracted_barcodes_2d[1].buffer).size == (193, 201) From 39bb365ca4c8fce7c8fcca7fa9ee7adbf4c0ea98 Mon Sep 17 00:00:00 2001 From: sebastianMindee Date: Mon, 17 Jun 2024 14:05:13 +0200 Subject: [PATCH 2/7] update page indexes to start at one for pdfs --- mindee/image_extraction/common/image_extractor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mindee/image_extraction/common/image_extractor.py b/mindee/image_extraction/common/image_extractor.py index ed6edaa2..3cb9790a 100644 --- a/mindee/image_extraction/common/image_extractor.py +++ b/mindee/image_extraction/common/image_extractor.py @@ -76,7 +76,7 @@ def extract_multiple_images_from_source( ExtractedImage( BytesInput( buffer.read(), - f"{input_source.filename}_p{page_id}_e{element_id}.png", + f"{input_source.filename}_p{page_id + 1}_e{element_id}.png", ), page_id, element_id, From 0afb6d66cd2903a6468f70f0d92147f2f551d6cb Mon Sep 17 00:00:00 2001 From: sebastianMindee Date: Mon, 17 Jun 2024 15:19:39 +0200 Subject: [PATCH 3/7] fix savename syntax --- mindee/image_extraction/common/image_extractor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mindee/image_extraction/common/image_extractor.py b/mindee/image_extraction/common/image_extractor.py index 3cb9790a..772e5abe 100644 --- a/mindee/image_extraction/common/image_extractor.py +++ b/mindee/image_extraction/common/image_extractor.py @@ -76,7 +76,7 @@ def extract_multiple_images_from_source( ExtractedImage( BytesInput( buffer.read(), - f"{input_source.filename}_p{page_id + 1}_e{element_id}.png", + f"{input_source.filename}_page{page_id + 1}-{element_id}.png", ), page_id, element_id, From be89102ec453951457150a53148d221114cfe0c7 Mon Sep 17 00:00:00 2001 From: sebastianMindee Date: Mon, 17 Jun 2024 16:06:19 +0200 Subject: [PATCH 4/7] fix extension formats --- .../common/image_extractor.py | 92 ++++++++++++++++--- .../image_extraction/test_image_extractor.py | 6 +- 2 files changed, 80 insertions(+), 18 deletions(-) diff --git a/mindee/image_extraction/common/image_extractor.py b/mindee/image_extraction/common/image_extractor.py index 772e5abe..01976b04 100644 --- a/mindee/image_extraction/common/image_extractor.py +++ b/mindee/image_extraction/common/image_extractor.py @@ -4,6 +4,7 @@ import pypdfium2 as pdfium from PIL import Image +from mindee.error import MindeeError from mindee.geometry import Point, get_min_max_x, get_min_max_y from mindee.image_extraction.common import ExtractedImage from mindee.input import BytesInput, LocalInputSource @@ -41,6 +42,75 @@ def attach_image_as_new_file( # type: ignore return pdf +def extract_image_from_polygon( + page_content: Image.Image, + polygon: List[Point], + width: float, + height: float, + file_format: str, +) -> bytes: + """ + Crops the image from the given polygon. + + :param page_content: + :param polygon: + :param width: + :param height: + :param file_format: + :return: + """ + min_max_x = get_min_max_x(polygon) + min_max_y = get_min_max_y(polygon) + cropped_image = page_content.crop( + ( + int(min_max_x.min * width), + int(min_max_y.min * height), + int(min_max_x.max * width), + int(min_max_y.max * height), + ) + ) + return save_image_to_buffer(cropped_image, file_format) + + +def save_image_to_buffer(image: Image.Image, file_format: str) -> bytes: + """ + Saves an image as a buffer. + + :param image: Pillow wrapper for the image. + :param file_format: Format to save the file as. + :return: + """ + buffer = io.BytesIO() + image.save(buffer, format=file_format) + buffer.seek(0) + return buffer.read() + + +def determine_file_format(input_source: LocalInputSource): + """ + Retrieves the file format from an input source. + + :param input_source: Local input source to retrieve the format from. + :return: + """ + if input_source.is_pdf(): + return "JPEG" + img = Image.open(input_source.file_object) + if img.format is None: + raise MindeeError("Image format was not found.") + return img.format + + +def get_file_extension(file_format: str): + """ + Extract the correct file extension. + + :param file_format: Format of the file. + :return: + """ + return file_format.lower() if file_format != "JPEG" else "jpg" + + def extract_multiple_images_from_source( input_source: LocalInputSource, page_id: int, polygons: List[List[Point]] ) -> List[ExtractedImage]: @@ -56,27 +126,19 @@ def extract_multiple_images_from_source( page_content = page.render().to_pil() width, height = page.get_size() + file_format = determine_file_format(input_source) + file_extension = get_file_extension(file_format) + extracted_elements = [] for element_id, polygon in enumerate(polygons): - min_max_x = get_min_max_x(polygon) - min_max_y = get_min_max_y(polygon) - - pillow_page = page_content.crop( - ( - int(min_max_x.min * width), - int(min_max_y.min * height), - int(min_max_x.max * width), - int(min_max_y.max * height), - ) + image_data = extract_image_from_polygon( + page_content, polygon, width, height, file_format ) - buffer = io.BytesIO() - pillow_page.save(buffer, format="PNG") - buffer.seek(0) extracted_elements.append( ExtractedImage( BytesInput( - buffer.read(), - f"{input_source.filename}_page{page_id + 1}-{element_id}.png", + image_data, + f"{input_source.filename}_page{page_id + 1}-{element_id}.{file_extension}", ), page_id, element_id, diff --git a/tests/image_extraction/test_image_extractor.py b/tests/image_extraction/test_image_extractor.py index c7771e10..08126fed 100644 --- a/tests/image_extraction/test_image_extractor.py +++ b/tests/image_extraction/test_image_extractor.py @@ -35,9 +35,9 @@ def test_barcode_image_extraction(barcode_path, barcode_json_path): assert len(extracted_barcodes_1d) == 1 assert len(extracted_barcodes_2d) == 2 - assert extracted_barcodes_1d[0].as_source().filename.endswith("png") + assert extracted_barcodes_1d[0].as_source().filename.endswith("jpg") assert Image.open(extracted_barcodes_1d[0].buffer).size == (353, 200) assert Image.open(extracted_barcodes_2d[0].buffer).size == (214, 216) - assert extracted_barcodes_2d[0].as_source().filename.endswith("png") - assert extracted_barcodes_2d[1].as_source().filename.endswith("png") + assert extracted_barcodes_2d[0].as_source().filename.endswith("jpg") + assert extracted_barcodes_2d[1].as_source().filename.endswith("jpg") assert Image.open(extracted_barcodes_2d[1].buffer).size == (193, 201) From 1dc5b7a9573dcbfaacb7c220c0eaa5092b2244a5 Mon Sep 17 00:00:00 2001 From: sebastianMindee Date: Mon, 17 Jun 2024 17:02:53 +0200 Subject: [PATCH 5/7] fix imporper extension assignment --- mindee/image_extraction/common/extracted_image.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/mindee/image_extraction/common/extracted_image.py b/mindee/image_extraction/common/extracted_image.py index 39ae0dd8..b9bcb315 100644 --- a/mindee/image_extraction/common/extracted_image.py +++ b/mindee/image_extraction/common/extracted_image.py @@ -29,8 +29,16 @@ def __init__( """ self.buffer = io.BytesIO(input_source.file_object.read()) self.buffer.name = input_source.filename + if input_source.is_pdf(): + extension = "jpg" + else: + extension = Path(input_source.filename).resolve().suffix self.buffer.seek(0) - self.internal_file_name = f"{input_source.filename}_p{page_id}_{element_id}.pdf" + pg_number = str(page_id).zfill(3) + elem_number = str(element_id).zfill(3) + self.internal_file_name = ( + f"{input_source.filename}_page{pg_number}-{elem_number}.{extension}" + ) self._page_id = page_id self._element_id = 0 if element_id is None else element_id From 71ea2119fbf0d226e3f3c994527f9410d7fdfc62 Mon Sep 17 00:00:00 2001 From: sebastianMindee Date: Mon, 17 Jun 2024 17:36:29 +0200 Subject: [PATCH 6/7] fix docs --- .../common/image_extractor.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/mindee/image_extraction/common/image_extractor.py b/mindee/image_extraction/common/image_extractor.py index 01976b04..f049131a 100644 --- a/mindee/image_extraction/common/image_extractor.py +++ b/mindee/image_extraction/common/image_extractor.py @@ -52,12 +52,12 @@ def extract_image_from_polygon( """ Crops the image from the given polygon. - :param page_content: - :param polygon: - :param width: - :param height: - :param file_format: - :return: + :param page_content: Contents of the page as a Pillow 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. + :return: A generated image as a buffer. """ min_max_x = get_min_max_x(polygon) min_max_y = get_min_max_y(polygon) @@ -78,7 +78,7 @@ def save_image_to_buffer(image: Image.Image, file_format: str) -> bytes: :param image: Pillow wrapper for the image. :param file_format: Format to save the file as. - :return: + :return: A valid buffer. """ buffer = io.BytesIO() image.save(buffer, format=file_format) @@ -86,12 +86,12 @@ def save_image_to_buffer(image: Image.Image, file_format: str) -> bytes: return buffer.read() -def determine_file_format(input_source: LocalInputSource): +def determine_file_format(input_source: LocalInputSource) -> str: """ Retrieves the file format from an input source. :param input_source: Local input source to retrieve the format from. - :return: + :return: A valid pillow file format. """ if input_source.is_pdf(): return "JPEG" @@ -106,7 +106,7 @@ def get_file_extension(file_format: str): Extract the correct file extension. :param file_format: Format of the file. - :return: + :return: A valid file extension. """ return file_format.lower() if file_format != "JPEG" else "jpg" From 46870029c885f5ff7477955fe4d962882322882b Mon Sep 17 00:00:00 2001 From: sebastianMindee Date: Mon, 17 Jun 2024 17:37:07 +0200 Subject: [PATCH 7/7] fix more doc --- mindee/image_extraction/common/extracted_image.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mindee/image_extraction/common/extracted_image.py b/mindee/image_extraction/common/extracted_image.py index b9bcb315..43c434fb 100644 --- a/mindee/image_extraction/common/extracted_image.py +++ b/mindee/image_extraction/common/extracted_image.py @@ -82,15 +82,15 @@ def page_id(self): """ ID of the page the receipt was found on. - :return: + :return: A valid page ID. """ return self._page_id @property def element_id(self): """ - Id of the element on a given page. + ID of the element on a given page. - :return: + :return: A valid element ID. """ return self._element_id