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

from itertools import combinations, starmap, product
from typing import Iterable


def is_near_enough(rect_pair, max_gap=14):
    x1, y1, w1, h1 = rect_pair[0]
    x2, y2, w2, h2 = rect_pair[1]

    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_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]
    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_position1(rect_pair):
    x1, y1, w1, h1 = rect_pair[0]
    x2, y2, w2, h2 = rect_pair[1]
    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(rect_pair):
    return (is_near_enough(rect_pair) and has_correct_position1(rect_pair)) or is_overlapping(
        rect_pair)


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 rects_not_the_same(r):
    return r[0] != r[1]


def find_related_rects(rects):
    rect_pairs = list(filter(is_related, combinations(rects, 2)))
    rect_pairs = list(filter(rects_not_the_same, rect_pairs))
    if not rect_pairs:
        return [], rects
    rel_rects = list(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_rects(rects):
    rects_to_connect, rects_new = find_related_rects(rects)

    while len(rects_to_connect) > 0:
        rects_fused = list(starmap(fuse_rects, rects_to_connect))
        rects_fused = list(dict.fromkeys(rects_fused))

        if len(rects_fused) == 1:
            rects_new += rects_fused
            rects_fused = []

        rects_to_connect, connected_rects = find_related_rects(rects_fused)
        rects_new += connected_rects

        if len(rects_to_connect) > 1 and len(set(rects_to_connect)) == 1:
            rects_new.append(rects_fused[0])
            rects_to_connect = []

    return rects_new


def connect_related_rects2(rects: Iterable[tuple]):
    rects = list(rects)
    current_idx = 0

    while True:
        if current_idx + 1 >= len(rects) or len(rects) <= 1:
            break
        merge_happened = False
        current_rect = rects.pop(current_idx)
        for idx, maybe_related_rect in enumerate(rects):
            if is_related((current_rect, maybe_related_rect)):
                current_rect = fuse_rects(current_rect, maybe_related_rect)
                rects.pop(idx)
                merge_happened = True
                break
        rects.insert(0, current_rect)
        if not merge_happened:
            current_idx += 1
        elif merge_happened:
            current_idx = 0

    return rects