Python框排序

按照上下左右顺序对一系列框进行排序。

from easydict import EasyDict

def bbox4pto2p(box4p):
    box2p = [
        min(box4p[0], box4p[2], box4p[4], box4p[6]),
        min(box4p[1], box4p[3], box4p[5], box4p[7]),
        max(box4p[0], box4p[2], box4p[4], box4p[6]),
        max(box4p[1], box4p[3], box4p[5], box4p[7]),
    ]
    return box2p


class CompareBoxes(object):
    def __init__(self, coord_key:str):
        self.coord_key = coord_key

    def _structuralize(self, box1, box2):
        if len(box1) == 8:
            box1 = bbox4pto2p(box1)
        if len(box2) == 8:
            box2 = bbox4pto2p(box2)
        b1 = EasyDict()
        b2 = EasyDict()
        b1.minx = min(box1[0], box1[2])
        b1.maxx = max(box1[0], box1[2])
        b1.miny = min(box1[1], box1[3])
        b1.maxy = max(box1[1], box1[3])
        b2.minx = min(box2[0], box2[2])
        b2.maxx = max(box2[0], box2[2])
        b2.miny = min(box2[1], box2[3])
        b2.maxy = max(box2[1], box2[3])
        return b1, b2
    
    @staticmethod
    def calc_x_type(a, b):
        x_type = 0
        if a.maxx < b.minx-1e-4: 
            x_type = 1 # left: a b
        elif a.minx > b.maxx+1e-4:
            x_type = 2  # right: b a
        elif a.minx < b.minx and a.maxx < b.maxx:
            x_type = 3  # near left
        elif a.minx > b.minx and a.maxx > b.maxx:
            x_type = 4  # near right
        elif a.minx < b.minx and a.maxx > b.maxx:
            x_type = 5  # contain
        elif a.minx > b.minx and a.maxx < b.maxx:
            x_type = 6  # inside
        else: x_type = 6
        return x_type

    @staticmethod
    def calc_y_type(a, b):
        y_type = 0
        if a.maxy < b.miny-1e-4:
            y_type = 1  # up
        elif a.miny > b.maxy+1e-4:
            y_type = 2  # down
        elif a.miny < b.miny and a.maxy < b.maxy:
            y_type = 3  # near up
        elif a.miny > b.miny and a.maxy > b.maxy:
            y_type = 4  # near down
        elif a.miny < b.miny and a.maxy > b.maxy:
            y_type = 5  # contain
        elif a.miny > b.miny and a.maxy < b.maxy:
            y_type = 6  # inside
        else: y_type = 6
        return y_type
    
    @staticmethod
    def bool2int(b):
        if b: return -1
        else: return 1
    
    def cmp_boxes(self, a, b):
        y_type = self.calc_y_type(a, b)
        x_type = self.calc_x_type(a, b)
        y_near_rate = 0.0
        x_near_rate = 0.0
        if y_type==3: y_near_rate = (a.maxy-b.miny) / min(a.maxy-a.miny, b.maxy-b.miny)
        elif y_type==4: y_near_rate = (b.maxy-a.miny) / min(a.maxy-a.miny, b.maxy-b.miny)
        if x_type==3: x_near_rate = (a.maxx-b.minx) / min(a.maxx-a.minx, b.maxx-b.minx)
        elif x_type==4: x_near_rate = (b.maxx-a.minx) / min(a.maxx-a.minx, b.maxx-b.minx)

        # fprintf(stderr, "%s vs %s, x_type %d, y_type %d\n", a.text.c_str(), b.text.c_str(), x_type, y_type)
        if y_type==1: res = True
        elif y_type==2: res = False
        elif y_type==3:
            if x_type==1: res = True
            elif x_type==2: res = y_near_rate<0.5
            elif x_type==3: res = True
            elif x_type==4: res = y_near_rate<0.5
            else: res = True
        elif y_type==4:
            if x_type==1: res = y_near_rate>0.5
            elif x_type==2: res = False
            elif x_type==3: res = y_near_rate>0.5
            elif x_type==4: res = False
            else: res = False
        else:
            if x_type==1 or x_type==3: res = True
            elif x_type==2 or x_type==4: res = False
            else: res = (a.minx+a.maxx)<(b.minx+b.maxx)
        return self.bool2int(res)

    def __call__(self, seg1, seg2):
        """
        segi: {..., coord_key: [2 pts or 4 pts]}
        """
        a, b = self._structuralize(seg1[self.coord_key], seg2[self.coord_key])
        return self.cmp_boxes(a, b)


# Usage for example
from functools import cmp_to_key
cmp_box_fn = CompareBoxes(coord_key="box")
line_info = sorted(line_info, key=cmp_to_key(cmp_box_fn))