yomitoku 0.5.3__tar.gz → 0.6.0__tar.gz

{yomitoku-0.5.3 → yomitoku-0.6.0}/.github/workflows/lint-and-test.yml

@@ -37,4 +37,4 @@ jobs:
       - name: Run linter
         run: tox -e lint
       - name: Run tests
-        run:  tox -p -e py39,py310,py311,py312
+        run:  tox -p -e py310,py311,py312

{yomitoku-0.5.3 → yomitoku-0.6.0}/.gitignore

@@ -20,4 +20,6 @@ dataset/
 weights/
 results/
 
-.coverage*
+.coverage*
+
+*.onnx

yomitoku-0.6.0/.python-version

@@ -0,0 +1 @@
+3.10

{yomitoku-0.5.3 → yomitoku-0.6.0}/PKG-INFO

@@ -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 → yomitoku-0.6.0}/README.md

@@ -2,7 +2,7 @@
 
 <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)
@@ -48,19 +48,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` 検出した図表に含まれる文字も出力ファイルにエクスポートします。
@@ -73,6 +74,7 @@ yomitoku --help
 ```
 
 **NOTE**
+
 - GPU での実行を推奨します。CPU を用いての推論向けに最適化されておらず、処理時間が長くなります。
 - 活字のみ識別をサポートしております。手書き文字に関しては、読み取れる場合もありますが、公式にはサポートしておりません。
 - Yomitoku は文書 OCR 向けに最適化されており、情景 OCR(看板など紙以外にプリントされた文字の読み取り)向けには最適化されていません。
@@ -86,6 +88,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 → yomitoku-0.6.0}/README_EN.md

@@ -2,7 +2,7 @@
 
 <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)
@@ -63,6 +63,7 @@ yomitoku ${path_data} -f md -o results -v --figure
 - `-f`, `--format`: Specify the output file format. Supported formats are json, csv, html, and md.
 - `-o`, `--outdir`: Specify the name of the output directory. If it does not exist, it will be created.
 - `-v`, `--vis`: If specified, outputs visualized images of the analysis results.
+- `-l`, `--lite`: inference is performed using a lightweight model. This enables fast inference even on a CPU.
 - `-d`, `--device`: Specify the device for running the model. If a GPU is unavailable, inference will be executed on the CPU. (Default: cuda)
 - `--ignore_line_break`: Ignores line breaks in the image and concatenates sentences within a paragraph. (Default: respects line breaks as they appear in the image.)
 - `--figure_letter`: Exports characters contained within detected figures and tables to the output file.

yomitoku-0.6.0/configs/table_structure_recognitizer.yaml

@@ -0,0 +1,67 @@
+hf_hub_repo: KotaroKinoshita/yomitoku-table-structure-recognizer-rtdtrv2-open-beta
+thresh_score: 0.4
+data:
+  img_size:
+  - 640
+  - 640
+PResNet:
+  depth: 50
+  variant: d
+  freeze_at: 0
+  return_idx:
+  - 1
+  - 2
+  - 3
+  num_stages: 4
+  freeze_norm: true
+HybridEncoder:
+  in_channels:
+  - 512
+  - 1024
+  - 2048
+  feat_strides:
+  - 8
+  - 16
+  - 32
+  hidden_dim: 256
+  use_encoder_idx:
+  - 2
+  num_encoder_layers: 1
+  nhead: 8
+  dim_feedforward: 1024
+  dropout: 0.0
+  enc_act: gelu
+  expansion: 1.0
+  depth_mult: 1
+  act: silu
+RTDETRTransformerv2:
+  num_classes: 3
+  feat_channels:
+  - 256
+  - 256
+  - 256
+  feat_strides:
+  - 8
+  - 16
+  - 32
+  hidden_dim: 256
+  num_levels: 3
+  num_layers: 6
+  num_queries: 300
+  num_denoising: 100
+  label_noise_ratio: 0.5
+  box_noise_scale: 1.0
+  eval_spatial_size:
+  - 640
+  - 640
+  eval_idx: -1
+  num_points:
+  - 4
+  - 4
+  - 4
+  cross_attn_method: default
+  query_select_method: default
+category:
+- row
+- col
+- span

yomitoku-0.6.0/configs/text_recognizer.yaml

@@ -0,0 +1,32 @@
+hf_hub_repo: KotaroKinoshita/yomitoku-text-recognizer-parseq-open-beta
+charset: /home/kinoshita/Projects/know-how/yomitoku/src/yomitoku/resource/charset.txt
+num_tokens: 7312
+max_label_length: 100
+decode_ar: 1
+refine_iters: 1
+data:
+  num_workers: 4
+  batch_size: 128
+  img_size:
+  - 32
+  - 800
+encoder:
+  patch_size:
+  - 8
+  - 8
+  num_heads: 8
+  embed_dim: 512
+  mlp_ratio: 4
+  depth: 12
+decoder:
+  embed_dim: 512
+  num_heads: 8
+  mlp_ratio: 4
+  depth: 1
+visualize:
+  font: /home/kinoshita/Projects/know-how/yomitoku/src/yomitoku/resource/MPLUS1p-Medium.ttf
+  color:
+  - 0
+  - 0
+  - 255
+  font_size: 18

{yomitoku-0.5.3 → yomitoku-0.6.0}/docs/installation.en.md

@@ -1,7 +1,7 @@
 # Installation
 
 
