yomitoku 0.5.3__py3-none-any.whl → 0.6.0__py3-none-any.whl

yomitoku/cli/main.py

@@ -104,6 +104,12 @@ def main():
         default="results",
         help="output directory",
     )
+    parser.add_argument(
+        "-l",
+        "--lite",
+        action="store_true",
+        help="if set, use lite model",
+    )
     parser.add_argument(
         "-d",
         "--device",
@@ -197,6 +203,15 @@ def main():
         },
     }
 
+    if args.lite:
+        configs["ocr"]["text_recognizer"]["model_name"] = "parseq-small"
+        configs["ocr"]["text_detector"]["infer_onnx"] = True
+
+        # Note: Text Detector以外はONNX推論よりもPyTorch推論の方が速いため、ONNX推論は行わない
+        # configs["ocr"]["text_recognizer"]["infer_onnx"] = True
+        # configs["layout_analyzer"]["table_structure_recognizer"]["infer_onnx"] = True
+        # configs["layout_analyzer"]["layout_parser"]["infer_onnx"] = True
+
     analyzer = DocumentAnalyzer(
         configs=configs,
         visualize=args.vis,

yomitoku/configs/__init__.py

@@ -4,10 +4,12 @@ from .cfg_table_structure_recognizer_rtdtrv2 import (
 )
 from .cfg_text_detector_dbnet import TextDetectorDBNetConfig
 from .cfg_text_recognizer_parseq import TextRecognizerPARSeqConfig
+from .cfg_text_recognizer_parseq_small import TextRecognizerPARSeqSmallConfig
 
 __all__ = [
     "TextDetectorDBNetConfig",
     "TextRecognizerPARSeqConfig",
     "LayoutParserRTDETRv2Config",
     "TableStructureRecognizerRTDETRv2Config",
+    "TextRecognizerPARSeqSmallConfig",
 ]

yomitoku/configs/cfg_text_recognizer_parseq_small.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass, field
+from typing import List
+
+from ..constants import ROOT_DIR
+
+
+@dataclass
+class Data:
+    num_workers: int = 4
+    batch_size: int = 128
+    img_size: List[int] = field(default_factory=lambda: [32, 800])
+
+
+@dataclass
+class Encoder:
+    patch_size: List[int] = field(default_factory=lambda: [16, 16])
+    num_heads: int = 8
+    embed_dim: int = 384
+    mlp_ratio: int = 4
+    depth: int = 9
+
+
+@dataclass
+class Decoder:
+    embed_dim: int = 384
+    num_heads: int = 8
+    mlp_ratio: int = 4
+    depth: int = 1
+
+
+@dataclass
+class Visualize:
+    font: str = str(ROOT_DIR + "/resource/MPLUS1p-Medium.ttf")
+    color: List[int] = field(default_factory=lambda: [0, 0, 255])  # RGB
+    font_size: int = 18
+
+
+@dataclass
+class TextRecognizerPARSeqSmallConfig:
+    hf_hub_repo: str = "KotaroKinoshita/yomitoku-text-recognizer-parseq-small-open-beta"
+    charset: str = str(ROOT_DIR + "/resource/charset.txt")
+    num_tokens: int = 7312
+    max_label_length: int = 100
+    decode_ar: int = 1
+    refine_iters: int = 1
+
+    data: Data = field(default_factory=Data)
+    encoder: Encoder = field(default_factory=Encoder)
+    decoder: Decoder = field(default_factory=Decoder)
+
+    visualize: Visualize = field(default_factory=Visualize)

yomitoku/layout_parser.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 LayoutParserRTDETRv2Config
 from .models import RTDETRv2
@@ -91,6 +96,7 @@ class LayoutParser(BaseModule):
         device="cuda",
         visualize=False,
         from_pretrained=True,
+        infer_onnx=False,
     ):
         super().__init__()
         self.load_model(model_name, path_cfg, from_pretrained)
@@ -119,11 +125,44 @@ class LayoutParser(BaseModule):
         }
 
         self.role = self._cfg.role
+        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)
+
+            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())
+
+    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):
         cv_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         img = Image.fromarray(cv_img)
-        img_tensor = self.transforms(img)[None].to(self.device)
+        img_tensor = self.transforms(img)[None]
         return img_tensor
 
     def postprocess(self, preds, image_size):
@@ -175,8 +214,19 @@ class LayoutParser(BaseModule):
         ori_h, ori_w = img.shape[:2]
         img_tensor = self.preprocess(img)
 
