yomitoku 0.4.0.post1.dev0__py3-none-any.whl

yomitoku/postprocessor/rtdetr_postprocessor.py

@@ -0,0 +1,107 @@
+"""Copyright(c) 2023 lyuwenyu. All Rights Reserved."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+
+
+def mod(a, b):
+    out = a - a // b * b
+    return out
+
+
+class RTDETRPostProcessor(nn.Module):
+    __share__ = [
+        "num_classes",
+        "use_focal_loss",
+        "num_top_queries",
+        "remap_mscoco_category",
+    ]
+
+    def __init__(
+        self,
+        num_classes=80,
+        use_focal_loss=True,
+        num_top_queries=300,
+        remap_mscoco_category=False,
+    ) -> None:
+        super().__init__()
+        self.use_focal_loss = use_focal_loss
+        self.num_top_queries = num_top_queries
+        self.num_classes = int(num_classes)
+        self.remap_mscoco_category = remap_mscoco_category
+        self.deploy_mode = False
+
+    def extra_repr(self) -> str:
+        return f"use_focal_loss={self.use_focal_loss}, num_classes={self.num_classes}, num_top_queries={self.num_top_queries}"
+
+    # def forward(self, outputs, orig_target_sizes):
+    def forward(self, outputs, orig_target_sizes: torch.Tensor, threshold):
+        logits, boxes = outputs["pred_logits"], outputs["pred_boxes"]
+        # orig_target_sizes = torch.stack([t["orig_size"] for t in targets], dim=0)
+
+        bbox_pred = torchvision.ops.box_convert(boxes, in_fmt="cxcywh", out_fmt="xyxy")
+        bbox_pred *= orig_target_sizes.repeat(1, 2).unsqueeze(1)
+
+        if self.use_focal_loss:
+            scores = F.sigmoid(logits)
+            scores, index = torch.topk(scores.flatten(1), self.num_top_queries, dim=-1)
+            # TODO for older tensorrt
+            # labels = index % self.num_classes
+            labels = mod(index, self.num_classes)
+            index = index // self.num_classes
+            boxes = bbox_pred.gather(
+                dim=1,
+                index=index.unsqueeze(-1).repeat(1, 1, bbox_pred.shape[-1]),
+            )
+
+        else:
+            scores = F.softmax(logits)[:, :, :-1]
+            scores, labels = scores.max(dim=-1)
+            if scores.shape[1] > self.num_top_queries:
+                scores, index = torch.topk(scores, self.num_top_queries, dim=-1)
+                labels = torch.gather(labels, dim=1, index=index)
+                boxes = torch.gather(
+                    boxes,
+                    dim=1,
+                    index=index.unsqueeze(-1).tile(1, 1, boxes.shape[-1]),
+                )
+
+        # TODO for onnx export
+        if self.deploy_mode:
+            return labels, boxes, scores
+
+        # TODO
+        if self.remap_mscoco_category:
+            from ...data.dataset import mscoco_label2category
+
+            labels = (
+                torch.tensor(
+                    [mscoco_label2category[int(x.item())] for x in labels.flatten()]
+                )
+                .to(boxes.device)
+                .reshape(labels.shape)
+            )
+
+        results = []
+        for lab, box, sco in zip(labels, boxes, scores):
+            lab = lab[sco > threshold]
+            box = box[sco > threshold]
+            sco = sco[sco > threshold]
+
+            lab = lab.cpu().numpy()
+            box = box.cpu().numpy()
+            sco = sco.cpu().numpy()
+
+            result = dict(labels=lab, boxes=box, scores=sco)
+            results.append(result)
+
+        return results
+
+    def deploy(
+        self,
+    ):
+        self.eval()
+        self.deploy_mode = True
+        return self

yomitoku/reading_order.py

@@ -0,0 +1,214 @@
+import cv2
+
+from .utils.graph import Node
+from .utils.misc import (
+    is_intersected_vertical,
+    is_intersected_horizontal,
+)
+
+
+def is_locked_node(node):
+    return all([child.is_locked for child in node.children])
+
+
+def _priority_dfs(nodes, direction):
+    if len(nodes) == 0:
+        return []
+
+    pending_nodes = sorted(nodes, key=lambda x: x.prop["distance"])
+    visited = [False] * len(nodes)
+
+    start = pending_nodes.pop(0)
+    stack = [start]
+
+    order = []
+    open_list = []
+
+    while not all(visited):
+        while stack:
+            is_updated = False
+            current = stack.pop()
+            if not visited[current.id]:
+                parents = current.parents
+                if all([visited[parent.id] for parent in parents]) or len(parents) == 0:
+                    visited[current.id] = True
+                    order.append(current.id)
+                    is_updated = True
+                else:
+                    if current not in open_list:
+                        open_list.append(current)
+
+            if is_updated:
+                for open_node in reversed(open_list):
+                    stack.append(open_node)
+                    open_list.remove(open_node)
+
+            if len(current.children) > 0:
+                stack.append(current)
+
+            if len(current.children) == 0:
+                children = []
+                for node in stack:
+                    if current in node.parents:
+                        children.append(node)
+                        stack.remove(node)
+
+                if direction == "horizontal":
+                    children = sorted(
+                        children, key=lambda x: x.prop["box"][0], reverse=True
+                    )
+                else:
+                    children = sorted(
+                        children, key=lambda x: x.prop["box"][1], reverse=True
+                    )
+
+                stack.extend(children)
+                continue
+
+            child = current.children.pop(0)
+            stack.append(child)
+
+        for node in pending_nodes:
+            if node in open_list:
+                continue
+            stack.append(node)
+            pending_nodes.remove(node)
+            break
+        else:
+            if not all(visited) and len(open_list) != 0:
+                node = open_list.pop(0)
+                visited[node.id] = True
+                order.append(node.id)
+
+    return order
+
+
+def _exist_other_node_between_vertical(node, other_node, nodes):
+    for search_node in nodes:
+        if search_node == node or search_node == other_node:
+            continue
+
+        _, sy1, _, sy2 = search_node.prop["box"]
+        _, oy1, _, oy2 = other_node.prop["box"]
+        _, ny1, _, ny2 = node.prop["box"]
+
+        if is_intersected_vertical(search_node.prop["box"], node.prop["box"]):
+            if ny2 < sy1 < oy1 and ny2 < sy2 < oy1:
+                return True
+
+            if oy2 < sy1 < ny1 and oy2 < sy2 < ny1:
+                return True
+
+    return False
+
+
+def _exist_other_node_between_horizontal(node, other_node, nodes):
+    for search_node in nodes:
+        if search_node == node or search_node == other_node:
+            continue
+
+        sx1, _, sx2, _ = search_node.prop["box"]
+        ox1, _, ox2, _ = other_node.prop["box"]
+        nx1, _, nx2, _ = node.prop["box"]
+
+        if is_intersected_horizontal(search_node.prop["box"], node.prop["box"]):
+            if nx2 < sx1 < ox1 and nx2 < sx2 < ox1:
+                return True
+
+            if ox2 < sx1 < nx1 and ox2 < sx2 < nx1:
+                return True
+
+    return False
+
+
+def _create_graph_horizontal(nodes):
+    for i, node in enumerate(nodes):
+        for j, other_node in enumerate(nodes):
+            if i == j:
+                continue
+
+            if is_intersected_vertical(node.prop["box"], other_node.prop["box"]):
+                ty = node.prop["box"][1]
+                oy = other_node.prop["box"][1]
+
+                if _exist_other_node_between_vertical(node, other_node, nodes):
+                    continue
+
+                if ty < oy:
+                    node.add_link(other_node)
+                else:
+                    other_node.add_link(node)
+
+            node_distance = node.prop["box"][0] + node.prop["box"][1]
+            node.prop["distance"] = node_distance
+
+    for node in nodes:
+        node.children = sorted(node.children, key=lambda x: x.prop["box"][0])
+
+
+def _create_graph_vertical(nodes):
+    max_x = max([node.prop["box"][2] for node in nodes])
+
+    for i, node in enumerate(nodes):
+        for j, other_node in enumerate(nodes):
+            if i == j:
+                continue
+
+            if is_intersected_horizontal(node.prop["box"], other_node.prop["box"]):
+                tx = node.prop["box"][2]
+                ox = other_node.prop["box"][2]
+
+                if _exist_other_node_between_horizontal(node, other_node, nodes):
+                    continue
+
+                if tx < ox:
+                    other_node.add_link(node)
+                else:
+                    node.add_link(other_node)
+
+            node.prop["distance"] = (max_x - node.prop["box"][2]) + node.prop["box"][1]
+
+    for node in nodes:
+        node.children = sorted(node.children, key=lambda x: x.prop["box"][1])
+
+
+def prediction_reading_order(elements, direction, img=None):
+    if len(elements) < 2:
+        return elements
+
+    nodes = [Node(i, element.dict()) for i, element in enumerate(elements)]
+    if direction == "horizontal":
+        _create_graph_horizontal(nodes)
+    else:
+        _create_graph_vertical(nodes)
+
+    # For debugging
+    # if img is not None:
+    #    visualize_graph(img, nodes)
+
+    order = _priority_dfs(nodes, direction)
+    for i, index in enumerate(order):
+        elements[index].order = i
+
+    return elements
+
+
+def visualize_graph(img, nodes):
+    out = img.copy()
+    for node in nodes:
+        for child in node.children:
+            nx1, ny1, nx2, ny2 = node.prop["box"]
+            cx1, cy1, cx2, cy2 = child.prop["box"]
+
+            node_center = nx1 + (nx2 - nx1) // 2, ny1 + (ny2 - ny1) // 2
+            child_center = cx1 + (cx2 - cx1) // 2, cy1 + (cy2 - cy1) // 2
+
+            cv2.arrowedLine(
+                out,
+                (int(node_center[0]), int(node_center[1])),
+                (int(child_center[0]), int(child_center[1])),
+                (0, 0, 255),
+                2,
+            )
+
+    cv2.imwrite("graph.jpg", out)

yomitoku/resource/charset.txt

@@ -0,0 +1 @@
+ !"#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\]^_`abcdefghijklmnopqrstuvwxyz{|}~¡¢£¤¥¦§©«¬­®°±¶·»¿Å×ð÷ÿħʒ̧́̅̈̋ΓΔΘΛΞΠΣΦΨΩαβγδεζηθικλμνξοπρστυφχψωДЩджзщъёאิლ‐–—‖‘’“”†‡•‰′※‿⁂⁄⁑€℧←↑→↓↔↖↗↘↙⇄⇒⇔⇦⇧⇨⇩∀∂∃∅∇∈∉∋−∓√∝∞∟∠∦∧∨∩∪∫∮∴∵∽≃≅≈≒≠≡≢≦≧≪≫≶≷⊂⊃⊄⊅⊆⊇⊊⊋⊕⊖⊗⊥⊿⋚⋛⌅⌆⌒⌘⎾⎿⏀⏁⏂⏃⏄⏅⏆⏇⏈⏉⏊⏋⏌⏎⓫⓬⓭⓮⓯⓰⓱⓲⓳⓴⓵⓶⓷⓸⓹⓺⓻⓼⓽⓾─━┌┐┘├╹■□▱▲△▶▷▼▽◀◁◆◇◉○◎●◐◑◒◓◕◡◦◯☀☁☂☃★☆☎☖☗☞♀♂♠♡♢♣♤♥♦♧♨♩♪♫♬♭♮♯✓❖❶❷❸❹❺❻❼❽❾❿⤴⤵⦅⦆⦿⧺⧻、。〃〄々〆〇〈〉《》「」『』【】〒〓〔〕〖〗〘〙〜〝〟〠〳〴〵〻〼〽ぁあぃいぅうぇえぉおかがきぎくぐけげこごさざしじすずせぜそぞただちぢっつづてでとどなにぬねのはばぱひびぴふぶぷへべぺほぼぽまみむめもゃやゅゆょよらりるれろゎわゐゑをん゙゚ゝゞ゠ァアィイゥウェエォオカガキギクグケゲコゴサザシジスズセゼソゾタダチヂッツヅテデトドナニヌネノハバパヒビピフブプヘベペホボポマミムメモャヤュユョヨラリルレロヮワヰヱヲンヴヵヶヷヸヹ・ーヽヾㇲㇵㇶㇷㇻ㒵㓛㕘㗂㩁㮶㷀㷔䑓䒾䠖䰗一丁七万丈三上下不与丐丑且丕世丘丙丞両並个中丰串丸丹主丼丿乂乃久之乍乎乏乕乖乗乘乙九乞也乢乱乳乾亀亂了予争亊事二于云互五井亘亙些亜亞亟亡亢交亥亦亨享京亭亮亳亶亷亹人亻什仁仄仆仇今介仍从仏仔仕他仗付仙仝仞代令以仭仮仰仲仵件价任份仿企伉伊伋伍伎伏伐休会伜伝伯估伴伶伸伺似伽佃但佇位低住佐佑体何佗余佚佛作佝佞佟佤佩佯佰佳併佶佻佼佾使侃來侈例侍侏侑侖侗侘供依侠価侫侭侮侯侵侶便係促俄俅俊俎俐俑俔俗俘俚俛保俟俠信俣俤俥修俯俱俳俵俶俸俺俾倂倅倆倉個倍倏們倒倔倖候倚倞借倣値倥倦倨倩倪倫倭倶倹倻偃假偈偉偏偐偕偖做停健偬偲側偵偶偸偽傀傅傍傑傕傘備傚傪催傭傲傳傴債傷傾僂僄僅僉僊働像僑僔僕僖僙僚僞僣僥僧僭僮價僻儀儁儂億儉儋儒儔儕儘儚儛儞償儡優儲儷儺儻儼兀允元兄充兆兇先光兊克兌免兎児兒兔兕兗党兜兢入內全兩兪八公六兮共兵其具典兼冀内円冉冊册再冐冑冒冕冗写冠冢冤冥冦冨冩冪冬冰冱冲决冴况冶冷冼冽凄准凉凋凌凍凖凛凜凝凞几凡処凧凩凪凭凮凰凱凶凸凹出函凾刀刃刄分切刈刊刋刎刑刕列初判別利刪刮到刳制刷券刹刺刻剃剄則剉削剋剌前剔剖剛剜剝剣剤剥剩剪副剰剱割剴創剽剿劃劇劈劉劍劒劔力功加劣助努劫劬劭励労劵効劻劾勁勃勅勇勉勒動勘務勛勝勞募勢勣勤勦勧勰勲勳勵勸勺勾勿匀匁匂包匆匈匍匏匐匕化北匙匝匠匡匣匪匯匱匹区医匾匿區十千卅卆升午卉半卍卑卒卓協南単博卜卞占卣卦卮卯印危即却卵卷卸卻卽卿厄厓厖厘厚原厠厥厦厨厩厭厰厲厳厷去参參又叉及友双反収叔取受叙叛叟叡叢口古句叩只叫召叭叮可台叱史右叶号司叺吁吃各合吉吊同名后吏吐向君吝吞吟吠否吧吩含听吶吸吹吻吼吽吾呀呂呆呈呉告呍呎呑呕呟周呪呰呱味呵呶呷呻呼命咀咄咆咈咋和咎咏咐咒咖咜咢咤咥咨咩咫咬咯咲咳咸咽哀品哂哄哆哇哈哉員哥哨哩哭哮哲哺哿唄唆唇唎唐唖售唯唱唳唸唹唼唾啀啄商問啐啓啖啗啜啞啡啻啼啾喀喃善喆喇喉喊喋喎喘喙喚喜喝喟喧喨喩喪喫喬單喰喱営嗄嗅嗇嗒嗔嗚嗜嗟嗣嗤嗧嗽嗾嘆嘈嘉嘎嘔嘖嘗嘘嘩嘯嘰嘱嘲嘴嘶嘸嘻噂噉噌噎噓噛噤噦器噪噫噬噯噲噴噶噸噺嚀嚆嚇嚊嚏嚔嚙嚝嚠嚢嚥嚬嚭嚮嚴嚶嚼囀囁囂囃囅囈囊囎囑囓囗囘囚四回因団囮困囲図囹固国囿圀圃圄國圍圏園圓圖團圜土圦圧在圭地圳圷圻址坂均坊坌坍坎坏坐坑坡坤坦坨坩坪垂垈型垓垜垠垢垣垤垰垸埃埋城埒埓埔埜域埠埣埤埴埵執培基埼堀堂堅堆堊堕堙堝堡堤堪堯堰報場堵堺塀塁塊塋塌塑塒塔塗塘塙塚塞塡塢塩填塰塲塵塹塼塾境墅墉墓増墜增墟墨墮墳墺墻墾壁壅壇壊壌壑壒壓壕壘壙壜壞壟壤士壬壮壯声壱売壷壹壺壻壼壽変夋夏夔夕外夘夙多夛夜夢夤夥大天太夫夬夭央失夲夷夸夾奄奇奈奉奎奏奐契奔奕套奘奚奝奠奢奥奧奨奪奬奭奮女奴奶奸好妁如妃妄妊妍妓妖妙妝妣妤妥妨妬妮妲妹妺妻妾姆姉始姐姑姒姓委姙姚姜姣姥姦姧姨姪姫姮姶姸姻姿威娃娉娍娑娘娜娟娠娥娩娯娵娶娼婀婁婆婉婕婚婢婣婦婪婬婿媄媒媚媛媧媱媳媼媽媾嫁嫂嫄嫉嫋嫌嫖嫗嫡嫣嫦嫩嫺嬉嬋嬌嬖嬢嬥嬪嬬嬰嬲嬴嬶孀孁孃孅孌子孑孔孕字存孚孛孜孝孟季孤学孩孫孰孵學孺孽宀宅宇守安宋完宍宏宓宕宗官宙定宛宜宝実客宣室宥宦宮宰害宴宵家宸容宿寂寃寄寅密寇富寐寒寓寔寛寝寞察寡寢寤寥實寧寨審寫寮寰寳寵寶寸寺対寿封専射尅将將專尉尊尋對導小少尓尖尚尠尢尤尨尩尫尭就尸尹尺尻尼尽尾尿局屁居屆屈届屋屍屎屏屐屑屓展屛属屠屡屢層履屬屮屯山屹岌岐岑岡岨岩岫岬岱岳岷岸岺岻峅峙峠峡峨峩峪峭峯峰峴島峻峽崇崋崎崐崑崔崕崖崗崘崙崚崛崟崢崧崩嵆嵋嵌嵎嵐嵜嵩嵬嵯嶂嶄嶇嶋嶌嶠嶬嶮嶲嶷嶸嶺嶼嶽巉巌巍巎巒巓巖川州巡巢巣工左巧巨巫差己已巳巴巵巷巻巽巾市布帆帋希帑帒帕帖帙帚帛帝帥師席帮帯帰帳帶帷常帽幀幄幅幇幌幎幔幕幖幗幘幟幡幢幣幤幫幮干平年并幷幸幹幻幼幽幾庁広庄庇床序底庖店庚府庠度座庫庭庵庶康庸庾廁廂廃廈廉廊廋廏廐廓廖廚廛廟廠廡廢廣廨廩廬廰廳廴延廷廸建廻廼廿弁弄弉弊弋弌弍式弐弑弓弔引弖弗弘弛弟弥弦弧弩弭弯弱弴張強强弻弼弾彅彈彊彌彎当彗彙彝彡形彣彤彦彧彩彪彫彬彭彰影彳彷役彼彽彿往征徂徃径待徇徊律後徐徑徒従得徘徙從徠御徧徨復循徭微徳徴徵德徹徼徽心必忉忌忍忖志忘忙応忝忠忡忤忩快忯忰忱念忸忻忽忿怏怐怒怕怖怙怛怜思怠怡急性怨怩怪怫怯怱怵怺恁恂恃恆恊恋恍恐恒恕恙恚恟恠恢恣恤恥恨恩恪恫恬恭息恰恵悃悄悅悉悋悌悍悔悖悚悛悝悟悠患悦悧悩悪悰悲悳悴悵悶悸悼悽情惇惋惑惓惔惕惘惚惜惟惠惡惣惧惨惰惱想惴惶惹惺惻愁愆愈愉愍愎意愔愕愚愛感愡愧愨愬愴愷愼愽愾愿慂慄慇慈慊態慌慍慎慓慕慘慙慚慟慢慣慥慧慨慫慮慰慳慴慶慷慾憂憇憊憍憎憐憑憔憖憙憚憤憧憩憫憬憮憲憶憺憾懃懆懇懈應懊懋懌懐懣懦懲懴懶懷懸懺懼懽懾懿戀戈戊戌戍戎成我戒戔戕或戚戛戞戟戡戦戩截戮戯戰戲戴戸戻房所扁扇扈扉手才扎打払托扛扞扠扣扨扭扮扱扶批扼承技抂抄抉把抑抒抓抔投抖抗折抛抜択披抬抱抵抹抻押抽拂担拆拇拈拉拋拌拍拏拐拒拓拔拕拗拘拙招拜拝拠拡括拭拮拯拱拳拵拶拷拼拾拿持挂指挈按挊挌挍挐挑挘挙挟挨挫振挵挹挺挽挾挿捉捌捍捏捐捕捗捜捥捧捨捩据捲捶捷捺捻掃授掉掌掏排掖掘掛掟掠採探掣接控推掩措掬掲掴掻掾揀揃揄揆揉描提插揖揚換握揣揩揭揮援揶揷揺搆損搏搔搖搗搜搤搦搨搩搬搭搶携搾摂摑摘摠摧摩摭摯摶摸摹摺撃撈撑撒撓撕撚撝撞撤撥撩撫播撮撰撲撹撻撼撾撿擁擂擅擇擉操擒擔擘據擠擡擢擣擤擦擧擬擯擱擲擴擻擽擾攀攅攘攝攢攣攪攫攬支攴攵收攷攸改攻放政故效敍敎敏救敕敖敗教敝敢散敦敧敬数敲整敵敷數斂斃文斉斌斎斐斑斗料斛斜斟斡斤斥斧斫斬断斯新斷方於施旁旄旅旆旋旌族旒旗旛无既日旦旧旨早旬旭旱旺旻旼昂昃昆昇昉昊昌明昏易昔昕星映春昧昨昭是昱昴昵昶昺昻昼晁時晃晄晉晋晏晒晝晟晡晢晤晦晧晨晩普景晰晴晶晷智暁暇暈暉暎暑暖暗暘暝暠暢暦暫暮暲暴暸暹暻暼暾曁曄曇曈曉曖曙曜曝曠曦曩曰曲曳更曵曷書曹曺曻曼曽曾替最會朅月有朋服朏朔朕朖朗望朝期朦朧木未末本札朮朱朴朶朸机朽杁杆杉李杏材村杓杖杙杜杞束杠条杢杣杤来杦杭杮杯杰東杲杳杵杷杼松板枇枉枋枌析枒枕林枘枚果枝枠枡枢枯枲枳枴架枷枸枹枻柁柄柊柏某柑染柔柘柚柝柞柢柩柬柯柱柳柴柵査柾柿栂栃栄栓栖栗栞校栢栩株栫栱栲栴核根格栽栾桀桁桂桃桅框案桎桐桑桒桓桔桙桛桜桝桟档桧桴桶桾桿梁梃梅梏梓梔梗梛條梟梠梢梦梧梨梭梯械梱梲梳梵梶梼棃棄棈棉棊棋棍棐棒棕棖棗棘棚棟棠棡棣棧棭森棰棱棲棹棺椀椅椋植椎椏椒椙椚椛検椥椨椰椴椵椹椽椿楂楊楓楔楕楚楝楞楠楡楢楣楤楨楪楫業楮楯楳極楷楸楹楼楽榀榁概榊榎榑榔榕榛榜榧榮榲榴榻榾槁槃槇槊構槌槍槎槐槓槖様槙槧槨槫槲槵槹槻槽槿樂樅樊樋樒樓樗標樝樞樟模樣権横樫樵樸樹樺樽橄橅橇橈橋橐橘橙橛機橡橢橫橳橿檀檄檍檎檐檔檗檜檞檠檢檣檪檫檬檮檳檸檻櫂櫃櫑櫓櫚櫛櫞櫟櫨櫪櫬櫺櫻欄欅權欒欖欝欠次欣欧欲欵欷欹欺欽款歃歇歉歌歎歐歓歔歙歛歟歡止正此步武歩歪歯歳歴歷歸死歿殂殃殄殆殉殊残殖殘殞殤殪殯殱殲殳殴段殷殺殻殼殿毀毅毆毉毋毌母毎每毒毓比毖毗毘毛毟毫毬毮毯毱毳氈氏氐民氓气気氣氤氳水氵氷永氾汀汁求汎汐汕汗汙汚汜汝汞江池汨汪汰汲汳汴汶決汽汾沁沂沃沅沈沌沍沐沒沓沖沙沚沛没沢沪沫沮沱河沸油治沼沽沾沿況泄泉泊泌泓泔法泗泛泝泡波泣泥注泪泮泯泰泳洋洌洎洒洗洙洛洞洟津洩洪洫洮洱洲洵洶洸洹活洽派流浄浅浙浚浜浣浦浩浪浬浮浴海浸涂涅涇消涌涎涓涕涙涛涜涪涯液涵涸涼涿淀淄淅淆淇淋淑淒淖淘淙淚淛淝淞淡淤淦淨淪淫淮深淳淵混淹淺添清渇済渉渋渓渕渗渙渚減渝渟渠渡渣渤渥渦渧温渫測渭港游渺渾湃湄湊湍湎湑湖湘湛湣湧湫湮湯湲湾湿満溂溌源準溘溜溝溟溢溥溪溫溯溲溶溷溺滁滂滄滅滇滉滋滌滎滑滓滔滕滙滝滞滬滯滲滴滸滾滿漁漂漆漉漏漑漓演漕漠漢漣漪漫漬漱漲漳漸漾漿潁潅潑潔潘潙潛潜潞潟潢潤潦潭潮潯潰潴潸潺潼潾澀澁澂澄澆澈澎澗澠澡澣澤澪澱澳澶澹激濁濂濃濊濕濘濛濟濠濡濤濫濮濯濱濳濵濶濹濾瀆瀉瀋瀏瀑瀕瀘瀚瀛瀝瀞瀟瀦瀧瀨瀬瀰瀲瀾灌灑灘灝灣灤火灯灰灸灼災炉炊炎炒炕炙炫炬炭炮炯炳炷炸点為炻烈烏烔烘烙烝烟烤烬烱烹烽焉焔焙焚焜無焦焮焰焱然焼煆煇煉煊煌煎煕煖煙煞煠煢煤煥煦照煨煩煬煮煽熄熇熈熊熏熒熔熕熖熙熟熨熬熱熹熾燁燃燄燈燎燐燒燔燕燗營燠燥燦燧燬燭燮燵燹燻燼燾燿爆爍爐爛爨爪爬爭爰爲爵父爸爹爺爻爼爽爾牀牆片版牋牌牒牓牖牘牙牛牝牟牡牢牧物牯牲牴特牽牾犀犁犂犇犍犒犠犢犧犬犯状狀狂狃狄狆狎狐狒狗狙狛狠狡狢狩独狭狷狸狹狼狽猊猖猗猙猛猜猝猟猥猨猩猪猫献猯猴猶猷猹猽猾猿獄獅獏獐獒獗獣獦獨獪獰獱獲獵獸獺獻玄玅率玉王玕玖玘玟玩玫玲玳玵玻珀珂珈珉珊珍珎珞珠珢珥珩珪班珮現球琅理琉琊琛琢琥琦琨琪琫琬琮琰琲琳琴琵琶琺瑀瑁瑄瑇瑋瑕瑗瑙瑚瑛瑜瑞瑟瑠瑣瑤瑩瑪瑫瑯瑰瑳瑶瑾璀璃璇璈璉璋璐璘璜璞璟璣璧璨璫環璽璿瓊瓏瓔瓘瓚瓜瓠瓢瓣瓦瓩瓮瓱瓲瓶瓷甃甄甌甍甎甑甕甗甘甚甜甞生産甥甦用甪甫甬田由甲申男甸町画畊界畏畑畔留畚畜畝畠畢畤略畦畧番畫畭畯異畳畴當畷畸畿疆疇疊疋疎疏疑疔疚疝疣疥疫疱疲疳疵疸疹疼疽疾痂痃病症痍痎痒痔痕痘痙痛痞痢痣痤痩痰痲痳痴痺痼痾痿瘀瘁瘉瘋瘍瘖瘙瘞瘠瘡瘢瘤瘦瘧瘭瘰瘳瘴瘻療癆癇癈癋癌癒癖癘癜癡癢癧癩癪癬癭癲癸発登發白百皀皃的皆皇皈皋皎皐皓皕皖皙皚皮皰皴皷皸皹皺皿盂盃盆盈盉益盍盎盒盔盖盗盙盛盜盞盟盡監盤盥盦盧盪盬目盱盲直相盻盼盾省眄眇眈眉看県眙眛眞真眠眥眦眩眶眷眸眺眼着睇睛睜睡睢督睥睦睨睫睹睺睾睿瞀瞋瞑瞞瞠瞥瞬瞭瞰瞳瞶瞻瞼瞽瞿矍矗矚矛矜矢矣知矧矩短矮矯石砂砉砌砒研砕砖砥砦砧砭砲破砺硅硎硏硝硨硫硬确硯硴硼碁碆碇碌碍碎碑碓碕碗碣碧碩碪確碼碾磁磅磊磋磌磐磔磚磧磨磬磯磲磴磽礀礁礎礑礒礙礦礪礫礬示礼社祀祁祅祆祇祈祉祐祓祕祖祗祚祛祜祝神祟祠祢祥票祭祷祺祿禀禁禄禅禊禍禎福禖禛禝禦禧禪禮禰禱禳禴禹禺禽禾禿秀私秉秋种科秒秕秘秞租秡秣秤秦秧秩称移稀稈程稍税稔稗稙稚稜稟稠種稱稲稷稹稻稼稽稿穀穂穆穉穌積穎穏穐穗穙穡穢穣穩穫穰穴穵究穹空穽穿突窃窄窈窒窓窕窖窗窘窟窠窣窩窪窬窮窯窰窶窺窻窿竃竄竅竇竈竊立竎站竜竝竟章竢竣童竦竪竫竭端競竸竹竺竿笄笈笊笋笏笑笒笘笙笛笞笠笥符笨第笮笳笵笹笻筅筆筇筈等筋筌筍筏筐筑筒筓答策筝筠筥筧筬筭筮筰筱筵筺筿箆箇箋箍箏箒箔箕算箙箚箝箞管箪箭箯箱箴箸節篁範篆篇築篊篋篏篙篝篠篤篥篦篩篭篳篶篷簀簁簃簇簋簍簏簑簒簓簔簗簞簟簠簡簣簧簪簫簳簶簷簸簽簾簿籀籃籌籍籏籐籔籖籙籟籠籤籩籬米籵籾粁粂粃粉粋粍粏粒粕粗粘粛粟粡粢粤粥粧粮粱粲粳粹粼粽精糀糂糅糊糎糒糕糖糜糝糞糟糠糢糧糯糲糴糶糸糺系糾紀紂約紅紆紇紈紊紋納紐紒純紕紗紘紙級紛紜素紡索紫紬紮累細紱紲紳紹紺絁終絃組絅絆絋経絎絏結絓絖絛絜絞絡絢絣給絨絪絮統絲絳絵絶絹絽綉綏經継続綛綜綝綟綠綢綣綦綫綬維綮綯綰綱網綴綵綸綺綻綽綾綿緇緊緋総緑緒緖緘線緜緝緞締緡緣編緩緬緯緲練緻縁縄縅縉縊縋縒縕縛縝縞縟縢縣縦縫縬縮縱縲縵縷縹縺縻總績繁繃繅繆繇繊繋繍織繕繙繚繝繞繡繦繧繩繪繫繭繰繹繻繼繽繿纂纈纉纊續纍纎纏纐纒纓纔纜缶缸缺罄罅罌罎罐罔罕罘罠罡罣罨罩罪罫置罰署罵罷罸罹羂羅羆羇羈羊羌美羚羞羣群羨義羮羯羲羸羹羽羿翁翅翆翊翌翎習翔翕翛翟翠翡翦翩翫翰翳翹翻翼耀老考耄者耆耋而耐耒耕耗耘耙耜耡耦耨耳耶耻耽耿聃聆聊聖聘聚聞聟聡聢聨聯聰聱聲聳聴聶職聽聾聿肄肅肆肇肉肋肌肓肖肘肚肛肝股肢肥肩肪肭肯肱育肴肺胃胄胆背胎胖胚胛胝胞胡胤胥胯胱胳胴胸胼能脂脅脆脇脈脉脊脘脚脛脞脣脩脯脱脳脹脾腅腆腊腋腎腐腑腓腔腕腟腥腦腫腮腰腱腴腸腹腺腿膀膂膃膅膈膊膏膓膕膚膜膝膠膣膨膩膰膳膵膸膺膻膽膾膿臀臂臆臈臉臍臏臑臓臘臙臚臟臠臣臥臧臨自臬臭至致臺臼臾舁舂舅與興舉舊舌舍舎舐舒舖舗舘舛舜舞舟舩航舫般舲舳舵舶舷舸船艀艇艘艙艚艜艝艟艣艤艦艨艪艫艮良艱色艶艷艸艾芋芍芎芒芙芝芟芥芦芨芩芬芭芮芯花芳芷芸芹芻芽芾芿苅苑苒苓苔苕苗苛苜苞苟苡苣若苦苧苫英苳苴苹苺苻茂范茄茅茆茉茎茗茘茜茣茨茫茯茱茲茴茵茶茸茹荀荃荆草荊荏荒荔荘荳荵荷荻荼莅莆莉莊莎莒莓莕莖莘莚莞莟莠莢莧莨莩莪莫莱莵莽菀菁菅菇菊菌菎菓菖菘菜菟菠菩菫華菰菱菲菴菹菽萃萄萇萊萌萍萎萓萠萢萩萪萬萱萵萸萼落葆葈葉葎著葛葡葢董葦葩葫葬葭葯葱葵葷葺蒂蒅蒋蒐蒔蒙蒜蒟蒡蒯蒲蒸蒻蒼蒾蒿蓀蓁蓂蓄蓆蓉蓊蓋蓍蓐蓑蓓蓖蓙蓚蓬蓮蓯蓴蓼蓿蔀蔑蔓蔔蔕蔗蔘蔚蔞蔟蔡蔣蔦蔬蔭蔲蔵蔽蕁蕃蕈蕉蕊蕋蕎蕓蕗蕙蕚蕝蕞蕣蕨蕩蕪蕭蕷蕾薀薄薇薈薊薏薐薑薔薗薙薛薜薟薤薦薨薩薪薫薬薭薮薯薹薺藁藉藍藏藐藕藜藝藤藥藦藩藪藷藹藺藻藿蘂蘄蘆蘇蘊蘋蘐蘑蘖蘗蘚蘞蘩蘭蘯蘸蘿虎虐虔處虗虚虜虞號虢虧虫虬虱虹虻蚊蚋蚌蚓蚕蚣蚤蚩蚪蚫蚯蚰蛄蛆蛇蛉蛋蛍蛎蛔蛙蛛蛞蛟蛤蛩蛭蛮蛯蛸蛹蛻蛼蛾蜀蜂蜃蜆蜈蜉蜊蜍蜑蜒蜓蜘蜚蜜蜥蜩蜱蜴蜷蜺蜻蜾蜿蝀蝃蝉蝋蝌蝎蝓蝕蝗蝘蝙蝟蝠蝣蝦蝨蝪蝮蝯蝴蝶蝸蝿螂融螟螢螫螭螳螺螻螽蟀蟄蟆蟇蟋蟎蟐蟒蟖蟠蟥蟬蟯蟲蟷蟹蟻蟾蠃蠅蠍蠑蠓蠕蠖蠟蠡蠢蠣蠧蠱蠲蠶蠹蠻血衄衆行衍衒術街衙衛衝衞衡衢衣衤表衫衰衲衵衷衽衾衿袁袂袈袋袍袒袖袗袘袙袞袢袤袪被袰袱袴袵袷袿裀裁裂裃裄装裎裏裒裔裕裘裙補裝裟裡裨裰裱裲裳裴裵裸裹裼製裾褂褄複褊褌褐褒褓褘褙褚褜褥褪褫褰褲褶褸褻襁襄襖襞襟襠襤襦襯襲襴襷襾西要覃覆覇覈覊見規視覗覘覚覡覦覧覩親覬覯覲観覺覽覿觀角觔觚觜觝解触觫觳觴觸觿言訂訃計訊訌討訓託記訛訝訟訢訣訥訪設許訳訴訶診註証詁詆詈詐詑詒詔評詛詞詠詡詢詣試詩詫詬詭詮詰話該詳詵詹詼誂誄誅誇誉誌認誑誓誕誘語誠誡誣誤誥誦誨説読誰課誹誼誾調諂諄談請諌諍諏諒論諗諚諛諜諡諢諤諦諧諫諭諮諱諳諶諷諸諺諾謀謁謂謄謇謌謎謐謔謖謗謙謚講謝謠謡謦謨謫謬謳謹謾譃證譌譎譏譔識譙譚譜譞譟警譫譬譯議譲譴護譽讀讃變讌讎讐讒讓讖讙讚讜谷谺谿豁豆豇豈豉豊豌豎豐豕豚象豨豪豫豬豳豸豹豺豼貂貅貉貌貎貔貘貝貞負財貢貧貨販貪貫責貮貯貰貲貴貶買貸費貼貽貿賀賁賂賃賄資賈賊賍賎賑賓賚賛賜賞賠賡賢賣賤賦質賭賴賺購賽贄贅贇贈贊贋贍贏贐贓贔贖贛赤赦赧赫赭走赳赴起趁超越趙趣趨足趺趾跋跌跎跏跑跖跚跛距跟跡跣跨跪跫路跳践跼踈踉踊踏踐踔踝踞踠踣踪踰踴踵蹂蹄蹇蹈蹉蹊蹋蹌蹐蹕蹙蹟蹠蹣蹤蹬蹰蹲蹴蹶蹼躁躄躅躇躊躋躍躑躓躔躙躡躪身躬躯躰躱躶躾軀軈車軋軌軍軒軛軟転軫軸軺軻軼軽軾較載輊輌輒輓輔輕輙輛輜輝輞輟輦輩輪輯輳輸輻輾輿轀轂轄轅轆轉轌轍轎轔轗轜轟轡轢轤辛辜辞辟辣辦辧辨辭辮辯辰辱農辶辷辺辻込辿迂迄迅迎近返迚迢迤迦迩迪迫迭迯述迴迷迸迹迺追退送适逃逅逆逈逋逍透逐逑逓途逕逗這通逝逞速造逡逢連逮週進逵逶逸逹逼遁遂遅遇遉遊運遍過遏遐遑道達違遙遜遞遠遡遣遥適遭遮遯遲遵遶遷選遺遼遽避邀邁邂邃還邅邇邈邉邊邏邑邕邗邙邢那邦邨邪邯邱邳邵邶邸邾郁郃郈郊郎郗郝郞郡郢部郭郯郵郷都郿鄂鄄鄒鄕鄘鄙鄧鄭鄯鄰鄲鄴鄽酈酉酊酋酌配酎酒酔酖酗酘酛酡酢酣酤酥酩酪酬酲酵酷酸醂醇醉醋醍醐醒醗醜醢醤醪醫醬醱醴醵醸醺釀釁采釈釉釋里重野量釐金釗釘釜針釟釡釣釦釧釭釵釼鈍鈎鈐鈑鈔鈕鈞鈴鈷鈸鈺鈿鉀鉄鉅鉇鉈鉉鉊鉋鉏鉗鉛鉞鉢鉤鉦鉧鉱鉾銀銃銅銈銍銑銓銕銖銘銙銚銛銜銭銷銹鋏鋒鋤鋧鋩鋪鋭鋲鋳鋸鋹鋺鋼錄錆錐錕錘錙錚錠錡錢錣錥錦錨錫錬錮錯録錵錺錻鍈鍋鍍鍑鍔鍛鍜鍬鍮鍵鍼鍾鎌鎔鎖鎗鎚鎛鎧鎬鎭鎮鎰鎹鏃鏈鏑鏖鏗鏘鏝鏞鏟鏡鏢鏤鏧鏨鏱鏻鏽鐃鐐鐓鐔鐗鐘鐙鐚鐡鐫鐵鐶鐸鐺鑁鑄鑅鑑鑒鑓鑚鑛鑞鑠鑢鑫鑰鑵鑷鑼鑽鑾鑿長镹門閂閃閇閉閊開閎閏閑閒間閔閖閘閙関閣閤閥閦閧閨閩閫閭閲閹閻閼閽閾闇闈闊闌闍闐闔闕闖闘關闡闢闥阜阡阨阪阮阯防阳阻阿陀陂附陋陌降限陖陘陛陜陝陞陟院陣除陥陪陬陰陲陳陵陶陷陸険陽隂隄隅隆隈隊隋隍階随隔隕隗隘隙際障隠隣隧隨險隱隲隴隷隸隹隺隻隼雀雁雄雅集雇雉雋雌雍雎雑雒雕雖雙雛雜雝雞離難雧雨雩雪雫雯雰雲零雷雹電需霄霆震霊霍霏霑霓霖霙霜霞霤霧霰露霸霹霽霾靁靂靄靆靈靉靏靑青靖静靚靛靜非靠靡面靤靨革靫靭靱靳靴靹靺靼鞁鞄鞅鞆鞋鞍鞏鞘鞚鞜鞠鞢鞣鞦鞨鞭鞮鞳鞴鞺韃韆韋韓韜韞韮音韴韵韶韻響頁頂頃項順須頊頌頎頏預頑頒頓頗領頚頞頡頣頤頫頬頭頰頴頷頸頻頼頽顆顋題額顎顏顓顔顕顗願顚顛類顥顧顫顬顯顰顱顳顴風颪颯颱颶颷颺颼飃飄飆飛飜食飡飢飩飫飭飮飯飲飴飼飽飾餃餅餉養餌餐餒餓餘餝餞餠餡餤館餬餮餼饂饅饉饋饌饐饑饒饕饗饜首馗馘香馞馥馨馬馭馮馳馴駁駄駅駆駈駐駑駒駕駘駛駝駟駢駭駮駱駸駻駿騁騅騎騏騒験騖騙騨騫騰騷騾驀驃驅驍驒驕驗驚驛驟驢驤驥驩驪骨骰骶骸骼髀髄髏髑髓體高髙髞髠髢髣髦髪髫髭髮髯髱髴髷髻鬂鬆鬐鬒鬘鬚鬟鬠鬢鬣鬧鬨鬩鬪鬮鬱鬲鬳鬻鬼魁魂魃魄魅魋魍魎魏魑魔魘魚魯魳魴魹鮃鮄鮎鮐鮑鮒鮓鮞鮟鮠鮧鮨鮪鮫鮭鮮鮴鮹鯆鯉鯊鯎鯒鯔鯖鯛鯡鯣鯥鯨鯯鯰鯱鯵鯷鯸鯺鰆鰈鰉鰊鰌鰍鰐鰒鰓鰕鰡鰣鰤鰥鰩鰭鰮鰯鰰鰱鰲鰶鰹鰺鰻鱁鱇鱈鱉鱏鱒鱗鱚鱠鱧鱮鱲鱶鱷鱸鳥鳧鳩鳬鳰鳳鳴鳶鴃鴇鴈鴉鴎鴒鴗鴛鴝鴟鴣鴦鴨鴫鴬鴻鴾鴿鵃鵄鵈鵐鵑鵒鵙鵜鵝鵞鵟鵠鵡鵤鵬鵯鵲鵺鵼鵾鶉鶏鶚鶩鶫鶯鶲鶴鶸鶺鶿鷀鷁鷄鷉鷓鷗鷙鷦鷧鷭鷯鷲鷸鷹鷺鷽鸊鸕鸚鸛鸝鸞鹵鹸鹹鹼鹽鹿麁麈麋麐麒麓麕麗麝麟麤麥麦麩麪麭麴麵麸麹麺麻麼麽麾麿黃黄黌黍黎黐黑黒黔默黙黛黜黝點黠黥黨黯黴黷黻黽鼇鼈鼎鼓鼕鼙鼠鼬鼯鼹鼻鼽鼾齅齊齋齎齒齕齝齟齡齢齣齦齧齪齬齲齶齷龍龐龔龕龗龜龝龠龢글한﨑﨟﹅﹆𠮟𡌛𡢳𢌞𢪸𥝱𧘔𩹉

yomitoku/table_structure_recognizer.py

@@ -0,0 +1,244 @@
+from typing import List, Union
+
+import cv2
+import torch
+import torchvision.transforms as T
+from PIL import Image
+from pydantic import conlist
+
+from .base import BaseModelCatalog, BaseModule, BaseSchema
+from .configs import TableStructureRecognizerRTDETRv2Config
+from .layout_parser import filter_contained_rectangles_within_category
+from .models import RTDETRv2
+from .postprocessor import RTDETRPostProcessor
+from .utils.misc import calc_intersection, filter_by_flag, is_contained
+from .utils.visualizer import table_visualizer
+
+
+class TableStructureRecognizerModelCatalog(BaseModelCatalog):
+    def __init__(self):
+        super().__init__()
+        self.register("rtdetrv2", TableStructureRecognizerRTDETRv2Config, RTDETRv2)
+
+
+class TableCellSchema(BaseSchema):
+    col: int
+    row: int
+    col_span: int
+    row_span: int
+    box: conlist(int, min_length=4, max_length=4)
+    contents: Union[str, None]
+
+
+class TableStructureRecognizerSchema(BaseSchema):
+    box: conlist(int, min_length=4, max_length=4)
+    n_row: int
+    n_col: int
+    cells: List[TableCellSchema]
+    order: int
+
+
+def extract_cells(row_boxes, col_boxes):
+    cells = []
+    for i, row_box in enumerate(row_boxes):
+        for j, col_box in enumerate(col_boxes):
+            intersection = calc_intersection(row_box, col_box)
+            if intersection is None:
+                continue
+
+            cells.append(
+                {
+                    "col": j + 1,
+                    "row": i + 1,
+                    "col_span": 1,
+                    "row_span": 1,
+                    "box": intersection,
+                    "contents": None,
+                }
+            )
+
+    return cells
+
+
+def filter_contained_cells_within_spancell(cells, span_boxes):
+    check_list = [True] * len(cells)
+    child_boxes = [[] for _ in range(len(span_boxes))]
+    for i, span_box in enumerate(span_boxes):
+        for j, sub_cell in enumerate(cells):
+            if is_contained(span_box, sub_cell["box"]):
+                check_list[j] = False
+                child_boxes[i].append(sub_cell)
+
+    cells = filter_by_flag(cells, check_list)
+
+    for i, span_box in enumerate(span_boxes):
+        child_box = child_boxes[i]
+
+        if len(child_box) == 0:
+            continue
+
+        row = min([box["row"] for box in child_box])
+        col = min([box["col"] for box in child_box])
+        row_span = max([box["row"] for box in child_box]) - row + 1
+        col_span = max([box["col"] for box in child_box]) - col + 1
+
+        span_box = list(map(int, span_box))
+
+        cells.append(
+            {
+                "col": col,
+                "row": row,
+                "col_span": col_span,
+                "row_span": row_span,
+                "box": span_box,
+                "contents": None,
+            }
+        )
+
+    cells = sorted(cells, key=lambda x: (x["row"], x["col"]))
+    return cells
+
+
+class TableStructureRecognizer(BaseModule):
+    model_catalog = TableStructureRecognizerModelCatalog()
+
+    def __init__(
+        self,
+        model_name="rtdetrv2",
+        path_cfg=None,
+        device="cuda",
+        visualize=False,
+        from_pretrained=True,
+    ):
+        super().__init__()
+        self.load_model(
+            model_name,
+            path_cfg,
+            from_pretrained=True,
+        )
+        self.device = device
+        self.visualize = visualize
+
+        self.model.eval()
+        self.model.to(self.device)
+
+        self.postprocessor = RTDETRPostProcessor(
+            num_classes=self._cfg.RTDETRTransformerv2.num_classes,
+            num_top_queries=self._cfg.RTDETRTransformerv2.num_queries,
+        )
+
+        self.transforms = T.Compose(
+            [
+                T.Resize(self._cfg.data.img_size),
+                T.ToTensor(),
+            ]
+        )
+
+        self.thresh_score = self._cfg.thresh_score
+
+        self.label_mapper = {
+            id: category for id, category in enumerate(self._cfg.category)
+        }
+
+    def preprocess(self, img, boxes):
+        cv_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        table_imgs = []
+        for box in boxes:
+            x1, y1, x2, y2 = map(int, box)
+            table_img = cv_img[y1:y2, x1:x2, :]
+            th, hw = table_img.shape[:2]
+            table_img = Image.fromarray(table_img)
+            img_tensor = self.transforms(table_img)[None].to(self.device)
+            table_imgs.append(
+                {
+                    "tensor": img_tensor,
+                    "size": (th, hw),
+                    "offset": (x1, y1),
+                }
+            )
+        return table_imgs
+
+    def postprocess(self, preds, data):
+        h, w = data["size"]
+        orig_size = torch.tensor([w, h])[None].to(self.device)
+        outputs = self.postprocessor(preds, orig_size, self.thresh_score)
+
+        preds = outputs[0]
+        scores = preds["scores"]
+        boxes = preds["boxes"]
+        labels = preds["labels"]
+
+        category_elements = {category: [] for category in self.label_mapper.values()}
+        for box, score, label in zip(boxes, scores, labels):
+            category = self.label_mapper[label.item()]
+            box = box.astype(int).tolist()
+
+            box[0] += data["offset"][0]
+            box[1] += data["offset"][1]
+            box[2] += data["offset"][0]
+            box[3] += data["offset"][1]
+
+            category_elements[category].append(
+                {
+                    "box": box,
+                    "score": float(score),
+                }
+            )
+
+        category_elements = filter_contained_rectangles_within_category(
+            category_elements
+        )
+
+        cells, n_row, n_col = self.extract_cell_elements(category_elements)
+
+        table_x, table_y = data["offset"]
+        table_x2 = table_x + data["size"][1]
+        table_y2 = table_y + data["size"][0]
+        table_box = [table_x, table_y, table_x2, table_y2]
+
+        table = {
+            "box": table_box,
+            "n_row": n_row,
+            "n_col": n_col,
+            "cells": cells,
+            "order": 0,
+        }
+
+        results = TableStructureRecognizerSchema(**table)
+
+        return results
+
+    def extract_cell_elements(self, elements):
+        row_boxes = [element["box"] for element in elements["row"]]
+        col_boxes = [element["box"] for element in elements["col"]]
+        span_boxes = [element["box"] for element in elements["span"]]
+
+        row_boxes = sorted(row_boxes, key=lambda x: x[1])
+        col_boxes = sorted(col_boxes, key=lambda x: x[0])
+
+        cells = extract_cells(row_boxes, col_boxes)
+        cells = filter_contained_cells_within_spancell(cells, span_boxes)
+
+        return cells, len(row_boxes), len(col_boxes)
+
+    def __call__(self, img, table_boxes, vis=None):
+        img_tensors = self.preprocess(img, table_boxes)
+        outputs = []
+        for data in img_tensors:
+            with torch.inference_mode():
+                pred = self.model(data["tensor"])
+            table = self.postprocess(pred, data)
+            outputs.append(table)
+
+        if vis is None and self.visualize:
+            vis = img.copy()
+
+        if self.visualize:
+            for table in outputs:
+                vis = table_visualizer(
+                    vis,
+                    table,
+                )
+
+        return outputs, vis