added parsable pdf image extractor

image_prediction/image_extractor/extractors/parsable.py

@@ -0,0 +1,32 @@
+from itertools import chain
+from operator import itemgetter
+
+import fitz
+
+from image_prediction.image_extractor.extractor import ImageExtractor, ImageMetadataPair
+
+
+class ParsablePDFImageExtractor(ImageExtractor):
+    def __init__(self):
+        self.doc: fitz.fitz.Document = None
+        self.page: fitz.fitz.Page = None
+
+    def __build_metadata(self, xref):
+        metadata = self.page.get_image_info(xref)
+        page_width, page_height = self.page.mediabox_size
+        metadata = {**metadata, "page_width": page_width, "page_height": page_height, "page_idx": self.page.number}
+        return metadata
+
+    def __process_images_on_page(self, page: fitz.fitz.Page):
+        self.page = page
+
+        image_handles = page.get_images(full=True)
+        xrefs = itemgetter(0)(image_handles)
+        images = map(lambda xref: self.doc.extract_image(xref)["image"], xrefs)
+        metadata = map(self.__build_metadata, xrefs)
+        return map(ImageMetadataPair, zip(images, metadata))
+
+    def extract(self, pdf: bytes):
+        self.doc = fitz.Document(stream=pdf)
+        image_metadata_pairs = chain(*map(self.__process_images_on_page, self.doc))
+        return image_metadata_pairs

requirements.txt

@@ -21,3 +21,4 @@ PyYAML~=5.4.1
 scikit_learn~=0.24.2
 pytest~=7.1.0
 funcy==1.17
+PyMuPDF==1.19.6