-        with torch.inference_mode():
-            preds = self.model(img_tensor)
+        if self.infer_onnx:
+            input = img_tensor.numpy()
+            results = self.sess.run(None, {"input": input})
+            preds = {
+                "pred_logits": torch.tensor(results[0]).to(self.device),
+                "pred_boxes": torch.tensor(results[1]).to(self.device),
+            }
+
+        else:
+            with torch.inference_mode():
+                img_tensor = img_tensor.to(self.device)
+                preds = self.model(img_tensor)
+
         results = self.postprocess(preds, (ori_h, ori_w))
 
         vis = None

yomitoku/models/layers/rtdetr_backbone.py

@@ -59,9 +59,7 @@ class ConvNormLayer(nn.Module):
 class BasicBlock(nn.Module):
     expansion = 1
 
-    def __init__(
-        self, ch_in, ch_out, stride, shortcut, act="relu", variant="b"
-    ):
+    def __init__(self, ch_in, ch_out, stride, shortcut, act="relu", variant="b"):
         super().__init__()
 
         self.shortcut = shortcut
@@ -100,9 +98,7 @@ class BasicBlock(nn.Module):
 class BottleNeck(nn.Module):
     expansion = 4
 
-    def __init__(
-        self, ch_in, ch_out, stride, shortcut, act="relu", variant="b"
-    ):
+    def __init__(self, ch_in, ch_out, stride, shortcut, act="relu", variant="b"):
         super().__init__()
 
         if variant == "a":
@@ -125,17 +121,13 @@ class BottleNeck(nn.Module):
                             ("pool", nn.AvgPool2d(2, 2, 0, ceil_mode=True)),
                             (
                                 "conv",
-                                ConvNormLayer(
-                                    ch_in, ch_out * self.expansion, 1, 1
-                                ),
+                                ConvNormLayer(ch_in, ch_out * self.expansion, 1, 1),
                             ),
                         ]
                     )
                 )
             else:
-                self.short = ConvNormLayer(
-                    ch_in, ch_out * self.expansion, 1, stride
-                )
+                self.short = ConvNormLayer(ch_in, ch_out * self.expansion, 1, stride)
 
         self.act = nn.Identity() if act is None else get_activation(act)
 
@@ -156,9 +148,7 @@ class BottleNeck(nn.Module):
 
 
 class Blocks(nn.Module):
-    def __init__(
-        self, block, ch_in, ch_out, count, stage_num, act="relu", variant="b"
-    ):
+    def __init__(self, block, ch_in, ch_out, count, stage_num, act="relu", variant="b"):
         super().__init__()
 
         self.blocks = nn.ModuleList()

yomitoku/models/layers/rtdetr_hybrid_encoder.py

@@ -252,9 +252,7 @@ class HybridEncoder(nn.Module):
         for in_channel in in_channels:
             if version == "v1":
                 proj = nn.Sequential(
-                    nn.Conv2d(
-                        in_channel, hidden_dim, kernel_size=1, bias=False
-                    ),
+                    nn.Conv2d(in_channel, hidden_dim, kernel_size=1, bias=False),
                     nn.BatchNorm2d(hidden_dim),
                 )
             elif version == "v2":
@@ -290,9 +288,7 @@ class HybridEncoder(nn.Module):
 
         self.encoder = nn.ModuleList(
             [
-                TransformerEncoder(
-                    copy.deepcopy(encoder_layer), num_encoder_layers
-                )
+                TransformerEncoder(copy.deepcopy(encoder_layer), num_encoder_layers)
                 for _ in range(len(use_encoder_idx))
             ]
         )
@@ -347,9 +343,7 @@ class HybridEncoder(nn.Module):
                 # self.register_buffer(f'pos_embed{idx}', pos_embed)
 
     @staticmethod
-    def build_2d_sincos_position_embedding(
-        w, h, embed_dim=256, temperature=10000.0
-    ):
+    def build_2d_sincos_position_embedding(w, h, embed_dim=256, temperature=10000.0):
         """ """
         grid_w = torch.arange(int(w), dtype=torch.float32)
         grid_h = torch.arange(int(h), dtype=torch.float32)
@@ -387,9 +381,7 @@ class HybridEncoder(nn.Module):
                         src_flatten.device
                     )
 
