cv-analysis-service/cv_analysis/pyinfra_compat.py

import gzip
from typing import Callable

from funcy import lmap
from pyinfra.server.utils import make_streamable_and_wrap_in_packing_logic

from cv_analysis.figure_detection import detect_figures
from cv_analysis.layout_parsing import parse_layout
from cv_analysis.redaction_detection import find_redactions
from cv_analysis.table_parsing import parse_tables
from cv_analysis.utils.preprocessing import open_img_from_bytes


def make_streamable_analysis_fn(analysis_fn: Callable):
    """Makes an analysis function streamable for pyinfra server logic. The wrapped function then
    works with data and metadata and returns a tuple or generator of tuples with data and metadata.
    For more information about the server logic, see the PyInfra documentation.

    Args:
        analysis_fn: cv-analysis function

    Returns:
        wrapped function
    """

    def analyse(data, metadata: dict):
        def format_results():
            return {
                **metadata,
                "pageWidth": image.shape[1],
                "pageHeight": image.shape[0],
                "cells": results,
            }

        image = open_img_from_bytes(gzip.decompress(data))
        results = lmap(lambda x: x.json_xywh(), analysis_fn(image))
        results_metadata = format_results() if results else {}
        return b"", results_metadata

    return make_streamable_and_wrap_in_packing_logic(analyse, batched=False)


def get_analysis_fn(operation):
    return {
        "table_parsing": parse_tables,
        "layout_parsing": parse_layout,
        "figure_detection": detect_figures,
        "redaction_detection": find_redactions,
    }[operation]