113 lines
3.2 KiB
Python
113 lines
3.2 KiB
Python
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(alpha: Rectangle, beta: Rectangle, max_gap=14):
|
|
x1, y1, w1, h1 = alpha
|
|
x2, y2, w2, h2 = beta
|
|
|
|
return any(
|
|
[
|
|
abs(x1 - (x2 + w2)) <= max_gap,
|
|
abs(x2 - (x1 + w1)) <= max_gap,
|
|
abs(y2 - (y1 + h1)) <= max_gap,
|
|
abs(y1 - (y2 + h2)) <= max_gap,
|
|
]
|
|
)
|
|
|
|
|
|
def is_on_same_line(rect_pair):
|
|
x1, y1, w1, h1 = rect_pair[0]
|
|
x2, y2, w2, h2 = rect_pair[1]
|
|
return any(
|
|
[
|
|
any([abs(y1 - y2) <= 10, abs(y1 + h1 - (y2 + h2)) <= 10]),
|
|
any([y2 <= y1 and y1 + h1 <= y2 + h2, y1 <= y2 and y2 + h2 <= y1 + h1]),
|
|
]
|
|
)
|
|
|
|
|
|
def has_correct_position(alpha: Rectangle, beta: Rectangle):
|
|
x1, y1, w1, h1 = alpha
|
|
x2, y2, w2, h2 = beta
|
|
return any(
|
|
[
|
|
any(
|
|
[abs(x1 - x2) <= 10, abs(y1 - y2) <= 10, abs(x1 + w1 - (x2 + w2)) <= 10, abs(y1 + h1 - (y2 + h2)) <= 10]
|
|
),
|
|
any(
|
|
[
|
|
y2 <= y1 and y1 + h1 <= y2 + h2,
|
|
y1 <= y2 and y2 + h2 <= y1 + h1,
|
|
x2 <= x1 and x1 + w1 <= x2 + w2,
|
|
x1 <= x2 and x2 + w2 <= x1 + w1,
|
|
]
|
|
),
|
|
]
|
|
)
|
|
|
|
|
|
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):
|
|
if rect1 == rect2:
|
|
return rect1
|
|
x1, y1, w1, h1 = rect1
|
|
x2, y2, w2, h2 = rect2
|
|
|
|
topleft = list(min(product([x1, x2], [y1, y2])))
|
|
bottomright = list(max(product([x1 + w1, x2 + w2], [y1 + h1, y2 + h2])))
|
|
|
|
w = [bottomright[0] - topleft[0]]
|
|
h = [bottomright[1] - topleft[1]]
|
|
return tuple(topleft + w + h)
|
|
|
|
|
|
def rectangles_differ(r):
|
|
return r[0] != r[1]
|
|
|
|
|
|
def find_related_rects(rects):
|
|
rect_pairs = lfilter(star(is_related), combinations(rects, 2))
|
|
rect_pairs = lfilter(rectangles_differ, rect_pairs)
|
|
if not rect_pairs:
|
|
return [], rects
|
|
rel_rects = set([rect for pair in rect_pairs for rect in pair])
|
|
unrel_rects = [rect for rect in rects if rect not in rel_rects]
|
|
return rect_pairs, unrel_rects
|
|
|
|
|
|
def connect_related_rectangles(rectangles: Iterable[Rectangle]):
|
|
boxes = lmap(rectangle_to_box, rectangles)
|
|
|
|
current_idx = 0
|
|
|
|
while True:
|
|
if current_idx + 1 >= len(boxes) or len(boxes) <= 1:
|
|
break
|
|
merge_happened = False
|
|
current_rect = boxes.pop(current_idx)
|
|
|
|
for idx, maybe_related_rect in enumerate(boxes):
|
|
if is_related(current_rect, maybe_related_rect):
|
|
current_rect = fuse_rects(current_rect, maybe_related_rect)
|
|
boxes.pop(idx)
|
|
merge_happened = True
|
|
break
|
|
boxes.insert(0, current_rect)
|
|
if not merge_happened:
|
|
current_idx += 1
|
|
elif merge_happened:
|
|
current_idx = 0
|
|
|
|
return lmap(box_to_rectangle, boxes)
|