diff --git a/image_prediction/image_extractor/extractors/parsable.py b/image_prediction/image_extractor/extractors/parsable.py index 9a2d6d6..eac09e1 100644 --- a/image_prediction/image_extractor/extractors/parsable.py +++ b/image_prediction/image_extractor/extractors/parsable.py @@ -5,7 +5,7 @@ import traceback from functools import partial, lru_cache from itertools import chain, starmap, filterfalse from operator import itemgetter, truth -from typing import List +from typing import List, Iterable, Iterator import fitz from PIL import Image @@ -59,7 +59,7 @@ class ParsablePDFImageExtractor(ImageExtractor): yield from image_metadata_pairs @staticmethod - def __filter_valid_images(image_metadata_pairs): + def __filter_valid_images(image_metadata_pairs: Iterable[ImageMetadataPair]) -> Iterator[ImageMetadataPair]: def validate(image: Image.Image, metadata: dict): try: # TODO: stand-in heuristic for testing if image is valid => find cleaner solution (RED-5148) diff --git a/image_prediction/transformer/transformers/response.py b/image_prediction/transformer/transformers/response.py index b67825f..5a37670 100644 --- a/image_prediction/transformer/transformers/response.py +++ b/image_prediction/transformer/transformers/response.py @@ -27,7 +27,7 @@ def build_image_info(data: dict) -> dict: quotient = round(compute_geometric_quotient(), 4) min_image_to_page_quotient_breached = bool(quotient < CONFIG.filters.image_to_page_quotient.min) - max_image_to_page_quotient_breached = is_max_image_to_page_quotient_breached( + max_image_to_page_quotient_breached = __is_max_image_to_page_quotient_breached( quotient, data["classification"]["label"] ) min_image_width_to_height_quotient_breached = bool( @@ -77,7 +77,7 @@ def build_image_info(data: dict) -> dict: return image_info -def is_max_image_to_page_quotient_breached(quotient, label): +def __is_max_image_to_page_quotient_breached(quotient: float, label: str) -> bool: default_max_quotient = CONFIG.filters.image_to_page_quotient.max customized_entries = CONFIG.filters.image_to_page_quotient.customized.max max_quotient = customized_entries.get(label, default_max_quotient) diff --git a/test/unit_tests/response_transformer_test.py b/test/unit_tests/response_transformer_test.py index 59fb72d..d8c58a7 100644 --- a/test/unit_tests/response_transformer_test.py +++ b/test/unit_tests/response_transformer_test.py @@ -1,6 +1,6 @@ import pytest -from image_prediction.transformer.transformers.response import is_max_image_to_page_quotient_breached +from image_prediction.transformer.transformers.response import __is_max_image_to_page_quotient_breached @pytest.fixture @@ -18,4 +18,4 @@ def expected_is_breached(quotient, label): @pytest.mark.parametrize("quotient", [0.1, 0.5]) @pytest.mark.parametrize("label", ["logo", "signature"]) def test_customized_per_label_ratio_breach(quotient, label, expected_is_breached): - assert is_max_image_to_page_quotient_breached(quotient, label) == expected_is_breached + assert __is_max_image_to_page_quotient_breached(quotient, label) == expected_is_breached