-This package requires Python 3.9 or later and PyTorch 2.5 or later for execution. PyTorch must be installed according to your CUDA version. A GPU with more than 8GB of VRAM is recommended. While it can run on a CPU, please note that the processing is not currently optimized for CPUs, which may result in longer execution times.
+This package requires Python 3.10 or later and PyTorch 2.5 or later for execution. PyTorch must be installed according to your CUDA version. A GPU with more than 8GB of VRAM is recommended. While it can run on a CPU, please note that the processing is not currently optimized for CPUs, which may result in longer execution times.
 
 ## from PYPI
 

{yomitoku-0.5.3 → yomitoku-0.6.0}/docs/installation.ja.md

@@ -1,6 +1,6 @@
 # Installation
 
-本パッケージは Python3.9+, Pytorch が実行に必要です。Pytorch はご自身の環境に合わせて、インストールが必要です。計算機は GPU(> VRAM 8G)を推奨しています。CPU でも動作しますが、現在、CPU 向けに処理が最適化されておらず、実行に時間がかかりますのでご注意ください。
+本パッケージは Python3.10+, Pytorch が実行に必要です。Pytorch はご自身の環境に合わせて、インストールが必要です。計算機は GPU(> VRAM 8G)を推奨しています。CPU でも動作しますが、現在、CPU 向けに処理が最適化されておらず、実行に時間がかかりますのでご注意ください。
 
 ## PYPI からインストール
 

{yomitoku-0.5.3 → yomitoku-0.6.0}/docs/usage.en.md

@@ -12,6 +12,7 @@ yomitoku ${path_data} -f md -o results -v
 - `-f`, `--format`: Specify the output file format. Supported formats are json, csv, html, and md.
 - `-o`, `--outdir`: Specify the name of the output directory. If it does not exist, it will be created.
 - `-v`, `--vis`: If specified, outputs visualized images of the analysis results.
+- `-l`, `--lite`: inference is performed using a lightweight model. This enables fast inference even on a CPU.
 - `-d`, `--device`: Specify the device for running the model. If a GPU is unavailable, inference will be executed on the CPU. (Default: cuda)
 - `--ignore_line_break`: Ignores line breaks in the image and concatenates sentences within a paragraph. (Default: respects line breaks as they appear in the image.)
 - `--figure_letter`: Exports characters contained within detected figures and tables to the output file.

{yomitoku-0.5.3 → yomitoku-0.6.0}/docs/usage.ja.md

@@ -10,6 +10,7 @@ yomitoku ${path_data} -f md -o results -v
 
 - `${path_data}` 解析対象の画像が含まれたディレクトリか画像ファイルのパスを直接して指定してください。かディレクトリを対象とした場合はディレクトリのサブディレクトリ内の画像も含めて処理を実行します。入力をサポートしているファイル形式は jpeg, png, bmp, tiff, pdf です。
 - `-f` 出力形式のファイルフォーマットを指定します。(json, csv, html, md をサポート)
+- `-l` 指定すると軽量モデルで推論を実行します。CPUでも高速に推論可能です。
 - `-o` 出力先のディレクトリ名を指定します。存在しない場合は新規で作成されます。
 - `-v` を指定すると解析結果を可視化した画像を出力します。
 - `-d` モデルを実行するためのデバイスを指定します。gpu が利用できない場合は cpu で推論が実行されます。(デフォルト: cuda)

{yomitoku-0.5.3 → yomitoku-0.6.0}/pyproject.toml

@@ -12,7 +12,7 @@ authors = [{name = "Kotaro Kinoshita", email = "kotaro.kinoshita@mlism.com"}]
 description = "Yomitoku is an AI-powered document image analysis package designed specifically for the Japanese language."
 readme = "README.md"
 license = {text = "CC BY-NC-SA 4.0"}
-requires-python = ">=3.9,<3.13"
+requires-python = ">=3.10,<3.13"
 keywords = ["Japanese", "OCR", "Deep Learning"]
 dependencies = [
     "huggingface-hub>=0.26.1",
@@ -26,6 +26,9 @@ dependencies = [
     "torchvision>=0.20.0",
     "torch>=2.5.0",
     "pypdfium2>=4.30.0",
+    "onnx>=1.17.0",
+    "onnxruntime>=1.20.1",
+    "onnxruntime-gpu>=1.20.1",
 ]
 
 [tool.uv-dynamic-versioning]
@@ -71,7 +74,7 @@ yomitoku = "yomitoku.cli.main:main"
 [tool.tox]
 legacy_tox_ini = """
 [tox]
-envlist = lint, py39, py310, py311, py312, docs
+envlist = lint, py310, py311, py312, docs
 
 [testenv]
 deps = pytest

{yomitoku-0.5.3 → yomitoku-0.6.0}/scripts/register_hugging_face_hub.py

@@ -14,7 +14,7 @@ def get_module(module_name):
         return module
 
     elif module_name == "text_recognizer":
-        module = TextRecognizer(from_pretrained=False)
+        module = TextRecognizer(from_pretrained=False, model_name="parseq-small")
         return module
 
     elif module_name == "layout_parser":

{yomitoku-0.5.3 → yomitoku-0.6.0}/src/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-0.5.3 → yomitoku-0.6.0}/src/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-0.6.0/src/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-0.5.3 → yomitoku-0.6.0}/src/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-0.5.3 → yomitoku-0.6.0}/src/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-0.5.3 → yomitoku-0.6.0}/src/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)
             )