yomitoku 0.4.1__py3-none-any.whl → 0.7.4__py3-none-any.whl

yomitoku/ocr.py

@@ -16,16 +16,37 @@ class WordPrediction(BaseSchema):
     )
     content: str
     direction: str
-    det_score: float
     rec_score: float
+    det_score: float
 
 
 class OCRSchema(BaseSchema):
     words: List[WordPrediction]
 
 
+def ocr_aggregate(det_outputs, rec_outputs):
+    words = []
+    for points, det_score, pred, rec_score, direction in zip(
+        det_outputs.points,
+        det_outputs.scores,
+        rec_outputs.contents,
+        rec_outputs.scores,
+        rec_outputs.directions,
+    ):
+        words.append(
+            {
+                "points": points,
+                "content": pred,
+                "direction": direction,
+                "det_score": det_score,
+                "rec_score": rec_score,
+            }
+        )
+    return words
+
+
 class OCR:
-    def __init__(self, configs=None, device="cuda", visualize=False):
+    def __init__(self, configs={}, device="cuda", visualize=False):
         text_detector_kwargs = {
             "device": device,
             "visualize": visualize,
@@ -36,10 +57,6 @@ class OCR:
         }
 
         if isinstance(configs, dict):
-            assert (
-                "text_detector" in configs or "text_recognizer" in configs
-            ), "Invalid config key. Please check the config keys."
-
             if "text_detector" in configs:
                 text_detector_kwargs.update(configs["text_detector"])
             if "text_recognizer" in configs:
@@ -52,26 +69,6 @@ class OCR:
         self.detector = TextDetector(**text_detector_kwargs)
         self.recognizer = TextRecognizer(**text_recognizer_kwargs)
 
