yomitoku 0.7.1__py3-none-any.whl → 0.7.2__py3-none-any.whl

yomitoku/base.py

@@ -54,7 +54,7 @@ class BaseSchema(BaseModel):
         validate_assignment = True
 
     def to_json(self, out_path: str, **kwargs):
-        export_json(self, out_path, **kwargs)
+        return export_json(self, out_path, **kwargs)
 
 
 class BaseModule:

yomitoku/cli/main.py

@@ -1,19 +1,65 @@
 import argparse
 import os
-import torch
+import time
 from pathlib import Path
 
 import cv2
-import time
+import torch
 
 from ..constants import SUPPORT_OUTPUT_FORMAT
 from ..data.functions import load_image, load_pdf
 from ..document_analyzer import DocumentAnalyzer
 from ..utils.logger import set_logger
 
+from ..export import save_csv, save_html, save_json, save_markdown
+
 logger = set_logger(__name__, "INFO")
 
 
+def merge_all_pages(results):
+    out = None
+    for result in results:
+        format = result["format"]
+        data = result["data"]
+
+        if format == "json":
+            if out is None:
+                out = [data]
+            else:
+                out.append(data)
+
+        elif format == "csv":
+            if out is None:
+                out = data
+            else:
+                out.extend(data)
+
+        elif format == "html":
+            if out is None:
+                out = data
+            else:
+                out += "\n" + data
+
+        elif format == "md":
+            if out is None:
+                out = data
+            else:
+                out += "\n" + data
+
+    return out
+
+
+def save_merged_file(out_path, args, out):
+    if args.format == "json":
+        save_json(out_path, args.encoding, out)
+    elif args.format == "csv":
+        save_csv(out_path, args.encoding, out)
+    elif args.format == "html":
+        save_html(out_path, args.encoding, out)
+    elif args.format == "md":
+        save_markdown(out_path, args.encoding, out)
+
+
 def validate_encoding(encoding):
     if encoding not in [
         "utf-8",
@@ -32,8 +78,9 @@ def process_single_file(args, analyzer, path, format):
     else:
         imgs = [load_image(path)]
 
+    results = []
     for page, img in enumerate(imgs):
-        results, ocr, layout = analyzer(img)
+        result, ocr, layout = analyzer(img)
         dirname = path.parent.name
         filename = path.stem
 
@@ -56,7 +103,7 @@ def process_single_file(args, analyzer, path, format):
         out_path = os.path.join(args.outdir, f"{dirname}_{filename}_p{page+1}.{format}")
 
         if format == "json":
-            results.to_json(
+            json = result.to_json(
                 out_path,
                 ignore_line_break=args.ignore_line_break,
                 encoding=args.encoding,
@@ -64,8 +111,19 @@ def process_single_file(args, analyzer, path, format):
                 export_figure=args.figure,
                 figure_dir=args.figure_dir,
             )
+
+            results.append(
+                {
+                    "format": format,
+                    "data": json,
+                }
+            )
+
+            if not args.combine:
+                save_json(out_path, args.encoding, json)
+
         elif format == "csv":
-            results.to_csv(
+            csv = result.to_csv(
                 out_path,
                 ignore_line_break=args.ignore_line_break,
                 encoding=args.encoding,
@@ -73,8 +131,19 @@ def process_single_file(args, analyzer, path, format):
                 export_figure=args.figure,
                 figure_dir=args.figure_dir,
             )
+
+            results.append(
+                {
+                    "format": format,
+                    "data": csv,
+                }
+            )
+
+            if not args.combine:
+                save_csv(out_path, args.encoding, csv)
+
         elif format == "html":
-            results.to_html(
+            html = result.to_html(
                 out_path,
                 ignore_line_break=args.ignore_line_break,
                 img=img,
@@ -84,8 +153,19 @@ def process_single_file(args, analyzer, path, format):
                 figure_dir=args.figure_dir,
                 encoding=args.encoding,
             )
+
+            results.append(
+                {
+                    "format": format,
+                    "data": html,
+                }
+            )
+
+            if not args.combine:
+                save_html(out_path, args.encoding, html)
+
         elif format == "md":
-            results.to_markdown(
+            md = result.to_markdown(
                 out_path,
                 ignore_line_break=args.ignore_line_break,
                 img=img,
@@ -96,7 +176,24 @@ def process_single_file(args, analyzer, path, format):
                 encoding=args.encoding,
             )
 
-        logger.info(f"Output file: {out_path}")
+            results.append(
+                {
+                    "format": format,
+                    "data": md,
+                }
+            )
+
+            if not args.combine:
+                save_markdown(out_path, args.encoding, md)
+
+    out = merge_all_pages(results)
+    if args.combine:
+        out_path = os.path.join(args.outdir, f"{dirname}_{filename}.{format}")
+        save_merged_file(
+            out_path,
+            args,
+            out,
+        )
 
 
 def main():
@@ -196,6 +293,16 @@ def main():
         default="utf-8",
         help="Specifies the character encoding for the output file to be exported. If unsupported characters are included, they will be ignored.",
     )
+    parser.add_argument(
+        "--combine",
+        action="store_true",
+        help="if set, merge all pages in the output",
+    )
+    parser.add_argument(
+        "--ignore_meta",
+        action="store_true",
+        help="if set, ignore meta information(header, footer) in the output",
+    )
 
     args = parser.parse_args()
 
@@ -248,6 +355,7 @@ def main():
         configs=configs,
         visualize=args.vis,
         device=args.device,
+        ignore_meta=args.ignore_meta,
     )
 
     os.makedirs(args.outdir, exist_ok=True)

yomitoku/document_analyzer.py

@@ -3,25 +3,19 @@ from concurrent.futures import ThreadPoolExecutor
 from typing import List, Union
 
 import numpy as np
-
 from pydantic import conlist
 
+from yomitoku.text_detector import TextDetector
+from yomitoku.text_recognizer import TextRecognizer
+
 from .base import BaseSchema
 from .export import export_csv, export_html, export_markdown
 from .layout_analyzer import LayoutAnalyzer
 from .ocr import OCRSchema, WordPrediction, ocr_aggregate
 from .reading_order import prediction_reading_order
 from .table_structure_recognizer import TableStructureRecognizerSchema
-from .utils.misc import (
-    is_contained,
-    quad_to_xyxy,
-    calc_overlap_ratio,
-)
-from .utils.visualizer import reading_order_visualizer
-from yomitoku.text_detector import TextDetector
-from yomitoku.text_recognizer import TextRecognizer
-
-from .utils.visualizer import det_visualizer
+from .utils.misc import calc_overlap_ratio, is_contained, quad_to_xyxy
+from .utils.visualizer import det_visualizer, reading_order_visualizer
 
 
 class ParagraphSchema(BaseSchema):
@@ -47,13 +41,13 @@ class DocumentAnalyzerSchema(BaseSchema):
     figures: List[FigureSchema]
 
     def to_html(self, out_path: str, **kwargs):
-        export_html(self, out_path, **kwargs)
+        return export_html(self, out_path, **kwargs)
 
     def to_markdown(self, out_path: str, **kwargs):
-        export_markdown(self, out_path, **kwargs)
+        return export_markdown(self, out_path, **kwargs)
 
     def to_csv(self, out_path: str, **kwargs):
-        export_csv(self, out_path, **kwargs)
+        return export_csv(self, out_path, **kwargs)
 
 
 def combine_flags(flag1, flag2):
@@ -328,7 +322,13 @@ def _split_text_across_cells(results_det, results_layout):
 
 
 class DocumentAnalyzer:
-    def __init__(self, configs={}, device="cuda", visualize=False):
+    def __init__(
+        self,
+        configs={},
+        device="cuda",
+        visualize=False,
+        ignore_meta=False,
+    ):
         default_configs = {
             "ocr": {
                 "text_detector": {
@@ -371,6 +371,8 @@ class DocumentAnalyzer:
         )
         self.visualize = visualize
 
+        self.ignore_meta = ignore_meta
+
     def aggregate(self, ocr_res, layout_res):
         paragraphs = []
         check_list = [False] * len(ocr_res.words)
@@ -431,11 +433,15 @@ class DocumentAnalyzer:
         page_direction = judge_page_direction(paragraphs)
 
         headers = [
-            paragraph for paragraph in paragraphs if paragraph.role == "page_header"
+            paragraph
+            for paragraph in paragraphs
+            if paragraph.role == "page_header" and not self.ignore_meta
         ]
 
         footers = [
-            paragraph for paragraph in paragraphs if paragraph.role == "page_footer"
+            paragraph
+            for paragraph in paragraphs
+            if paragraph.role == "page_footer" and not self.ignore_meta
         ]
 
         page_contents = [
@@ -503,9 +509,9 @@ class DocumentAnalyzer:
 
     def __call__(self, img):
         self.img = img
-        resutls, ocr, layout = asyncio.run(self.run(img))
+        results, ocr, layout = asyncio.run(self.run(img))
 
         if self.visualize:
-            layout = reading_order_visualizer(layout, resutls)
+            layout = reading_order_visualizer(layout, results)
 
-        return resutls, ocr, layout
+        return results, ocr, layout

yomitoku/export/__init__.py

@@ -1,6 +1,15 @@
-from .export_csv import export_csv
-from .export_html import export_html
-from .export_json import export_json
-from .export_markdown import export_markdown
+from .export_csv import export_csv, save_csv
+from .export_html import export_html, save_html
+from .export_json import export_json, save_json
+from .export_markdown import export_markdown, save_markdown
 
-__all__ = ["export_html", "export_markdown", "export_csv", "export_json"]
+__all__ = [
+    "export_html",
+    "export_markdown",
+    "export_csv",
+    "export_json",
+    "save_html",
+    "save_markdown",
+    "save_csv",
+    "save_json",
+]

yomitoku/export/export_csv.py

@@ -1,7 +1,8 @@
 import csv
-import cv2
 import os
 
+import cv2
+
 
 def table_to_csv(table, ignore_line_break):
     num_rows = table.n_row
@@ -98,7 +99,10 @@ def export_csv(
         )
 
     elements = sorted(elements, key=lambda x: x["order"])
+    return elements
+
 
+def save_csv(out_path, encoding, elements):
     with open(out_path, "w", newline="", encoding=encoding, errors="ignore") as f:
         writer = csv.writer(f, quoting=csv.QUOTE_MINIMAL)
         for element in elements:

yomitoku/export/export_html.py

@@ -1,9 +1,8 @@
-import re
 import os
-import cv2
-
+import re
 from html import escape
 
+import cv2
 from lxml import etree, html
 
 
@@ -182,10 +181,13 @@ def export_html(
     elements = sorted(elements, key=lambda x: x["order"])
 
     html_string = "".join([element["html"] for element in elements])
-    html_string = add_html_tag(html_string)
+    # html_string = add_html_tag(html_string)
 
     parsed_html = html.fromstring(html_string)
     formatted_html = etree.tostring(parsed_html, pretty_print=True, encoding="unicode")
+    return formatted_html
+
 
+def save_html(out_path, encoding, html):
     with open(out_path, "w", encoding=encoding, errors="ignore") as f:
-        f.write(formatted_html)
+        f.write(html)

yomitoku/export/export_json.py

@@ -1,7 +1,7 @@
 import json
+import os
 
 import cv2
-import os
 
 
 def paragraph_to_json(paragraph, ignore_line_break):
@@ -63,9 +63,13 @@ def export_json(
                 figure_dir=figure_dir,
             )
 
+    return inputs.model_dump()
+
+
+def save_json(out_path, encoding, data):
     with open(out_path, "w", encoding=encoding, errors="ignore") as f:
         json.dump(
-            inputs.model_dump(),
+            data,
             f,
             ensure_ascii=False,
             indent=4,

yomitoku/export/export_markdown.py

@@ -1,6 +1,7 @@
+import os
 import re
+
 import cv2
-import os
 
 
 def escape_markdown_special_chars(text):
@@ -144,5 +145,9 @@ def export_markdown(
     elements = sorted(elements, key=lambda x: x["order"])
     markdown = "\n".join([element["md"] for element in elements])
 
+    return markdown
+
+
+def save_markdown(out_path, encoding, markdown):
     with open(out_path, "w", encoding=encoding, errors="ignore") as f:
         f.write(markdown)

{yomitoku-0.7.1.dist-info → yomitoku-0.7.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: yomitoku
-Version: 0.7.1
+Version: 0.7.2
 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
@@ -19,8 +19,6 @@ Requires-Dist: shapely>=2.0.6
 Requires-Dist: timm>=1.0.11
 Requires-Dist: torch>=2.5.0
 Requires-Dist: torchvision>=0.20.0
-Provides-Extra: gpu
-Requires-Dist: onnxruntime-gpu>=1.20.1; extra == 'gpu'
 Description-Content-Type: text/markdown
 
 日本語版 | [English](README_EN.md)
@@ -32,6 +30,7 @@ Description-Content-Type: text/markdown
 ![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)
 [![Document](https://img.shields.io/badge/docs-live-brightgreen)](https://kotaro-kinoshita.github.io/yomitoku/)
+[![PyPI Downloads](https://static.pepy.tech/badge/yomitoku)](https://pepy.tech/projects/yomitoku)
 
 ## 🌟 概要
 
@@ -73,18 +72,13 @@ Markdown でエクスポートした結果は関してはリポジトリ内の[s
 pip install yomitoku
 ```
 
-onnxruntimeの実行にGPUを使用する場合
-```
-pip install yomitoku[gpu]
-```
-
 - pytorch はご自身の CUDA のバージョンにあったものをインストールしてください。デフォルトでは CUDA12.4 以上に対応したものがインストールされます。
 - pytorch は 2.5 以上のバージョンに対応しています。その関係で CUDA11.8 以上のバージョンが必要になります。対応できない場合は、リポジトリ内の Dockerfile を利用してください。
 
 ## 🚀 実行方法
 
 ```
-yomitoku ${path_data} -f md -o results -v --figure　--lite
+yomitoku ${path_data} -f md -o results -v --figure --lite
 ```
 
 - `${path_data}` 解析対象の画像が含まれたディレクトリか画像ファイルのパスを直接して指定してください。ディレクトリを対象とした場合はディレクトリのサブディレクトリ内の画像も含めて処理を実行します。
@@ -97,6 +91,8 @@ yomitoku ${path_data} -f md -o results -v --figure　--lite
 - `--figure_letter` 検出した図表に含まれる文字も出力ファイルにエクスポートします。
 - `--figure` 検出した図、画像を出力ファイルにエクスポートします。
 - `--encoding` エクスポートする出力ファイルの文字エンコーディングを指定します。サポートされていない文字コードが含まれる場合は、その文字を無視します。(utf-8, utf-8-sig, shift-jis, enc-jp, cp932)
+- `--combine` PDFを入力に与えたときに、複数ページが含まれる場合に、それらの予測結果を一つのファイルに統合してエクスポートします。
+- `--ignore_meta` 文章のheater, fotterなどの文字情報を出力ファイルに含めません。
 
 その他のオプションに関しては、ヘルプを参照
 

{yomitoku-0.7.1.dist-info → yomitoku-0.7.2.dist-info}/RECORD

@@ -1,7 +1,7 @@
 yomitoku/__init__.py,sha256=kXOM8RbpwwLABG3p3vPT3dJWBk4JX2MFGrOeBEW0hKM,543
-yomitoku/base.py,sha256=lzR_V8t87aRasmFdFwD-8KAeSahSTI3AZaEn6g8sOv8,3871
+yomitoku/base.py,sha256=9U3sfe69O6vuO430JzzKQQNkgPsLM9WdLfOUUhp3Ljs,3878
 yomitoku/constants.py,sha256=zlW5QRc_u_F3C2RAgBFWyHJZexBnJT5N15GC-9d3iLo,686
-yomitoku/document_analyzer.py,sha256=B2F_MXFKbq58ePDCgcZKk_bgQUkno1ehYb6CZmAekCk,16234
+yomitoku/document_analyzer.py,sha256=wQMmXACDsDmyaxg2OnG9Og5Nx53WPUkQdUmgYtljACQ,16412
 yomitoku/layout_analyzer.py,sha256=VhNf1ZQFoozj6WUGk5ll1p2p1jk5X3j-JPcDbTAoSl4,1856
 yomitoku/layout_parser.py,sha256=V_mAkZxke1gwHfnxBFMTOJ8hnz2X_kfZu2lLiMd8cAs,7610
 yomitoku/ocr.py,sha256=JSTjkupcxHITQm6ERnzU7As0c3KWf8-oxc0AqNoWHXo,2272
@@ -10,7 +10,7 @@ yomitoku/table_structure_recognizer.py,sha256=Eam9t7OjW4a-UWk_dl-ylbOcinN_Te_ovu
 yomitoku/text_detector.py,sha256=XgqhtbNcJww2x3BrH8EFz45qC6kqPKCX9hsa-dzRoIA,4274
 yomitoku/text_recognizer.py,sha256=LVMjy-PaGlDQqfJrjKX_7vOQXDyFg6FaCeIQIyWUJX8,5833
 yomitoku/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-yomitoku/cli/main.py,sha256=Dcu-jBql5dP4DiN0QWbPAQscnJ7tTtbWUnaohz3FFbs,7868
+yomitoku/cli/main.py,sha256=WvQO9V5HzxxvRAIsGkrDl9OGrmaKsAbDBrg4ApCSy_c,10527
 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
@@ -20,11 +20,11 @@ yomitoku/configs/cfg_text_recognizer_parseq_small.py,sha256=uCm_VC_G79IbZpOiK8fg
 yomitoku/data/__init__.py,sha256=KAofFc9rk9ZdTKBjemu9RM8Vj9XnKbWC2MPZ2RWtOdE,82
 yomitoku/data/dataset.py,sha256=-I4f-FDtgsPnJ2MnXB7FtwihMW3koDaSI1OEoqKneIg,1014
 yomitoku/data/functions.py,sha256=eOyxo8S6EoAf1xGSPLWQFb9-t5Rg52NggD9MFIrOSpY,7506
-yomitoku/export/__init__.py,sha256=aANEfuovH2aevFjb2pGrBLFP-4iRzEzD9wcriCR-M7I,229
-yomitoku/export/export_csv.py,sha256=gKIhhFKOyZaxwmjbtxk8rkFOv3uPaRwfOnKwTC8d7K0,2935
-yomitoku/export/export_html.py,sha256=RsFU-IMtBOJWCdqHhp5btswf2fzfow01ypY1h6E6Vvo,4979
-yomitoku/export/export_json.py,sha256=-57hcT2ENTa1HcT7YMjXiv6tZZf_Y_1q1xu2Jt9T1P4,1976
-yomitoku/export/export_markdown.py,sha256=kGOyopq_vT2NqM4LoAu4JIPbjV_SLg9H0MIOAKpS-Gk,4099
+yomitoku/export/__init__.py,sha256=fkwOtqH0lh6eZQW5b4EMSjIH1FmWYLKKszahR-jQYSg,366
+yomitoku/export/export_csv.py,sha256=B234jlNeO4n5kQ_lwxxAZe_O2ipTbeDYlWU1zyyaVrw,3001
+yomitoku/export/export_html.py,sha256=pCLoxV10_SzRWmZlDnHuyfPFIuUGB3ZkqSdABVU7DTs,5038
+yomitoku/export/export_json.py,sha256=D6dD04gcPR5lmfHFVX-iGOYapsOVaJ_kH1Qhs6d2O0M,2035
+yomitoku/export/export_markdown.py,sha256=4mk_7P4kta6RguThnQUh7hesvZNLuXFnPlxhA951qVE,4171
 yomitoku/models/__init__.py,sha256=Enxq9sjJWusZuxecTori8IQa8NEYKaiiptDluHX1avg,144
 yomitoku/models/dbnet_plus.py,sha256=jeWJZm0ihbxoJeAXBFK7uVIwoosx2IUNk7Ut5wRH0vA,7998
 yomitoku/models/parseq.py,sha256=-DQMQuON2jwtb4Ib2V0O19un9w-WG4rXS0SiscydrXU,8593
@@ -48,7 +48,7 @@ yomitoku/utils/graph.py,sha256=LKNB8ZhSQwOZMfeAimPMF5UCVVr2ZaUWoGDkz8z-uGU,456
 yomitoku/utils/logger.py,sha256=uOmtQDr0A0JD7wyFshedL08BiNrQorHnpktRXba8bjU,424
 yomitoku/utils/misc.py,sha256=FbwPLeIYYBvNf9wQh2RoEonTM5BF7_IwaEqmRsYHKA8,2673
 yomitoku/utils/visualizer.py,sha256=DjDwHiAu1iFRKh96H3Egq4vuI2s_-9dLCDeykhKi8jo,5251
-yomitoku-0.7.1.dist-info/METADATA,sha256=NgweyozOKBkIn9-yxX0hOV313-dBh29SsFQ22BSvoD8,8488
-yomitoku-0.7.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-yomitoku-0.7.1.dist-info/entry_points.txt,sha256=nFV3S11zgBNW0Qq_D0XQNg2R4lNXU_9XUFr6rdJoyF8,52
-yomitoku-0.7.1.dist-info/RECORD,,
+yomitoku-0.7.2.dist-info/METADATA,sha256=4nyP_4zWDOxWgt7qfwRbgtlHL9xiv_FLzdw8YtfJNYU,8717
+yomitoku-0.7.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+yomitoku-0.7.2.dist-info/entry_points.txt,sha256=nFV3S11zgBNW0Qq_D0XQNg2R4lNXU_9XUFr6rdJoyF8,52
+yomitoku-0.7.2.dist-info/RECORD,,