import math from operator import itemgetter from image_prediction.config import CONFIG from image_prediction.transformer.transformer import Transformer from image_prediction.utils import get_logger logger = get_logger() class ResponseTransformer(Transformer): def transform(self, data): logger.debug("ResponseTransformer.transform") return build_image_info(data) def build_image_info(data: 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 def is_max_image_to_page_quotient_breached(quotient, label): 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) max_quotient = max_quotient if max_quotient else default_max_quotient return bool(quotient > max_quotient) page_width, page_height, x1, x2, y1, y2, width, height, alpha = itemgetter( "page_width", "page_height", "x1", "x2", "y1", "y2", "width", "height", "alpha" )(data) 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( quotient, data["classification"]["label"] ) 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 ) classification = data["classification"] representation = data["representation"] min_confidence_breached = bool(max(classification["probabilities"].values()) < CONFIG.filters.min_confidence) image_info = { "classification": classification, "representation": representation, "position": {"x1": x1, "x2": x2, "y1": y1, "y2": y2, "pageNumber": data["page_idx"] + 1}, "geometry": {"width": width, "height": height}, "alpha": alpha, "filters": { "geometry": { "imageSize": { "quotient": quotient, "tooLarge": max_image_to_page_quotient_breached, "tooSmall": min_image_to_page_quotient_breached, }, "imageFormat": { "quotient": round(width / height, 4), "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