-    def aggregate(self, det_outputs, rec_outputs):
-        words = []
-        for points, det_score, pred, rec_score, direction in zip(
-            det_outputs.points,
-            det_outputs.scores,
-            rec_outputs.contents,
-            rec_outputs.scores,
-            rec_outputs.directions,
-        ):
-            words.append(
-                {
-                    "points": points,
-                    "content": pred,
-                    "direction": direction,
-                    "det_score": det_score,
-                    "rec_score": rec_score,
-                }
-            )
-        return words
-
     def __call__(self, img):
         """_summary_
 
@@ -82,6 +79,6 @@ class OCR:
         det_outputs, vis = self.detector(img)
         rec_outputs, vis = self.recognizer(img, det_outputs.points, vis=vis)
 
-        outputs = {"words": self.aggregate(det_outputs, rec_outputs)}
+        outputs = {"words": ocr_aggregate(det_outputs, rec_outputs)}
         results = OCRSchema(**outputs)
         return results, vis

yomitoku/postprocessor/dbnet_postporcessor.py

@@ -1,13 +1,12 @@
 import cv2
+import math
 import numpy as np
 import pyclipper
 from shapely.geometry import Polygon
 
 
 class DBnetPostProcessor:
-    def __init__(
-        self, min_size, thresh, box_thresh, max_candidates, unclip_ratio
-    ):
+    def __init__(self, min_size, thresh, box_thresh, max_candidates, unclip_ratio):
         self.min_size = min_size
         self.thresh = thresh
         self.box_thresh = box_thresh
@@ -24,9 +23,7 @@ class DBnetPostProcessor:
         pred = preds["binary"][0]
         segmentation = self.binarize(pred)[0]
         height, width = image_size
-        quads, scores = self.boxes_from_bitmap(
-            pred, segmentation, width, height
-        )
+        quads, scores = self.boxes_from_bitmap(pred, segmentation, width, height)
         return quads, scores
 
     def binarize(self, pred):
@@ -65,9 +62,7 @@ class DBnetPostProcessor:
             if self.box_thresh > score:
                 continue
 
-            box = self.unclip(points, unclip_ratio=self.unclip_ratio).reshape(
-                -1, 1, 2
-            )
+            box = self.unclip(points, unclip_ratio=self.unclip_ratio).reshape(-1, 1, 2)
             box, sside = self.get_mini_boxes(box)
             if sside < self.min_size + 2:
                 continue
@@ -76,9 +71,7 @@ class DBnetPostProcessor:
                 dest_width = dest_width.item()
                 dest_height = dest_height.item()
 
-            box[:, 0] = np.clip(
-                np.round(box[:, 0] / width * dest_width), 0, dest_width
-            )
+            box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
             box[:, 1] = np.clip(
                 np.round(box[:, 1] / height * dest_height), 0, dest_height
             )
@@ -88,9 +81,17 @@ class DBnetPostProcessor:
 
         return boxes, scores
 
-    def unclip(self, box, unclip_ratio=1.5):
+    def unclip(self, box, unclip_ratio=7):
+        # 小さい文字が見切れやすい、大きい文字のマージンが過度に大きくなる等の課題がある
+        # 対応として、文字の大きさに応じて、拡大パラメータを動的に変更する
+        # Note: こののルールはヒューリスティックで理論的根拠はない
         poly = Polygon(box)
-        distance = poly.area * unclip_ratio / poly.length
+        width = box[:, 0].max() - box[:, 0].min()
+        height = box[:, 1].max() - box[:, 1].min()
+        box_dist = min(width, height)
+        ratio = unclip_ratio / math.sqrt(box_dist)
+
+        distance = poly.area * ratio / poly.length
         offset = pyclipper.PyclipperOffset()
         offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
         expanded = np.array(offset.Execute(distance))

yomitoku/postprocessor/parseq_tokenizer.py

@@ -122,7 +122,5 @@ class ParseqTokenizer(BaseTokenizer):
             eos_idx = len(ids)  # Nothing to truncate.
         # Truncate after EOS
         ids = ids[:eos_idx]
-        probs = probs[
-            : eos_idx + 1
-        ]  # but include prob. for EOS (if it exists)
+        probs = probs[: eos_idx + 1]  # but include prob. for EOS (if it exists)
         return probs, ids

yomitoku/postprocessor/rtdetr_postprocessor.py

@@ -1,4 +1,17 @@
-"""Copyright(c) 2023 lyuwenyu. All Rights Reserved."""
+# Copyright 2023 lyuwenyu
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 
 import torch
 import torch.nn as nn

yomitoku/table_structure_recognizer.py

@@ -1,11 +1,16 @@
 from typing import List, Union
 
 import cv2
+import os
+import onnx
+import onnxruntime
 import torch
 import torchvision.transforms as T
 from PIL import Image
 from pydantic import conlist
 
+from .constants import ROOT_DIR
+
 from .base import BaseModelCatalog, BaseModule, BaseSchema
 from .configs import TableStructureRecognizerRTDETRv2Config
 from .layout_parser import filter_contained_rectangles_within_category
@@ -30,11 +35,19 @@ class TableCellSchema(BaseSchema):
     contents: Union[str, None]
 
 
+class TableLineSchema(BaseSchema):
+    box: conlist(int, min_length=4, max_length=4)
+    score: float
+
+
 class TableStructureRecognizerSchema(BaseSchema):
     box: conlist(int, min_length=4, max_length=4)
     n_row: int
     n_col: int
+    rows: List[TableLineSchema]
+    cols: List[TableLineSchema]
     cells: List[TableCellSchema]
+    spans: List[TableLineSchema]
     order: int
 
 
@@ -109,12 +122,13 @@ class TableStructureRecognizer(BaseModule):
         device="cuda",
         visualize=False,
         from_pretrained=True,
+        infer_onnx=False,
     ):
         super().__init__()
         self.load_model(
             model_name,
             path_cfg,
-            from_pretrained=True,
+            from_pretrained=from_pretrained,
         )
         self.device = device
         self.visualize = visualize
@@ -140,6 +154,45 @@ class TableStructureRecognizer(BaseModule):
             id: category for id, category in enumerate(self._cfg.category)
         }
 
+        self.infer_onnx = infer_onnx
+        if infer_onnx:
+            name = self._cfg.hf_hub_repo.split("/")[-1]
+            path_onnx = f"{ROOT_DIR}/onnx/{name}.onnx"
+            if not os.path.exists(path_onnx):
+                self.convert_onnx(path_onnx)
+
+            self.model = None
+
+            model = onnx.load(path_onnx)
+            if torch.cuda.is_available() and device == "cuda":
+                self.sess = onnxruntime.InferenceSession(
+                    model.SerializeToString(), providers=["CUDAExecutionProvider"]
+                )
+            else:
+                self.sess = onnxruntime.InferenceSession(model.SerializeToString())
+
+        if self.model is not None:
+            self.model.to(self.device)
+
+    def convert_onnx(self, path_onnx):
+        dynamic_axes = {
+            "input": {0: "batch_size"},
+            "output": {0: "batch_size"},
+        }
+
+        img_size = self._cfg.data.img_size
+        dummy_input = torch.randn(1, 3, *img_size, requires_grad=True)
+
+        torch.onnx.export(
+            self.model,
+            dummy_input,
+            path_onnx,
+            opset_version=16,
+            input_names=["input"],
+            output_names=["pred_logits", "pred_boxes"],
+            dynamic_axes=dynamic_axes,
+        )
+
     def preprocess(self, img, boxes):
         cv_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
@@ -149,7 +202,7 @@ class TableStructureRecognizer(BaseModule):
             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)
+            img_tensor = self.transforms(table_img)[None]
             table_imgs.append(
                 {
                     "tensor": img_tensor,
@@ -190,7 +243,7 @@ class TableStructureRecognizer(BaseModule):
             category_elements
         )
 
-        cells, n_row, n_col = self.extract_cell_elements(category_elements)
+        cells, rows, cols, spans = self.extract_cell_elements(category_elements)
 
         table_x, table_y = data["offset"]
         table_x2 = table_x + data["size"][1]
@@ -199,8 +252,11 @@ class TableStructureRecognizer(BaseModule):
 
         table = {
             "box": table_box,
-            "n_row": n_row,
-            "n_col": n_col,
+            "n_row": len(rows),
+            "n_col": len(cols),
+            "rows": rows,
+            "cols": cols,
+            "spans": spans,
             "cells": cells,
             "order": 0,
         }
@@ -220,16 +276,33 @@ class TableStructureRecognizer(BaseModule):
         cells = extract_cells(row_boxes, col_boxes)
         cells = filter_contained_cells_within_spancell(cells, span_boxes)
 
-        return cells, len(row_boxes), len(col_boxes)
+        rows = sorted(elements["row"], key=lambda x: x["box"][1])
+        cols = sorted(elements["col"], key=lambda x: x["box"][0])
+        spans = sorted(elements["span"], key=lambda x: x["box"][1])
+
+        return cells, rows, cols, spans
 
     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"])
+            if self.infer_onnx:
+                input = data["tensor"].numpy()
+                results = self.sess.run(None, {"input": input})
+                pred = {
+                    "pred_logits": torch.tensor(results[0]).to(self.device),
+                    "pred_boxes": torch.tensor(results[1]).to(self.device),
+                }
+
+            else:
+                with torch.inference_mode():
+                    data["tensor"] = data["tensor"].to(self.device)
+                    pred = self.model(data["tensor"])
+
             table = self.postprocess(pred, data)
-            outputs.append(table)
+
+            if table.n_row > 0 and table.n_col > 0:
+                outputs.append(table)
 
         if vis is None and self.visualize:
             vis = img.copy()

yomitoku/text_detector.py

@@ -2,6 +2,7 @@ from typing import List
 
 import numpy as np
 import torch
+import os
 from pydantic import conlist
 
 from .base import BaseModelCatalog, BaseModule, BaseSchema
@@ -14,6 +15,10 @@ from .data.functions import (
 from .models import DBNet
 from .postprocessor import DBnetPostProcessor
 from .utils.visualizer import det_visualizer
+from .constants import ROOT_DIR
+
+import onnx
+import onnxruntime
 
 
 class TextDetectorModelCatalog(BaseModelCatalog):
@@ -43,21 +48,60 @@ class TextDetector(BaseModule):
         device="cuda",
         visualize=False,
         from_pretrained=True,
+        infer_onnx=False,
     ):
         super().__init__()
         self.load_model(
             model_name,
             path_cfg,
-            from_pretrained=True,
+            from_pretrained=from_pretrained,
         )
 
         self.device = device
         self.visualize = visualize
 
         self.model.eval()
-        self.model.to(self.device)
-
         self.post_processor = DBnetPostProcessor(**self._cfg.post_process)
+        self.infer_onnx = infer_onnx
+
+        if infer_onnx:
+            name = self._cfg.hf_hub_repo.split("/")[-1]
+            path_onnx = f"{ROOT_DIR}/onnx/{name}.onnx"
+            if not os.path.exists(path_onnx):
+                self.convert_onnx(path_onnx)
+
+            self.model = None
+
+            model = onnx.load(path_onnx)
+            if torch.cuda.is_available() and device == "cuda":
+                self.sess = onnxruntime.InferenceSession(
+                    model.SerializeToString(), providers=["CUDAExecutionProvider"]
+                )
+            else:
+                self.sess = onnxruntime.InferenceSession(model.SerializeToString())
+
+            self.model = None
+
+        if self.model is not None:
+            self.model.to(self.device)
+
+    def convert_onnx(self, path_onnx):
+        dynamic_axes = {
+            "input": {0: "batch_size", 2: "height", 3: "width"},
+            "output": {0: "batch_size", 2: "height", 3: "width"},
+        }
+
+        dummy_input = torch.randn(1, 3, 256, 256, requires_grad=True)
+
+        torch.onnx.export(
+            self.model,
+            dummy_input,
+            path_onnx,
+            opset_version=14,
+            input_names=["input"],
+            output_names=["output"],
+            dynamic_axes=dynamic_axes,
+        )
 
     def preprocess(self, img):
         img = img.copy()
@@ -81,9 +125,15 @@ class TextDetector(BaseModule):
 
         ori_h, ori_w = img.shape[:2]
         tensor = self.preprocess(img)
-        tensor = tensor.to(self.device)
-        with torch.inference_mode():
-            preds = self.model(tensor)
+
+        if self.infer_onnx:
+            input = tensor.numpy()
+            results = self.sess.run(["output"], {"input": input})
+            preds = {"binary": torch.tensor(results[0])}
+        else:
+            with torch.inference_mode():
+                tensor = tensor.to(self.device)
+                preds = self.model(tensor)
 
         quads, scores = self.postprocess(preds, (ori_h, ori_w))
         outputs = {"points": quads, "scores": scores}
@@ -93,9 +143,9 @@ class TextDetector(BaseModule):
         vis = None
         if self.visualize:
             vis = det_visualizer(
-                preds,
                 img,
                 quads,
+                preds=preds,
                 vis_heatmap=self._cfg.visualize.heatmap,
                 line_color=tuple(self._cfg.visualize.color[::-1]),
             )

yomitoku/text_recognizer.py

@@ -2,21 +2,28 @@ from typing import List
 
 import numpy as np
 import torch
+import os
+import unicodedata
 from pydantic import conlist
 
 from .base import BaseModelCatalog, BaseModule, BaseSchema
-from .configs import TextRecognizerPARSeqConfig
+from .configs import TextRecognizerPARSeqConfig, TextRecognizerPARSeqSmallConfig
 from .data.dataset import ParseqDataset
 from .models import PARSeq
 from .postprocessor import ParseqTokenizer as Tokenizer
 from .utils.misc import load_charset
 from .utils.visualizer import rec_visualizer
 
+from .constants import ROOT_DIR
+import onnx
+import onnxruntime
+
 
 class TextRecognizerModelCatalog(BaseModelCatalog):
     def __init__(self):
         super().__init__()
         self.register("parseq", TextRecognizerPARSeqConfig, PARSeq)
+        self.register("parseq-small", TextRecognizerPARSeqSmallConfig, PARSeq)
 
 
 class TextRecognizerSchema(BaseSchema):
@@ -42,36 +49,91 @@ class TextRecognizer(BaseModule):
         device="cuda",
         visualize=False,
         from_pretrained=True,
+        infer_onnx=False,
     ):
         super().__init__()
         self.load_model(
             model_name,
             path_cfg,
-            from_pretrained=True,
+            from_pretrained=from_pretrained,
         )
         self.charset = load_charset(self._cfg.charset)
         self.tokenizer = Tokenizer(self.charset)
 
         self.device = device
 
+        self.model.tokenizer = self.tokenizer
         self.model.eval()
-        self.model.to(self.device)
 
         self.visualize = visualize
 
+        self.infer_onnx = infer_onnx
+
+        if infer_onnx:
+            name = self._cfg.hf_hub_repo.split("/")[-1]
+            path_onnx = f"{ROOT_DIR}/onnx/{name}.onnx"
+            if not os.path.exists(path_onnx):
+                self.convert_onnx(path_onnx)
+
+            self.model = None
+
+            model = onnx.load(path_onnx)
+            if torch.cuda.is_available() and device == "cuda":
+                self.sess = onnxruntime.InferenceSession(
+                    model.SerializeToString(), providers=["CUDAExecutionProvider"]
+                )
+            else:
+                self.sess = onnxruntime.InferenceSession(model.SerializeToString())
+
+        if self.model is not None:
+            self.model.to(self.device)
+
     def preprocess(self, img, polygons):
         dataset = ParseqDataset(self._cfg, img, polygons)
-        dataloader = torch.utils.data.DataLoader(
-            dataset,
-            batch_size=self._cfg.data.batch_size,
-            shuffle=False,
-            num_workers=self._cfg.data.num_workers,
-        )
+        dataloader = self._make_mini_batch(dataset)
 
         return dataloader
 
+    def _make_mini_batch(self, dataset):
+        mini_batches = []
+        mini_batch = []
+        for data in dataset:
+            data = torch.unsqueeze(data, 0)
+            mini_batch.append(data)
+
+            if len(mini_batch) == self._cfg.data.batch_size:
+                mini_batches.append(torch.cat(mini_batch, 0))
+                mini_batch = []
+        else:
+            if len(mini_batch) > 0:
+                mini_batches.append(torch.cat(mini_batch, 0))
+
+        return mini_batches
+
+    def convert_onnx(self, path_onnx):
+        img_size = self._cfg.data.img_size
+        input = torch.randn(1, 3, *img_size, requires_grad=True)
+        dynamic_axes = {
+            "input": {0: "batch_size"},
+            "output": {0: "batch_size"},
+        }
+
+        self.model.export_onnx = True
+        torch.onnx.export(
+            self.model,
+            input,
+            path_onnx,
+            opset_version=14,
+            input_names=["input"],
+            output_names=["output"],
+            do_constant_folding=True,
+            dynamic_axes=dynamic_axes,
+        )
+
     def postprocess(self, p, points):
         pred, score = self.tokenizer.decode(p)
+        pred = [unicodedata.normalize("NFKC", x) for x in pred]
+
         directions = []
         for point in points:
             point = np.array(point)
@@ -98,13 +160,19 @@ class TextRecognizer(BaseModule):
         scores = []
         directions = []
         for data in dataloader:
-            data = data.to(self.device)
-            with torch.inference_mode():
-                p = self.model(self.tokenizer, data).softmax(-1)
-                pred, score, direction = self.postprocess(p, points)
-                preds.extend(pred)
-                scores.extend(score)
-                directions.extend(direction)
+            if self.infer_onnx:
+                input = data.numpy()
+                results = self.sess.run(["output"], {"input": input})
+                p = torch.tensor(results[0])
+            else:
+                with torch.inference_mode():
+                    data = data.to(self.device)
+                    p = self.model(data).softmax(-1)
+
+            pred, score, direction = self.postprocess(p, points)
+            preds.extend(pred)
+            scores.extend(score)
+            directions.extend(direction)
 
         outputs = {
             "contents": preds,

yomitoku/utils/misc.py

@@ -1,5 +1,5 @@
 def load_charset(charset_path):
-    with open(charset_path, "r") as f:
+    with open(charset_path, "r", encoding="utf-8") as f:
         charset = f.read()
     return charset
 
@@ -9,6 +9,24 @@ def filter_by_flag(elements, flags):
     return [element for element, flag in zip(elements, flags) if flag]
 
 
+def calc_overlap_ratio(rect_a, rect_b):
+    intersection = calc_intersection(rect_a, rect_b)
+    if intersection is None:
+        return 0, None
+
+    ix1, iy1, ix2, iy2 = intersection
+
+    overlap_width = ix2 - ix1
+    overlap_height = iy2 - iy1
+    bx1, by1, bx2, by2 = rect_b
+
+    b_area = (bx2 - bx1) * (by2 - by1)
+    overlap_area = overlap_width * overlap_height
+
+    overlap_ratio = overlap_area / b_area
+    return overlap_ratio, intersection
+
+
 def is_contained(rect_a, rect_b, threshold=0.8):
     """二つの矩形A, Bが与えられたとき、矩形Bが矩形Aに含まれるかどうかを判定する。
     ずれを許容するため、重複率求め、thresholdを超える場合にTrueを返す。
@@ -23,20 +41,9 @@ def is_contained(rect_a, rect_b, threshold=0.8):
         bool: 矩形Bが矩形Aに含まれる場合True
     """
 
-    intersection = calc_intersection(rect_a, rect_b)
-    if intersection is None:
-        return False
-
-    ix1, iy1, ix2, iy2 = intersection
-
-    overlap_width = ix2 - ix1
-    overlap_height = iy2 - iy1
-    bx1, by1, bx2, by2 = rect_b
-
-    b_area = (bx2 - bx1) * (by2 - by1)
-    overlap_area = overlap_width * overlap_height
+    overlap_ratio, _ = calc_overlap_ratio(rect_a, rect_b)
 
-    if overlap_area / b_area > threshold:
+    if overlap_ratio > threshold:
         return True
 
     return False