From d80af336eb4ae6932f94d1fbe89e32339856f2bc Mon Sep 17 00:00:00 2001 From: Matthias Bisping Date: Mon, 11 Apr 2022 13:28:39 +0200 Subject: [PATCH] refactoring --- .../image_extractor/extractors/parsable.py | 38 ++++++++++++++----- 1 file changed, 28 insertions(+), 10 deletions(-) diff --git a/image_prediction/image_extractor/extractors/parsable.py b/image_prediction/image_extractor/extractors/parsable.py index c5892ca..d601d45 100644 --- a/image_prediction/image_extractor/extractors/parsable.py +++ b/image_prediction/image_extractor/extractors/parsable.py @@ -1,5 +1,5 @@ import io -from functools import partial +from functools import partial, lru_cache from itertools import chain, starmap from operator import itemgetter, truth @@ -43,21 +43,39 @@ def extract_pages(doc, page_range): return pages +def load_image_from_xref(doc, xref): + maybe_image = doc.extract_image(xref) + return Image.open(io.BytesIO(maybe_image["image"])) if maybe_image else None + + +@lru_cache(maxsize=None) +def get_image_infos(page: fitz.Page): + return page.get_image_info(xrefs=True) + + +def get_images_on_page(doc, page: fitz.Page): + image_infos = get_image_infos(page) + xrefs = map(itemgetter("xref"), image_infos) + images = map(partial(load_image_from_xref, doc), xrefs) + return images + + +def get_metadata_for_images_on_page(page: fitz.Page): + image_infos = get_image_infos(page) + metadata = map(get_image_metadata, image_infos) + metadata = map(partial(merge, get_page_metadata(page)), metadata) + return metadata + + class ParsablePDFImageExtractor(ImageExtractor): def __init__(self, verbose=False): self.doc: fitz.fitz.Document = None self.verbose = verbose def __process_images_on_page(self, page: fitz.fitz.Page): - def load_image_from_xref(xref): - maybe_image = self.doc.extract_image(xref) - return Image.open(io.BytesIO(maybe_image["image"])) if maybe_image else None - - image_infos = page.get_image_info(xrefs=True) - xrefs = map(itemgetter("xref"), image_infos) - images = map(load_image_from_xref, xrefs) - metadata = map(get_image_metadata, image_infos) - metadata = map(partial(merge, get_page_metadata(page)), metadata) + images = get_images_on_page(self.doc, page) + metadata = get_metadata_for_images_on_page(page) + get_image_infos.cache_clear() return starmap(ImageMetadataPair, filter(compose(all, curry(map)(truth)), zip(images, metadata)))