33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
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
|