diff --git a/image_prediction/image_extractor/extractors/parsable.py b/image_prediction/image_extractor/extractors/parsable.py index caa5c97..2d79fac 100644 --- a/image_prediction/image_extractor/extractors/parsable.py +++ b/image_prediction/image_extractor/extractors/parsable.py @@ -12,7 +12,67 @@ from image_prediction.image_extractor.extractor import ImageExtractor, ImageMeta from image_prediction.info import Info from image_prediction.stitching.stitching import stitch_pairs -rounder = rcompose(round, int) + +class ParsablePDFImageExtractor(ImageExtractor): + def __init__(self, verbose=False): + self.doc: fitz.fitz.Document = None + self.verbose = verbose + + def extract(self, pdf: bytes, page_range: range = None): + self.doc = fitz.Document(stream=pdf) + + pages = extract_pages(self.doc, page_range) if page_range else self.doc + + image_metadata_pairs = chain.from_iterable( + map( + self.__process_images_on_page, + tqdm(pages, desc="Extracting", disable=not self.verbose, total=len(page_range) if page_range else None), + ) + ) + + return image_metadata_pairs + + def __process_images_on_page(self, page: fitz.fitz.Page): + images = get_images_on_page(self.doc, page) + metadata = get_metadata_for_images_on_page(page) + get_image_infos.cache_clear() + + image_metadata_pairs = starmap( + ImageMetadataPair, filter(compose(all, curry(map)(truth)), zip(images, metadata)) + ) + image_metadata_pairs = stitch_pairs(list(image_metadata_pairs)) + + return image_metadata_pairs + + +def extract_pages(doc, page_range): + page_range = range(page_range.start + 1, page_range.stop + 1) + pages = map(doc.load_page, page_range) + return pages + + +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 + + +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_image_metadata(image_info): @@ -38,63 +98,4 @@ def get_page_metadata(page): } -def extract_pages(doc, page_range): - page_range = range(page_range.start + 1, page_range.stop + 1) - pages = map(doc.load_page, 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): - images = get_images_on_page(self.doc, page) - metadata = get_metadata_for_images_on_page(page) - get_image_infos.cache_clear() - - image_metadata_pairs = starmap( - ImageMetadataPair, filter(compose(all, curry(map)(truth)), zip(images, metadata)) - ) - image_metadata_pairs = stitch_pairs(list(image_metadata_pairs)) - - return image_metadata_pairs - - def extract(self, pdf: bytes, page_range: range = None): - self.doc = fitz.Document(stream=pdf) - - pages = extract_pages(self.doc, page_range) if page_range else self.doc - - image_metadata_pairs = chain.from_iterable( - map( - self.__process_images_on_page, - tqdm(pages, desc="Extracting", disable=not self.verbose, total=len(page_range) if page_range else None), - ) - ) - - return image_metadata_pairs +rounder = rcompose(round, int)