utils/panel_ordering.py

import cv2

class BoundingBox(object):
    def __init__(self, xmin=None, ymin=None, xmax=None, ymax=None,
    panels=None,
    bbtype=None,
    id_=""):
        self.dict = {
            "@xmin": xmin,
            "@ymin": ymin,
            "@xmax": xmax,
            "@ymax": ymax,
            "@id": id_,
        }
        if panels is None:
            self.panels = [self]
        else:
            self.panels = panels
        self.bbtype = bbtype

    def init_dict(self, d):
        self.dict = d
        self.dict["@xmin"] = float(self.xmin)
        self.dict["@ymin"] = float(self.ymin)
        self.dict["@xmax"] = float(self.xmax)
        self.dict["@ymax"] = float(self.ymax)
        return self

    def __getitem__(self, index):
        return self.dict[index]

    @property
    def xmin(self):
        return self.dict["@xmin"]

    @property
    def xmax(self):
        return self.dict["@xmax"]

    @property
    def ymin(self):
        return self.dict["@ymin"]

    @property
    def ymax(self):
        return self.dict["@ymax"]

    @property
    def width(self):
        return self.xmax - self.xmin

    @property
    def height(self):
        return self.ymax - self.ymin

    @property
    def text(self):
        return self.dict["#text"]

    @property
    def id(self):
        return self.dict["@id"]

    @property
    def list(self):
        return [self.xmin, self.ymin, self.xmax, self.ymax]

    @property
    def is_null(self):
        return self.xmin is None or self.ymin is None or self.xmax is None or self.ymax is None

    @property
    def area(self):
        if self.xmax is None or self.xmin is None or self.ymax is None or self.ymin is None:
            return 0
        return (self.xmax - self.xmin) * (self.ymax - self.ymin)

    @property
    def base_panels(self):
        return len(self.panels)


    def __getitem__(self, item):
        return self.dict[item]

    def __add__(self, a):
        assert issubclass(type(a), BoundingBox)
        if a.is_null:
            return self
        elif self.is_null:
            return a
        return BoundingBox(xmin=min(self.xmin, a.xmin),
                           ymin=min(self.ymin, a.ymin),
                           xmax=max(self.xmax, a.xmax),
                           ymax=max(self.ymax, a.ymax),
                           panels=self.panels + a.panels)

    def __mul__(self, a):
        assert issubclass(type(a), BoundingBox)
        bb = BoundingBox(xmin=max(self.xmin, a.xmin),
                         ymin=max(self.ymin, a.ymin),
                         xmax=min(self.xmax, a.xmax),
                         ymax=min(self.ymax, a.ymax),
                         panels=self.panels + a.panels)
        if bb.xmin > bb.xmax or bb.ymin > bb.ymax:
            return BoundingBox()
        else:
            return bb

    def __repr__(self):
        return "<BoundingBox({},{}) {},{},{},{},{}>".format(self.bbtype, self.id, *self.list, self.base_panels)
    
    

def get_pivot_side(zmin, zmax, pivot):
    interception_ratio_threshold = 0.25
    if pivot <= zmin:
        return 1
    elif zmax <= pivot:
        return 0
    else:
        pivot_z_ratio = (pivot - zmin) / (zmax - zmin)
        interception_ratio = min(pivot_z_ratio, 1 - pivot_z_ratio)

        if interception_ratio > interception_ratio_threshold:
            return -1
        else:
            return 0 if pivot_z_ratio > 0.5 else 1


class BoxSet(set):
    def get_highest_priority_division(self):
        # Horizontal division
        ydivs = sorted([bb.ymin for bb in self] + [bb.ymax for bb in self])
        for pivot in ydivs:
            division = self.get_pivot_division(pivot,
                                               is_horizontal_division=True)
            if len(division) > 1:
                return division

        # Vertical division
        xdivs = sorted([bb.xmin for bb in self] + [bb.xmax for bb in self], reverse=True)
        for pivot in xdivs:
            division = self.get_pivot_division(pivot,
                                               is_horizontal_division=False)
            if len(division) > 1:
                return division

        # Undividable box set
        return [self]

    def get_pivot_division(self, pivot, is_horizontal_division):
        divs = [BoxSet(), BoxSet()]
        for bb in self:
            if is_horizontal_division:
                side = get_pivot_side(bb.ymin, bb.ymax, pivot)
            else:
                side = get_pivot_side(-bb.xmax, -bb.xmin, -pivot)

            if side == -1:
                return [self]
            else:
                divs[side].add(bb)
        if len(divs[0]) == 0 or len(divs[1]) == 0:
            return [self]
        return divs

    def get_multicut_division(self, cuts):
        curset = self
        cur_division = []
        for cut in cuts:
            pivot, is_horizontal_division = cut
            division = curset.get_pivot_division(pivot, is_horizontal_division)
            if len(division) > 1:
                cur_division.append(division[0])
                curset = division[1]
        if len(cur_division) > 0:
            return cur_division + [curset]
        else:
            return [self]

    def yield_ordered_bbs(self):
        if len(self) == 0:
            pass
        elif len(self) > 1:
            yield self.sum(), False
        else:
            yield next(iter(self)), True

    def sum(self):
        if len(self) == 0:
            return BoundingBox()
        else:
            l = list(self)
            return sum(l[1:], l[0])

class BoxNode(object):
    def __init__(self, bbset, initial_cuts=None):
        if initial_cuts:
            division = bbset.get_multicut_division(initial_cuts)
        else:
            division = [bbset]

        if len(division) == 1:
            division = bbset.get_highest_priority_division()

        isLeaf = len(division) <= 1
        self.division = division if isLeaf else [BoxNode(section) for section in division]

    def yield_ordered_bbs(self):
        for section in self.division:
            for bb in section.yield_ordered_bbs():
                yield bb

class BoxOrderEstimator(object):
    def __init__(self, bbs, pagewidth=None, initial_cut_option=None):
        if initial_cut_option == "two-page-four-panel":
            initial_cuts = [(pagewidth * n / 4, False)
                            for n in reversed(range(1, 4))]
        elif initial_cut_option == "two-page":
            initial_cuts = [(pagewidth / 2, False)]
        else:
            initial_cuts = None

        self.boxnode = BoxNode(BoxSet(bbs), initial_cuts)
        t = tuple(zip(*self.boxnode.yield_ordered_bbs()))
        if len(t) > 0:
            self.ordered_bbs, self.bb_estimation_statuses = t
        else:
            self.ordered_bbs, self.bb_estimation_statuses = (), ()


def panel_ordering(test_image,dets):
    image = cv2.imread(test_image)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB for display
    # interception_ratio_threshold = 0.25

    # print(predictions_frame)

    panels = set()
    # for panel in dets['frame']:
    for panel in dets.panels:
        panel = panel['bbox']
        panels.add(BoundingBox(panel[0],panel[1],panel[2],panel[3]))
    # print(panels)

    # image = page.get_image()
    pagewidth = image.size
    pageheight, pagewidth, pagechannels = image.shape

    # panels = page.get_bbs()["frame"]
    # print(panels)
    boxOrderEstimator = BoxOrderEstimator(
        panels,
        pagewidth=pagewidth,
        initial_cut_option="two-page")
    
    return boxOrderEstimator