diff --git a/cv_analysis/utils/connect_rects.py b/cv_analysis/utils/connect_rects.py index 7522fc0..6cc4613 100644 --- a/cv_analysis/utils/connect_rects.py +++ b/cv_analysis/utils/connect_rects.py @@ -1,10 +1,7 @@ -from itertools import combinations, product -from typing import Iterable +from itertools import combinations +from typing import List -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 import until from cv_analysis.utils.rectangle import Rectangle from cv_analysis.utils.spacial import overlap @@ -71,38 +68,35 @@ 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: List[Rectangle]): + assert isinstance(rectangles, list) + no_new_merges = make_merger_sentinel() + return until(no_new_merges, merge_rectangles_once, rectangles) -def connect_related_rectangles(rectangles: Iterable[Rectangle]): - rectangles = list(rectangles) - - current_idx = 0 - - while True: - if current_idx + 1 >= len(rectangles) or len(rectangles) <= 1: - break - merge_happened = False - current_rect = rectangles.pop(current_idx) - - for idx, maybe_related_rect in enumerate(rectangles): - if is_related(current_rect, maybe_related_rect): - current_rect = bounding_rect(current_rect, maybe_related_rect) - - rectangles.pop(idx) - merge_happened = True - break - rectangles.insert(0, current_rect) - if not merge_happened: - current_idx += 1 - elif merge_happened: - current_idx = 0 - +def merge_rectangles_once(rectangles: List[Rectangle]): + for alpha, beta in combinations(rectangles, 2): + if is_related(alpha, beta): + rectangles.remove(alpha) + rectangles.remove(beta) + rectangles.append(bounding_rect(alpha, beta)) + return rectangles return rectangles + + +def make_merger_sentinel(): + def no_new_mergers(records): + nonlocal number_of_records_so_far + + number_of_records_now = len(records) + + if number_of_records_now == number_of_records_so_far: + return True + + else: + number_of_records_so_far = number_of_records_now + return False + + number_of_records_so_far = -1 + + return no_new_mergers diff --git a/cv_analysis/utils/utils.py b/cv_analysis/utils/utils.py index a7f9c4d..6e1d697 100644 --- a/cv_analysis/utils/utils.py +++ b/cv_analysis/utils/utils.py @@ -1,6 +1,7 @@ from __future__ import annotations import cv2 +from funcy import first, iterate from numpy import generic @@ -33,3 +34,7 @@ def star(fn): return fn(*args) return starred + + +def until(cond, func, *args, **kwargs): + return first(filter(cond, iterate(func, *args, **kwargs)))