-                memory: torch.Tensor = self.encoder[i](
-                    src_flatten, pos_embed=pos_embed
-                )
+                memory: torch.Tensor = self.encoder[i](src_flatten, pos_embed=pos_embed)
                 proj_feats[enc_ind] = (
                     memory.permute(0, 2, 1)
                     .reshape(-1, self.hidden_dim, h, w)
@@ -401,13 +393,9 @@ class HybridEncoder(nn.Module):
         for idx in range(len(self.in_channels) - 1, 0, -1):
             feat_heigh = inner_outs[0]
             feat_low = proj_feats[idx - 1]
-            feat_heigh = self.lateral_convs[len(self.in_channels) - 1 - idx](
-                feat_heigh
-            )
+            feat_heigh = self.lateral_convs[len(self.in_channels) - 1 - idx](feat_heigh)
             inner_outs[0] = feat_heigh
-            upsample_feat = F.interpolate(
-                feat_heigh, scale_factor=2.0, mode="nearest"
-            )
+            upsample_feat = F.interpolate(feat_heigh, scale_factor=2.0, mode="nearest")
             inner_out = self.fpn_blocks[len(self.in_channels) - 1 - idx](
                 torch.concat([upsample_feat, feat_low], dim=1)
             )

yomitoku/models/layers/rtdetrv2_decoder.py

@@ -40,9 +40,7 @@ def inverse_sigmoid(x: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
 
 
 class MLP(nn.Module):
-    def __init__(
-        self, input_dim, hidden_dim, output_dim, num_layers, act="relu"
-    ):
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers, act="relu"):
         super().__init__()
         self.num_layers = num_layers
         h = [hidden_dim] * (num_layers - 1)
@@ -193,9 +191,7 @@ class MSDeformableAttention(nn.Module):
         elif reference_points.shape[-1] == 4:
             # reference_points [8, 480, None, 1,  4]
             # sampling_offsets [8, 480, 8,    12, 2]
-            num_points_scale = self.num_points_scale.to(
-                dtype=query.dtype
-            ).unsqueeze(-1)
+            num_points_scale = self.num_points_scale.to(dtype=query.dtype).unsqueeze(-1)
             offset = (
                 sampling_offsets
                 * num_points_scale
@@ -330,9 +326,7 @@ def deformable_attention_core_func_v2(
     _, Len_q, _, _, _ = sampling_locations.shape
 
     split_shape = [h * w for h, w in value_spatial_shapes]
-    value_list = (
-        value.permute(0, 2, 3, 1).flatten(0, 1).split(split_shape, dim=-1)
-    )
+    value_list = value.permute(0, 2, 3, 1).flatten(0, 1).split(split_shape, dim=-1)
 
     # sampling_offsets [8, 480, 8, 12, 2]
     if method == "default":
@@ -361,8 +355,7 @@ def deformable_attention_core_func_v2(
         elif method == "discrete":
             # n * m, seq, n, 2
             sampling_coord = (
-                sampling_grid_l * torch.tensor([[w, h]], device=value.device)
-                + 0.5
+                sampling_grid_l * torch.tensor([[w, h]], device=value.device) + 0.5
             ).to(torch.int64)
 
             # FIX ME? for rectangle input
@@ -389,9 +382,7 @@ def deformable_attention_core_func_v2(
     attn_weights = attention_weights.permute(0, 2, 1, 3).reshape(
         bs * n_head, 1, Len_q, sum(num_points_list)
     )
-    weighted_sample_locs = (
-        torch.concat(sampling_value_list, dim=-1) * attn_weights
-    )
+    weighted_sample_locs = torch.concat(sampling_value_list, dim=-1) * attn_weights
     output = weighted_sample_locs.sum(-1).reshape(bs, n_head * c, Len_q)
 
     return output.permute(0, 2, 1)
@@ -606,9 +597,7 @@ class RTDETRTransformerv2(nn.Module):
                         [
                             (
                                 "conv",
-                                nn.Conv2d(
-                                    in_channels, self.hidden_dim, 1, bias=False
-                                ),
+                                nn.Conv2d(in_channels, self.hidden_dim, 1, bias=False),
                             ),
                             (
                                 "norm",
@@ -689,13 +678,9 @@ class RTDETRTransformerv2(nn.Module):
                 torch.arange(h), torch.arange(w), indexing="ij"
             )
             grid_xy = torch.stack([grid_x, grid_y], dim=-1)
-            grid_xy = (grid_xy.unsqueeze(0) + 0.5) / torch.tensor(
-                [w, h], dtype=dtype
-            )
+            grid_xy = (grid_xy.unsqueeze(0) + 0.5) / torch.tensor([w, h], dtype=dtype)
             wh = torch.ones_like(grid_xy) * grid_size * (2.0**lvl)
-            lvl_anchors = torch.concat([grid_xy, wh], dim=-1).reshape(
-                -1, h * w, 4
-            )
+            lvl_anchors = torch.concat([grid_xy, wh], dim=-1).reshape(-1, h * w, 4)
             anchors.append(lvl_anchors)
 
         anchors = torch.concat(anchors, dim=1).to(device)
@@ -729,22 +714,18 @@ class RTDETRTransformerv2(nn.Module):
         )
 
         enc_topk_bboxes_list, enc_topk_logits_list = [], []
-        enc_topk_memory, enc_topk_logits, enc_topk_bbox_unact = (
-            self._select_topk(
-                output_memory,
-                enc_outputs_logits,
-                enc_outputs_coord_unact,
-                self.num_queries,
-            )
+        enc_topk_memory, enc_topk_logits, enc_topk_bbox_unact = self._select_topk(
+            output_memory,
+            enc_outputs_logits,
+            enc_outputs_coord_unact,
+            self.num_queries,
         )
 
         # if self.num_select_queries != self.num_queries:
         #     raise NotImplementedError('')
 
         if self.learn_query_content:
-            content = self.tgt_embed.weight.unsqueeze(0).tile(
-                [memory.shape[0], 1, 1]
-            )
+            content = self.tgt_embed.weight.unsqueeze(0).tile([memory.shape[0], 1, 1])
         else:
             content = enc_topk_memory.detach()
 
@@ -771,9 +752,7 @@ class RTDETRTransformerv2(nn.Module):
         topk: int,
     ):
         if self.query_select_method == "default":
-            _, topk_ind = torch.topk(
-                outputs_logits.max(-1).values, topk, dim=-1
-            )
+            _, topk_ind = torch.topk(outputs_logits.max(-1).values, topk, dim=-1)
 
         elif self.query_select_method == "one2many":
             _, topk_ind = torch.topk(outputs_logits.flatten(1), topk, dim=-1)
@@ -786,16 +765,12 @@ class RTDETRTransformerv2(nn.Module):
 
         topk_coords = outputs_coords_unact.gather(
             dim=1,
-            index=topk_ind.unsqueeze(-1).repeat(
-                1, 1, outputs_coords_unact.shape[-1]
-            ),
+            index=topk_ind.unsqueeze(-1).repeat(1, 1, outputs_coords_unact.shape[-1]),
         )
 
         topk_logits = outputs_logits.gather(
             dim=1,
-            index=topk_ind.unsqueeze(-1).repeat(
-                1, 1, outputs_logits.shape[-1]
-            ),
+            index=topk_ind.unsqueeze(-1).repeat(1, 1, outputs_logits.shape[-1]),
         )
 
         topk_memory = memory.gather(

yomitoku/models/parseq.py

@@ -22,7 +22,6 @@ from huggingface_hub import PyTorchModelHubMixin
 from timm.models.helpers import named_apply
 from torch import Tensor
 
-from ..postprocessor import ParseqTokenizer as Tokenizer
 from .layers.parseq_transformer import Decoder, Encoder, TokenEmbedding
 
 
@@ -123,7 +122,6 @@ class PARSeq(nn.Module, PyTorchModelHubMixin):
 
     def forward(
         self,
-        tokenizer: Tokenizer,
         images: Tensor,
         max_length: Optional[int] = None,
     ) -> Tensor:
@@ -150,11 +148,11 @@ class PARSeq(nn.Module, PyTorchModelHubMixin):
         if self.decode_ar:
             tgt_in = torch.full(
                 (bs, num_steps),
-                tokenizer.pad_id,
+                self.tokenizer.pad_id,
                 dtype=torch.long,
                 device=self._device,
             )
-            tgt_in[:, 0] = tokenizer.bos_id
+            tgt_in[:, 0] = self.tokenizer.bos_id
 
             logits = []
             for i in range(num_steps):
@@ -177,7 +175,7 @@ class PARSeq(nn.Module, PyTorchModelHubMixin):
                     # greedy decode. add the next token index to the target input
                     tgt_in[:, j] = p_i.squeeze().argmax(-1)
                     # Efficient batch decoding: If all output words have at least one EOS token, end decoding.
-                    if testing and (tgt_in == tokenizer.eos_id).any(dim=-1).all():
+                    if testing and (tgt_in == self.tokenizer.eos_id).any(dim=-1).all():
                         break
 
             logits = torch.cat(logits, dim=1)
@@ -185,7 +183,7 @@ class PARSeq(nn.Module, PyTorchModelHubMixin):
             # No prior context, so input is just <bos>. We query all positions.
             tgt_in = torch.full(
                 (bs, 1),
-                tokenizer.bos_id,
+                self.tokenizer.bos_id,
                 dtype=torch.long,
                 device=self._device,
             )
@@ -200,7 +198,7 @@ class PARSeq(nn.Module, PyTorchModelHubMixin):
                     torch.ones(
                         num_steps,
                         num_steps,
-                        dtype=torch.bool,
+                        dtype=torch.int64,
                         device=self._device,
                     ),
                     2,
@@ -208,7 +206,7 @@ class PARSeq(nn.Module, PyTorchModelHubMixin):
             ] = 0
             bos = torch.full(
                 (bs, 1),
-                tokenizer.bos_id,
+                self.tokenizer.bos_id,
                 dtype=torch.long,
                 device=self._device,
             )
@@ -216,7 +214,9 @@ class PARSeq(nn.Module, PyTorchModelHubMixin):
                 # Prior context is the previous output.
                 tgt_in = torch.cat([bos, logits[:, :-1].argmax(-1)], dim=1)
                 # Mask tokens beyond the first EOS token.
-                tgt_padding_mask = (tgt_in == tokenizer.eos_id).int().cumsum(-1) > 0
+                tgt_padding_mask = (tgt_in == self.tokenizer.eos_id).int().cumsum(
+                    -1
+                ) > 0
                 tgt_out = self.decode(
                     tgt_in,
                     memory,

yomitoku/postprocessor/rtdetr_postprocessor.py

@@ -54,16 +54,12 @@ class RTDETRPostProcessor(nn.Module):
         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 = 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
-            )
+            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)
@@ -77,9 +73,7 @@ class RTDETRPostProcessor(nn.Module):
             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
-                )
+                scores, index = torch.topk(scores, self.num_top_queries, dim=-1)
                 labels = torch.gather(labels, dim=1, index=index)
                 boxes = torch.gather(
                     boxes,
@@ -97,10 +91,7 @@ class RTDETRPostProcessor(nn.Module):
 
             labels = (
                 torch.tensor(
-                    [
-                        mscoco_label2category[int(x.item())]
-                        for x in labels.flatten()
-                    ]
+                    [mscoco_label2category[int(x.item())] for x in labels.flatten()]
                 )
                 .to(boxes.device)
                 .reshape(labels.shape)

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
@@ -109,12 +114,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
@@ -127,6 +133,8 @@ class TableStructureRecognizer(BaseModule):
             num_top_queries=self._cfg.RTDETRTransformerv2.num_queries,
         )
 
+        self.save_config("table_structure_recognitizer.yaml")
+
         self.transforms = T.Compose(
             [
                 T.Resize(self._cfg.data.img_size),
@@ -140,6 +148,40 @@ 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)
+
+            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())
+
+    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 +191,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,
@@ -226,8 +268,19 @@ class TableStructureRecognizer(BaseModule):
         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)
 

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,12 +48,13 @@ 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
@@ -58,6 +64,39 @@ class TextDetector(BaseModule):
         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)
+
+            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())
+
+    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 +120,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}

