cv-analysis-service/cv_analysis/utils/postprocessing.py

from functools import partial
from itertools import starmap, compress
from typing import Iterable, List, Sequence

from funcy import lremove

from cv_analysis.utils.rectangle import Rectangle


def remove_overlapping(rectangles: Iterable[Rectangle]) -> List[Rectangle]:
    def overlap(a: Rectangle, rect2: Rectangle) -> float:
        return a.intersection(rect2) > 0

    def does_not_overlap(rect: Rectangle, rectangles: Iterable[Rectangle]) -> bool:
        return not any(overlap(rect, rect2) for rect2 in rectangles if not rect == rect2)

    rectangles = list(filter(partial(does_not_overlap, rectangles=rectangles), rectangles))
    return rectangles


def remove_included(rectangles: Iterable[Rectangle]) -> List[Rectangle]:
    rectangles_to_keep = [rect for rect in rectangles if not rect.is_included(rectangles)]
    return rectangles_to_keep


def remove_small(boxes: Iterable[Rectangle], page_width, page_height, min_percentage=0.13) -> List[Rectangle]:
    min_width = page_width * min_percentage
    min_height = page_height * min_percentage

    def small(box: Rectangle):
        return box.width < min_width or box.height < min_height

    return lremove(small, boxes)


def __remove_isolated_unsorted(rectangles: Iterable[Rectangle]) -> List[Rectangle]:
    def is_connected(rect: Rectangle, rectangles: Iterable[Rectangle]):
        return any(rect.adjacent(rect2) for rect2 in rectangles if not rect == rect2)

    if not isinstance(rectangles, list):
        rectangles = list(rectangles)

    rectangles = list(filter(partial(is_connected, rectangles=rectangles), rectangles))
    return rectangles


def __remove_isolated_sorted(rectangles: Iterable[Rectangle]) -> List[Rectangle]:
    def is_connected(left, center, right):
        return any([left.adjacent(center), center.adjacent(right)])

    rectangles = list(rectangles)
    lefts = [None, *rectangles[:-1]]
    rights = [*rectangles[1:], None]
    mask = starmap(is_connected, zip(lefts, rectangles, rights))

    rectangles = list(compress(rectangles, mask))
    return rectangles


def remove_isolated(rectangles: Iterable[Rectangle], input_unsorted: bool = True) -> List[Rectangle]:
    return (__remove_isolated_unsorted if input_unsorted else __remove_isolated_sorted)(rectangles)


def has_no_parent(hierarchy: Sequence[int]) -> bool:
    return hierarchy[-1] <= 0