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 "".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