yomitoku/text_recognizer.py

@@ -2,22 +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):
@@ -43,23 +49,41 @@ 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)
+
+            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())
+
     def preprocess(self, img, polygons):
         dataset = ParseqDataset(self._cfg, img, polygons)
         dataloader = torch.utils.data.DataLoader(
@@ -71,6 +95,25 @@ class TextRecognizer(BaseModule):
 
         return dataloader
 
+    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"},
+        }
+
+        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]
@@ -101,13 +144,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-0.5.3.dist-info → yomitoku-0.6.0.dist-info}/METADATA

@@ -1,14 +1,17 @@
 Metadata-Version: 2.3
 Name: yomitoku
-Version: 0.5.3
+Version: 0.6.0
 Summary: Yomitoku is an AI-powered document image analysis package designed specifically for the Japanese language.
 Author-email: Kotaro Kinoshita <kotaro.kinoshita@mlism.com>
 License: CC BY-NC-SA 4.0
 Keywords: Deep Learning,Japanese,OCR
-Requires-Python: <3.13,>=3.9
+Requires-Python: <3.13,>=3.10
 Requires-Dist: huggingface-hub>=0.26.1
 Requires-Dist: lxml>=5.3.0
 Requires-Dist: omegaconf>=2.3.0
+Requires-Dist: onnx>=1.17.0
+Requires-Dist: onnxruntime-gpu>=1.20.1
+Requires-Dist: onnxruntime>=1.20.1
 Requires-Dist: opencv-python>=4.10.0.84
 Requires-Dist: pyclipper>=1.3.0.post6
 Requires-Dist: pydantic>=2.9.2
