diff --git a/image_prediction/image_extractor/extractors/parsable.py b/image_prediction/image_extractor/extractors/parsable.py index 885cbae..d115059 100644 --- a/image_prediction/image_extractor/extractors/parsable.py +++ b/image_prediction/image_extractor/extractors/parsable.py @@ -1,22 +1,18 @@ import atexit -import json -import traceback from _operator import itemgetter from functools import partial, lru_cache -from itertools import chain, starmap, filterfalse +from itertools import chain, filterfalse from operator import itemgetter -from typing import Iterable, Iterator, List, Union +from typing import List, Union -import IPython import fitz import numpy as np from PIL import Image -from funcy import merge, compose, rcompose, keep, identity +from funcy import merge, compose, rcompose, keep from pymonad.maybe import Maybe, Nothing, Just from pymonad.tools import curry from image_prediction.exceptions import InvalidBox, BadXref -from image_prediction.formatter.formatters.enum import EnumFormatter from image_prediction.image_extractor.extractor import ImageExtractor, ImageMetadataPair from image_prediction.info import Info from image_prediction.stitching.stitching import stitch_pairs @@ -25,8 +21,6 @@ from image_prediction.utils import get_logger logger = get_logger() -Maybe - class ParsablePDFImageExtractor(ImageExtractor): def __init__(self, verbose=False, tolerance=0): @@ -57,29 +51,10 @@ class ParsablePDFImageExtractor(ImageExtractor): image_metadata_pairs = [pair.value for pair in maybe_image_metadata_pairs if pair.is_just()] clear_caches() - # TODO: In the future, consider to introduce an image validator as a pipeline component rather than doing the - # validation here. Invalid images can then be split into a different stream and joined with the intact images - # again for the formatting step. - image_metadata_pairs = self.__filter_valid_images(image_metadata_pairs) - image_metadata_pairs = stitch_pairs(list(image_metadata_pairs), tolerance=self.tolerance) yield from image_metadata_pairs - @staticmethod - def __filter_valid_images(image_metadata_pairs: Iterable[ImageMetadataPair]) -> Iterator[ImageMetadataPair]: - def validate(image: Image.Image, metadatum: dict): - try: - # TODO: stand-in heuristic for testing if image is valid => find cleaner solution (RED-5148) - image.resize((100, 100)).convert("RGB") - return ImageMetadataPair(image, metadatum) - except (OSError, Exception): - metadatum = json.dumps(EnumFormatter()(metadatum), indent=2) - logger.warning(f"Invalid image encountered. Image metadata:\n{metadatum}\n\n{traceback.format_exc()}") - return None - - return keep(starmap(validate, image_metadata_pairs)) - def extract_pages(doc, page_range): page_range = range(page_range.start + 1, page_range.stop + 1) @@ -88,6 +63,16 @@ def extract_pages(doc, page_range): yield from pages +def validate_image(image: Image.Image) -> Maybe: + try: + # TODO: stand-in heuristic for testing if image is valid => find cleaner solution (RED-5148) + image.resize((100, 100)).convert("RGB") + return Just(image) + except (OSError, Exception): + logger.warning(f"Invalid image encountered.") + return Nothing + + def extract_valid_metadata(doc: fitz.fitz.Document, page: fitz.fitz.Page): return compose( list, @@ -98,8 +83,9 @@ def extract_valid_metadata(doc: fitz.fitz.Document, page: fitz.fitz.Page): def metadatum_to_image_metadata_pair(doc, metadatum: dict) -> Maybe: - maybe_image = xref_to_maybe_image(doc, metadatum[Info.XREF]) - return make_maybe_image_metadata_pair(maybe_image, Just(metadatum)) + maybe_image = xref_to_maybe_image(doc, metadatum[Info.XREF]).bind(validate_image) + maybe_image_metadata_pair = make_maybe_image_metadata_pair(maybe_image, Just(metadatum)) + return maybe_image_metadata_pair def add_alpha_channel_info(doc, metadata):