diff --git a/cv_analysis/utils/connect_rects.py b/cv_analysis/utils/connect_rects.py index abc1fd2..21de49d 100644 --- a/cv_analysis/utils/connect_rects.py +++ b/cv_analysis/utils/connect_rects.py @@ -1,15 +1,17 @@ -from itertools import combinations, starmap, product +from itertools import combinations, product from typing import Iterable from funcy import lfilter, lmap +from cv_analysis.utils import star from cv_analysis.utils.conversion import rectangle_to_box, box_to_rectangle from cv_analysis.utils.rectangle import Rectangle +from cv_analysis.utils.spacial import overlap -def is_near_enough(rect_pair, max_gap=14): - x1, y1, w1, h1 = rect_pair[0] - x2, y2, w2, h2 = rect_pair[1] +def is_near_enough(alpha: Rectangle, beta: Rectangle, max_gap=14): + x1, y1, w1, h1 = alpha + x2, y2, w2, h2 = beta return any( [ @@ -21,14 +23,6 @@ def is_near_enough(rect_pair, max_gap=14): ) -def is_overlapping(rect_pair): - x1, y1, w1, h1 = rect_pair[0] - x2, y2, w2, h2 = rect_pair[1] - dx = min(x1 + w1, x2 + w2) - max(x1, x2) - dy = min(y1 + h1, y2 + h2) - max(y1, y2) - return True if (dx >= 0) and (dy >= 0) else False - - def is_on_same_line(rect_pair): x1, y1, w1, h1 = rect_pair[0] x2, y2, w2, h2 = rect_pair[1] @@ -60,8 +54,8 @@ def has_correct_position(alpha: Rectangle, beta: Rectangle): ) -def is_related(rect_pair): - return (is_near_enough(rect_pair) and has_correct_position(*rect_pair)) or is_overlapping(rect_pair) +def is_related(alpha: Rectangle, beta: Rectangle): + return (is_near_enough(alpha, beta) and has_correct_position(alpha, beta)) or overlap(alpha, beta) def fuse_rects(rect1, rect2): @@ -83,7 +77,7 @@ def rectangles_differ(r): def find_related_rects(rects): - rect_pairs = lfilter(is_related, combinations(rects, 2)) + rect_pairs = lfilter(star(is_related), combinations(rects, 2)) rect_pairs = lfilter(rectangles_differ, rect_pairs) if not rect_pairs: return [], rects @@ -104,7 +98,7 @@ def connect_related_rectangles(rectangles: Iterable[Rectangle]): current_rect = boxes.pop(current_idx) for idx, maybe_related_rect in enumerate(boxes): - if is_related((current_rect, maybe_related_rect)): + if is_related(current_rect, maybe_related_rect): current_rect = fuse_rects(current_rect, maybe_related_rect) boxes.pop(idx) merge_happened = True diff --git a/cv_analysis/utils/spacial.py b/cv_analysis/utils/spacial.py index 2a419c6..3d02340 100644 --- a/cv_analysis/utils/spacial.py +++ b/cv_analysis/utils/spacial.py @@ -179,3 +179,12 @@ def intersection_along_axis(alpha, beta, axis): 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 overlap(alpha: Rectangle, beta: Rectangle): + x1, y1, w1, h1 = alpha + x2, y2, w2, h2 = beta + + dx = min(x1 + w1, x2 + w2) - max(x1, x2) + dy = min(y1 + h1, y2 + h2) - max(y1, y2) + return True if (dx >= 0) and (dy >= 0) else False diff --git a/cv_analysis/utils/utils.py b/cv_analysis/utils/utils.py index e62769b..a7f9c4d 100644 --- a/cv_analysis/utils/utils.py +++ b/cv_analysis/utils/utils.py @@ -26,3 +26,10 @@ def lift(fn): yield from map(fn, coll) return lifted + + +def star(fn): + def starred(*args): + return fn(*args) + + return starred