@@ -23,7 +26,7 @@ Description-Content-Type: text/markdown
 
 <img src="static/logo/horizontal.png" width="800px">
 
-![Python](https://img.shields.io/badge/Python-3.9|3.10|3.11|3.12-F9DC3E.svg?logo=python&logoColor=&style=flat)
+![Python](https://img.shields.io/badge/Python-3.10|3.11|3.12-F9DC3E.svg?logo=python&logoColor=&style=flat)
 ![Pytorch](https://img.shields.io/badge/Pytorch-2.5-EE4C2C.svg?logo=Pytorch&style=fla)
 ![CUDA](https://img.shields.io/badge/CUDA->=11.8-76B900.svg?logo=NVIDIA&style=fla)
 ![OS](https://img.shields.io/badge/OS-Linux|Mac|Win-1793D1.svg?&style=fla)
@@ -69,19 +72,20 @@ Markdown でエクスポートした結果は関してはリポジトリ内の[s
 pip install yomitoku
 ```
 
-- pytorch はご自身の CUDAのバージョンにあったものをインストールしてください。デフォルトではCUDA12.4以上に対応したものがインストールされます。
-- pytorch は2.5以上のバージョンに対応しています。その関係でCUDA11.8以上のバージョンが必要になります。対応できない場合は、リポジトリ内のDockerfileを利用してください。
+- pytorch はご自身の CUDA のバージョンにあったものをインストールしてください。デフォルトでは CUDA12.4 以上に対応したものがインストールされます。
+- pytorch は 2.5 以上のバージョンに対応しています。その関係で CUDA11.8 以上のバージョンが必要になります。対応できない場合は、リポジトリ内の Dockerfile を利用してください。
 
 ## 🚀 実行方法
 
 ```
-yomitoku ${path_data} -f md -o results -v --figure
+yomitoku ${path_data} -f md -o results -v --figure　--lite
 ```
 
 - `${path_data}` 解析対象の画像が含まれたディレクトリか画像ファイルのパスを直接して指定してください。ディレクトリを対象とした場合はディレクトリのサブディレクトリ内の画像も含めて処理を実行します。
 - `-f`, `--format` 出力形式のファイルフォーマットを指定します。(json, csv, html, md をサポート)
 - `-o`, `--outdir` 出力先のディレクトリ名を指定します。存在しない場合は新規で作成されます。
 - `-v`, `--vis` を指定すると解析結果を可視化した画像を出力します。
+- `-l`, `--lite` を指定すると軽量モデルで推論を実行します。通常より高速に推論できますが、若干、精度が低下する可能性があります。
 - `-d`, `--device` モデルを実行するためのデバイスを指定します。gpu が利用できない場合は cpu で推論が実行されます。(デフォルト: cuda)
 - `--ignore_line_break` 画像の改行位置を無視して、段落内の文章を連結して返します。（デフォルト：画像通りの改行位置位置で改行します。）
 - `--figure_letter` 検出した図表に含まれる文字も出力ファイルにエクスポートします。
@@ -94,6 +98,7 @@ yomitoku --help
 ```
 
 **NOTE**
+
 - GPU での実行を推奨します。CPU を用いての推論向けに最適化されておらず、処理時間が長くなります。
 - 活字のみ識別をサポートしております。手書き文字に関しては、読み取れる場合もありますが、公式にはサポートしておりません。
 - Yomitoku は文書 OCR 向けに最適化されており、情景 OCR(看板など紙以外にプリントされた文字の読み取り)向けには最適化されていません。
@@ -107,6 +112,6 @@ yomitoku --help
 
 本リポジトリ内に格納されているソースコードおよび本プロジェクトに関連する HuggingFaceHub 上のモデルの重みファイルのライセンスは CC BY-NC-SA 4.0 に従います。
 非商用での個人利用、研究目的での利用はご自由にお使いください。
-商用目的での利用に関しては、別途、商用ライセンスを提供しますので、開発者にお問い合わせください。
+商用目的での利用に関しては、別途、商用ライセンスを提供しますので、https://www.mlism.com/ にお問い合わせください。
 
 YomiToku © 2024 by Kotaro Kinoshita is licensed under CC BY-NC-SA 4.0. To view a copy of this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/

{yomitoku-0.5.3.dist-info → yomitoku-0.6.0.dist-info}/RECORD

@@ -3,19 +3,20 @@ yomitoku/base.py,sha256=lzR_V8t87aRasmFdFwD-8KAeSahSTI3AZaEn6g8sOv8,3871
 yomitoku/constants.py,sha256=zlW5QRc_u_F3C2RAgBFWyHJZexBnJT5N15GC-9d3iLo,686
 yomitoku/document_analyzer.py,sha256=HIg-nVzDhJIP-h-tn4uU86KakgHdlAhosEqK_i-SWe4,9906
 yomitoku/layout_analyzer.py,sha256=QTeRcVd8aySz8u6dg2ikET77ar3sqlukRLBwYfTyMPM,2033
-yomitoku/layout_parser.py,sha256=V2jCNHE61jNp8ytYdKwPV34V5qEK7y-7-Mq7-AkoQhU,5898
+yomitoku/layout_parser.py,sha256=Yni1C_7j4fzHcdmBNNGRZPc23W_6J6HwPPQVjYvaztM,7539
 yomitoku/ocr.py,sha256=Rcojw0aGA6yDF2RjqfK23_rMw-xm61KGd8JmTCTOOVU,2516
 yomitoku/reading_order.py,sha256=OfhOS9ttPDoPSuHrIRKyOzG19GGeRufbuSKDqhsohh4,6404
-yomitoku/table_structure_recognizer.py,sha256=CouRzfdO_toZKUQbzQqocKdMcgA3Pr7glkZuqD5itpg,7280
-yomitoku/text_detector.py,sha256=okp0xuq4lXgEDcfgCzeJcrj8hfSI4NvAgorsNwi_NYI,2682
-yomitoku/text_recognizer.py,sha256=RHdq1M3-e3C1RECgbaoqPngtxicG3izAma12juD2ICQ,3789
+yomitoku/table_structure_recognizer.py,sha256=Wf_Ehmf6V27iVLmw2o9i7kJnbwEOhuExI-ljIO3a8NE,9043
+yomitoku/text_detector.py,sha256=fbwKelsVfwCt5YL4h-WEf4qkniv5cXmyaLR6oSYz0eA,4167
+yomitoku/text_recognizer.py,sha256=Iu-IzwaziNjmrTeSw9aoN9BDTHkNOzsZhViCv45yiN8,5422
 yomitoku/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-yomitoku/cli/main.py,sha256=MBD0S4sXgquJ8P2egkZjJcglXvCke5Uw46C28SDtr8g,6252
-yomitoku/configs/__init__.py,sha256=KBhb9S7xt22HZaIcoWSgZHfscXXj9YlimOwLH5z9CRo,454
+yomitoku/cli/main.py,sha256=qDB_YNK7abstIr9tYLiJjNU3xLSCd5x1UNDKqwUi2Rk,6885
+yomitoku/configs/__init__.py,sha256=e1Alss5QJLZSNfD6zLEG6xu5vDQDw-4Jayiqq8bq52s,571
 yomitoku/configs/cfg_layout_parser_rtdtrv2.py,sha256=8PRxB2Ar9UF7-DLtbgSokhrzdXb0veWI6Wc-X8qigRw,2329
 yomitoku/configs/cfg_table_structure_recognizer_rtdtrv2.py,sha256=o70GMHD8k-zeBeJtuhPS8x7vVB-ffucnJXeSyn-0AXo,2116
 yomitoku/configs/cfg_text_detector_dbnet.py,sha256=U9k48PON7haoOaytiELhbZRpv9RMiUm6nnfHmdxIa9Q,1153
 yomitoku/configs/cfg_text_recognizer_parseq.py,sha256=hpFs3nKqh4XdU3BZMTultegtLEGahEsCaZdjfKC_MO8,1247
+yomitoku/configs/cfg_text_recognizer_parseq_small.py,sha256=uCm_VC_G79IbZpOiK8fgYzAJ4b98H5pf328wyQomtfo,1259
 yomitoku/data/__init__.py,sha256=KAofFc9rk9ZdTKBjemu9RM8Vj9XnKbWC2MPZ2RWtOdE,82
 yomitoku/data/dataset.py,sha256=-I4f-FDtgsPnJ2MnXB7FtwihMW3koDaSI1OEoqKneIg,1014
 yomitoku/data/functions.py,sha256=eOyxo8S6EoAf1xGSPLWQFb9-t5Rg52NggD9MFIrOSpY,7506
@@ -26,19 +27,20 @@ yomitoku/export/export_json.py,sha256=1ChvCAHfCmMQvCfcAb1p3fSpr4elNAs3xBSIbpfn3b
 yomitoku/export/export_markdown.py,sha256=mCcsXUWBLrYc1NcRSBFfBT28d6eCddAF1oHp0qdBEnE,3986
 yomitoku/models/__init__.py,sha256=Enxq9sjJWusZuxecTori8IQa8NEYKaiiptDluHX1avg,144
 yomitoku/models/dbnet_plus.py,sha256=jeWJZm0ihbxoJeAXBFK7uVIwoosx2IUNk7Ut5wRH0vA,7998
-yomitoku/models/parseq.py,sha256=7QT-q5_oWqXTDXobRk1R6Lpap_AxdC4AzkSsOgXjOwM,8611
+yomitoku/models/parseq.py,sha256=-DQMQuON2jwtb4Ib2V0O19un9w-WG4rXS0SiscydrXU,8593
 yomitoku/models/rtdetr.py,sha256=oJsr8RHz3frslhLfXdVJve47lUsrmqLjfdTrZ41tlQ0,687
 yomitoku/models/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 yomitoku/models/layers/activate.py,sha256=S54GPssZBMloM2oFAXeDVMmBBZOWyjwU98Niq758txE,1244
 yomitoku/models/layers/dbnet_feature_attention.py,sha256=Vpp_PiLVuI7Zs30TTg4RNRn16KTb81ewonADpUHd4aE,6060
 yomitoku/models/layers/parseq_transformer.py,sha256=33eroJf8rmgIptP-NpZLJMhG7XOTwV4rXsq674VrKnU,6704
-yomitoku/models/layers/rtdetr_backbone.py,sha256=wn1jOI8oB0V4GWKaCB-WCNMJ7CqmjvkeO8v2GB3gtAQ,10054
-yomitoku/models/layers/rtdetr_hybrid_encoder.py,sha256=fCXNw8DmD5rXAQZkiVy2Ohj3v19TFUTUCohGlj7V408,14321
-yomitoku/models/layers/rtdetrv2_decoder.py,sha256=iAxZ-TknFuEcI1B6-UU8o0rvWnuBr20idqTWCpc-u7A,28456
+yomitoku/models/layers/rtdetr_backbone.py,sha256=VOWFW7XFfJl4cvPaupqqP4-I-YHdwlVltQEgliD69As,9904
+yomitoku/models/layers/rtdetr_hybrid_encoder.py,sha256=ZnpEzJLzHgu_hrx7YK6myXZ4F1CDHRM501RbAPQdzdQ,14125
+yomitoku/models/layers/rtdetrv2_decoder.py,sha256=ggUwTdWpBfyYHnZuLx8vyH8n0XfZkQFtxgpY-1YI2sI,28070
+yomitoku/onnx/.gitkeep,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 yomitoku/postprocessor/__init__.py,sha256=W4vUuqBaFtH5dlSBIYgyaCroGLMjpV6RrNGIBQ8NFVw,243
 yomitoku/postprocessor/dbnet_postporcessor.py,sha256=o_y8b5REd2dFEdIpRcr6o-XBfOCHo9rBYGwokP_uhTc,4948
 yomitoku/postprocessor/parseq_tokenizer.py,sha256=e89_g_bc4Au3SchuxoJfJNATJTxFmVYetzXyAzPWm28,4315
-yomitoku/postprocessor/rtdetr_postprocessor.py,sha256=ADT620nCs_OvHmoWwH01ylmbHuCNtVUN8pVYYq-vy0Q,4065
+yomitoku/postprocessor/rtdetr_postprocessor.py,sha256=TCv1t1zCxg2rSirsLm4sXlaltGubH-roVdEqnUoRs-8,3905
 yomitoku/resource/MPLUS1p-Medium.ttf,sha256=KLL1KkCumIBkgQtx1n4SffdaFuCNffThktEAbkB1OU8,1758908
 yomitoku/resource/charset.txt,sha256=sU91kSi-9Wk4733bCXy4j_UDmvcsj96sHOq1ppUJlOY,21672
 yomitoku/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -46,7 +48,7 @@ yomitoku/utils/graph.py,sha256=LKNB8ZhSQwOZMfeAimPMF5UCVVr2ZaUWoGDkz8z-uGU,456
 yomitoku/utils/logger.py,sha256=uOmtQDr0A0JD7wyFshedL08BiNrQorHnpktRXba8bjU,424
 yomitoku/utils/misc.py,sha256=2Eyy7-9K_h4Mal1zGXq6OlxubfNzhS0mEYwn_xt7xl8,2497
 yomitoku/utils/visualizer.py,sha256=2pSmbhUPylzVVJ0bXtGDoNmMdArAByab4Py7Xavvs_A,5230
-yomitoku-0.5.3.dist-info/METADATA,sha256=qLwgVjKd3AELsZu8k1JGbX2-VHgHq3Tn-eCaw11c0_s,7819
-yomitoku-0.5.3.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
-yomitoku-0.5.3.dist-info/entry_points.txt,sha256=nFV3S11zgBNW0Qq_D0XQNg2R4lNXU_9XUFr6rdJoyF8,52
-yomitoku-0.5.3.dist-info/RECORD,,
+yomitoku-0.6.0.dist-info/METADATA,sha256=XDmMBtDx9MjXPuzcARwOwJXRN7PMCsQDwc38jDSwX5g,8134
+yomitoku-0.6.0.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
+yomitoku-0.6.0.dist-info/entry_points.txt,sha256=nFV3S11zgBNW0Qq_D0XQNg2R4lNXU_9XUFr6rdJoyF8,52
+yomitoku-0.6.0.dist-info/RECORD,,