image-classification-service/image_prediction/response.py

"""Defines functions for constructing service responses."""


import math
from itertools import starmap
from operator import itemgetter

from image_prediction.config import CONFIG


def build_response(predictions: list, metadata: list) -> list:
    return list(starmap(build_image_info, zip(predictions, metadata)))


def build_image_info(prediction: dict, metadata: dict) -> dict:
    def compute_geometric_quotient():
        page_area_sqrt = math.sqrt(abs(page_width * page_height))
        image_area_sqrt = math.sqrt(abs(x2 - x1) * abs(y2 - y1))
        return image_area_sqrt / page_area_sqrt

    page_width, page_height, x1, x2, y1, y2, width, height = itemgetter(
        "page_width", "page_height", "x1", "x2", "y1", "y2", "width", "height"
    )(metadata)

    quotient = compute_geometric_quotient()

    min_image_to_page_quotient_breached = bool(quotient < CONFIG.filters.image_to_page_quotient.min)
    max_image_to_page_quotient_breached = bool(quotient > CONFIG.filters.image_to_page_quotient.max)
    min_image_width_to_height_quotient_breached = bool(
        width / height < CONFIG.filters.image_width_to_height_quotient.min
    )
    max_image_width_to_height_quotient_breached = bool(
        width / height > CONFIG.filters.image_width_to_height_quotient.max
    )

    min_confidence_breached = bool(max(prediction["probabilities"].values()) < CONFIG.filters.min_confidence)
    prediction["label"] = prediction.pop("class")  # "class" as field name causes problem for Java objectmapper
    prediction["probabilities"] = {klass: round(prob, 6) for klass, prob in prediction["probabilities"].items()}

    image_info = {
        "classification": prediction,
        "position": {"x1": x1, "x2": x2, "y1": y1, "y2": y2, "pageNumber": metadata["page_idx"] + 1},
        "geometry": {"width": width, "height": height},
        "filters": {
            "geometry": {
                "imageSize": {
                    "quotient": quotient,
                    "tooLarge": max_image_to_page_quotient_breached,
                    "tooSmall": min_image_to_page_quotient_breached,
                },
                "imageFormat": {
                    "quotient": width / height,
                    "tooTall": min_image_width_to_height_quotient_breached,
                    "tooWide": max_image_width_to_height_quotient_breached,
                },
            },
            "probability": {"unconfident": min_confidence_breached},
            "allPassed": not any(
                [
                    max_image_to_page_quotient_breached,
                    min_image_to_page_quotient_breached,
                    min_image_width_to_height_quotient_breached,
                    max_image_width_to_height_quotient_breached,
                    min_confidence_breached,
                ]
            ),
        },
    }

    return image_info