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

# See https://stackoverflow.com/a/39757388/3578468
from __future__ import annotations

from operator import attrgetter
from typing import TYPE_CHECKING, Iterable

from funcy import juxt, rpartial, compose, lflatten

from cv_analysis.utils import lift

if TYPE_CHECKING:
    from cv_analysis.utils.structures import Rectangle


def adjacent(alpha: Rectangle, beta: Rectangle, tolerance=7):
    """Check if the two rectangles are adjacent to each other."""
    return any(
        juxt(
            # +---+
            # |   | +---+
            # | a | | b |
            # |   | +___+
            # +___+
            right_left_aligned_and_vertically_overlapping,
            #       +---+
            # +---+ |   |
            # | b | | a |
            # +___+ |   |
            #       +___+
            left_right_aligned_and_vertically_overlapping,
            # +-----------+
            # |     a     |
            # +___________+
            #    +-----+
            #    |  b  |
            #    +_____+
            bottom_top_aligned_and_horizontally_overlapping,
            #    +-----+
            #    |  b  |
            #    +_____+
            # +-----------+
            # |     a     |
            # +___________+
            top_bottom_aligned_and_horizontally_overlapping,
        )(alpha, beta, tolerance)
    )


def right_left_aligned_and_vertically_overlapping(alpha: Rectangle, beta: Rectangle, tol):
    """Check if the first rectangle is left of the other within a tolerance and also overlaps the other's y range."""
    return adjacent_along_one_axis_and_overlapping_along_perpendicular_axis(
        alpha.x2, beta.x1, beta.y1, beta.y2, alpha.y1, alpha.y2, tolerance=tol
    )


def left_right_aligned_and_vertically_overlapping(alpha: Rectangle, beta: Rectangle, tol):
    """Check if the first rectangle is right of the other within a tolerance and also overlaps the other's y range."""
    return adjacent_along_one_axis_and_overlapping_along_perpendicular_axis(
        alpha.x1, beta.x2, beta.y1, beta.y2, alpha.y1, alpha.y2, tolerance=tol
    )


def bottom_top_aligned_and_horizontally_overlapping(alpha: Rectangle, beta: Rectangle, tol):
    """Check if the first rectangle is above the other within a tolerance and also overlaps the other's x range."""
    return adjacent_along_one_axis_and_overlapping_along_perpendicular_axis(
        alpha.y2, beta.y1, beta.x1, beta.x2, alpha.x1, alpha.x2, tolerance=tol
    )


def top_bottom_aligned_and_horizontally_overlapping(alpha: Rectangle, beta: Rectangle, tol):
    """Check if the first rectangle is below the other within a tolerance and also overlaps the other's x range."""
    return adjacent_along_one_axis_and_overlapping_along_perpendicular_axis(
        alpha.y1, beta.y2, beta.x1, beta.x2, alpha.x1, alpha.x2, tolerance=tol
    )


def adjacent_along_one_axis_and_overlapping_along_perpendicular_axis(
    axis_0_point_1,
    axis_1_point_2,
    axis_1_contained_point_1,
    axis_1_contained_point_2,
    axis_1_lower_bound,
    axis_1_upper_bound,
    tolerance,
):
    """Check if two points are adjacent along one axis and two other points overlap a range along the perpendicular
    axis."""
    return all(
        [
            abs(axis_0_point_1 - axis_1_point_2) <= tolerance,
            any(
                [
                    axis_1_lower_bound <= p <= axis_1_upper_bound
                    for p in [axis_1_contained_point_1, axis_1_contained_point_2]
                ]
            ),
        ]
    )


def contains(alpha: Rectangle, beta: Rectangle, tol=3):
    """Check if the first rectangle contains the second rectangle."""
    return (
        beta.x1 + tol >= alpha.x1
        and beta.y1 + tol >= alpha.y1
        and beta.x2 - tol <= alpha.x2
        and beta.y2 - tol <= alpha.y2
    )


def is_contained(rectangle: Rectangle, rectangles: Iterable[Rectangle]):
    """Check if the rectangle is contained within any of the other rectangles."""
    other_rectangles = filter(lambda r: r != rectangle, rectangles)
    return any(map(rpartial(contains, rectangle), other_rectangles))


def intersection(alpha, beta):
    return intersection_along_x_axis(alpha, beta) * intersection_along_y_axis(alpha, beta)


def intersection_along_x_axis(alpha, beta):
    return intersection_along_axis(alpha, beta, "x")


def intersection_along_y_axis(alpha, beta):
    return intersection_along_axis(alpha, beta, "y")


def intersection_along_axis(alpha, beta, axis):
    assert axis in ["x", "y"]

    def get_component_accessor(component):
        return attrgetter(f"{axis}{component}")

    def make_access_components_and_sort_fn(component):
        return compose(sorted, lift(get_component_accessor(component)))

    c1 = make_access_components_and_sort_fn(1)
    c2 = make_access_components_and_sort_fn(2)

    # Cases:
    #    a      b
    # [-----] (---)  => [a1, b1, a2, b2] => max(0, (a2 - b1)) = 0
    #   b      a
    # (---) [-----]  => [b1, a1, b2, a2] => max(0, (b2 - a1)) = 0
    #     a   b
    # [--(----]----) => [a1, b1, a2, b2] => max(0, (a2 - b1)) = (a2 - b1)
    #     a b
    # (-[---]----)   => [b1, a1, a2, b2] => max(0, (a2 - a1)) = (a2 - a1)
    #     b a
    # [-(---)----]   => [a1, b1, b2, a2] => max(0, (b2 - b1)) = (b2 - b1)
    #     b   a
    # (----[--)----] => [b1, a1, b2, a2] => max(0, (b2 - a1)) = (b2 - a1)

    coords = lflatten(juxt(c1, c2)((alpha, beta)))
    intersection = max(0, coords[2] - coords[1])
    return intersection