66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
from itertools import groupby, chain
|
|
from typing import Iterable, List
|
|
|
|
from funcy import compose, second
|
|
|
|
from image_prediction.image_extractor.extractor import ImageMetadataPair
|
|
from image_prediction.stitcher.utils import make_coord_getter, make_group_merger
|
|
from image_prediction.utils.generic import until
|
|
|
|
|
|
def group_by_coordinate(pairs, coord_getter):
|
|
pairs = sorted(pairs, key=coord_getter)
|
|
return map(compose(list, second), groupby(pairs, coord_getter))
|
|
|
|
|
|
class CoordGrouper:
|
|
def __init__(self, axis):
|
|
|
|
self.c1_getter = make_coord_getter(f"{other_axis(axis)}1")
|
|
self.c2_getter = make_coord_getter(f"{other_axis(axis)}2")
|
|
|
|
def get_c1(self, pair: ImageMetadataPair):
|
|
return self.c1_getter(pair)
|
|
|
|
def group_pairs_by_c1(self, pairs):
|
|
return group_by_coordinate(pairs, self.c1_getter)
|
|
|
|
def group_pairs_by_c2(self, pairs):
|
|
return group_by_coordinate(pairs, self.c2_getter)
|
|
|
|
|
|
def other_axis(axis):
|
|
return "y" if axis == "x" else "x"
|
|
|
|
|
|
class Stitcher:
|
|
def merge_along_axis(self, pairs, axis):
|
|
def group_pairs_within_groups_by_c2(groups):
|
|
return map(grouper.group_pairs_by_c2, groups)
|
|
|
|
def merge_groups_along_orthogonal_axis(groups):
|
|
return map(group_merger, groups)
|
|
|
|
grouper = CoordGrouper(axis)
|
|
group_merger = make_group_merger(axis)
|
|
|
|
groups_of_pairs_with_same_c1 = grouper.group_pairs_by_c1(pairs)
|
|
groups_of_groups_of_pairs_with_same_c1_and_c2 = group_pairs_within_groups_by_c2(groups_of_pairs_with_same_c1)
|
|
groups_of_pairs_with_matching_c1_and_c2 = chain(*groups_of_groups_of_pairs_with_same_c1_and_c2)
|
|
groups_of_merged_pairs = merge_groups_along_orthogonal_axis(groups_of_pairs_with_matching_c1_and_c2)
|
|
pairs = chain(*groups_of_merged_pairs)
|
|
|
|
return pairs
|
|
|
|
def merge_along_both_axes(self, pairs):
|
|
pairs = self.merge_along_axis(pairs, "x")
|
|
pairs = list(self.merge_along_axis(pairs, "y"))
|
|
|
|
return pairs
|
|
|
|
def stitch(self, pairs: Iterable[ImageMetadataPair]) -> List[ImageMetadataPair]:
|
|
def break_condition(pairs1, pairs2):
|
|
return len(pairs1) == len(pairs2)
|
|
|
|
return until(self.merge_along_both_axes, break_condition, pairs)
|