287 lines
10 KiB
Python

# See https://stackoverflow.com/a/39757388
from __future__ import annotations
from functools import lru_cache
from operator import attrgetter
from typing import TYPE_CHECKING, Iterable
from funcy import juxt, rpartial, compose, lflatten, first, second
from cv_analysis.utils import lift
if TYPE_CHECKING:
from cv_analysis.utils.rectangle import Rectangle
def adjacent(alpha: Rectangle, beta: Rectangle, tolerance=7, strict=False):
"""Checks if the two rectangles are adjacent to each other.
Args:
alpha: The first rectangle.
beta: The second rectangle.
tolerance: The maximum distance between the two rectangles.
strict: If True, the rectangles must be adjacent along one axis and contained within the other axis. Else, the
rectangles must be adjacent along one axis and overlapping the other axis.
Returns:
True if the two rectangles are adjacent to each other, False otherwise.
"""
select_strictness_variant = first if strict else second
test_candidates = [
# +---+
# | | +---+
# | a | | b |
# | | +___+
# +___+
(right_left_aligned_and_vertically_contained, right_left_aligned_and_vertically_overlapping),
# +---+
# +---+ | |
# | b | | a |
# +___+ | |
# +___+
(left_right_aligned_and_vertically_contained, left_right_aligned_and_vertically_overlapping),
# +-----------+
# | a |
# +___________+
# +-----+
# | b |
# +_____+
(bottom_top_aligned_and_horizontally_contained, bottom_top_aligned_and_horizontally_overlapping),
# +-----+
# | b |
# +_____+
# +-----------+
# | a |
# +___________+
(top_bottom_aligned_and_horizontally_contained, top_bottom_aligned_and_horizontally_overlapping),
]
tests = map(select_strictness_variant, test_candidates)
return any(juxt(*tests)(alpha, beta, tolerance))
def right_left_aligned_and_vertically_overlapping(alpha: Rectangle, beta: Rectangle, tol):
"""Checks 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):
"""Checks 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):
"""Checks 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):
"""Checks 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 right_left_aligned_and_vertically_contained(alpha: Rectangle, beta: Rectangle, tol):
"""Checks if the first rectangle is left of the other within a tolerance and also contains the other's y range."""
return adjacent_along_one_axis_and_contained_within_perpendicular_axis(
alpha.x2, beta.x1, beta.y1, beta.y2, alpha.y1, alpha.y2, tolerance=tol
)
def left_right_aligned_and_vertically_contained(alpha: Rectangle, beta: Rectangle, tol):
"""Checks if the first rectangle is right of the other within a tolerance and also contains the other's y range."""
return adjacent_along_one_axis_and_contained_within_perpendicular_axis(
alpha.x1, beta.x2, beta.y1, beta.y2, alpha.y1, alpha.y2, tolerance=tol
)
def bottom_top_aligned_and_horizontally_contained(alpha: Rectangle, beta: Rectangle, tol):
"""Checks if the first rectangle is above the other within a tolerance and also contains the other's x range."""
return adjacent_along_one_axis_and_contained_within_perpendicular_axis(
alpha.y2, beta.y1, beta.x1, beta.x2, alpha.x1, alpha.x2, tolerance=tol
)
def top_bottom_aligned_and_horizontally_contained(alpha: Rectangle, beta: Rectangle, tol):
"""Checks if the first rectangle is below the other within a tolerance and also contains the other's x range."""
return adjacent_along_one_axis_and_contained_within_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,
):
"""Checks if two points are adjacent along one axis and two other points overlap a range along the perpendicular
axis.
"""
return adjacent_along_one_axis_and_overlapping_or_contained_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,
mode="overlapping",
)
def adjacent_along_one_axis_and_contained_within_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,
):
"""Checks if two points are adjacent along one axis and two other points overlap a range along the perpendicular
axis.
"""
return adjacent_along_one_axis_and_overlapping_or_contained_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,
mode="contained",
)
def adjacent_along_one_axis_and_overlapping_or_contained_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,
mode,
):
"""Checks if two points are adjacent along one axis and two other points overlap a range along the perpendicular
axis or are contained in that range, depending on the mode specified.
"""
assert mode in ["overlapping", "contained"]
quantifier = any if mode == "overlapping" else all
return all(
[
abs(axis_0_point_1 - axis_1_point_2) <= tolerance,
quantifier(
[
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):
"""Checks 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]):
"""Checks 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 iou(alpha: Rectangle, beta: Rectangle):
"""Calculates the intersection area over the union area of two rectangles."""
return intersection(alpha, beta) / union(alpha, beta)
def area(rectangle: Rectangle):
"""Calculates the area of a rectangle."""
return abs((rectangle.x2 - rectangle.x1) * (rectangle.y2 - rectangle.y1))
def union(alpha: Rectangle, beta: Rectangle):
"""Calculates the union area of two rectangles."""
return area(alpha) + area(beta) - intersection(alpha, beta)
@lru_cache(maxsize=1000)
def intersection(alpha, beta):
"""Calculates the intersection of two rectangles."""
return intersection_along_x_axis(alpha, beta) * intersection_along_y_axis(alpha, beta)
def intersection_along_x_axis(alpha, beta):
"""Calculates the intersection along the x-axis."""
return intersection_along_axis(alpha, beta, "x")
def intersection_along_y_axis(alpha, beta):
"""Calculates the intersection along the y-axis."""
return intersection_along_axis(alpha, beta, "y")
def intersection_along_axis(alpha, beta, axis):
"""Calculates the intersection along the given axis.
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)
"""
assert axis in ["x", "y"]
def get_component_accessor(component):
"""Returns a function that accesses the given component of a rectangle."""
return attrgetter(f"{axis}{component}")
def make_access_components_and_sort_fn(component):
"""Returns a function that accesses and sorts the given component of multiple rectangles."""
assert component in [1, 2]
return compose(sorted, lift(get_component_accessor(component)))
sort_first_components, sort_second_components = map(make_access_components_and_sort_fn, [1, 2])
min_c1, max_c1, min_c2, max_c2 = lflatten(juxt(sort_first_components, sort_second_components)((alpha, beta)))
intersection = max(0, min_c2 - max_c1)
return intersection
def related(alpha: Rectangle, beta: Rectangle):
return close(alpha, beta) or overlap(alpha, beta)
def close(alpha: Rectangle, beta: Rectangle, max_gap=14):
# FIXME: Parameterize via factory
return adjacent(alpha, beta, tolerance=max_gap, strict=True)
def overlap(alpha: Rectangle, beta: Rectangle):
return intersection(alpha, beta) > 0