yomitoku 0.8.1__py3-none-any.whl → 0.9.0__py3-none-any.whl

yomitoku/cli/main.py

@@ -3,7 +3,6 @@ import os
 import time
 from pathlib import Path
 
-import cv2
 import torch
 
 from ..constants import SUPPORT_OUTPUT_FORMAT
@@ -14,6 +13,8 @@ from ..utils.logger import set_logger
 from ..export import save_csv, save_html, save_json, save_markdown
 from ..export import convert_json, convert_csv, convert_html, convert_markdown
 
+from ..utils.misc import save_image
+
 logger = set_logger(__name__, "INFO")
 
 
@@ -94,7 +95,7 @@ def process_single_file(args, analyzer, path, format):
                 args.outdir, f"{dirname}_{filename}_p{page+1}_ocr.jpg"
             )
 
-            cv2.imwrite(out_path, ocr)
+            save_image(ocr, out_path)
             logger.info(f"Output file: {out_path}")
 
         if layout is not None:
@@ -102,7 +103,7 @@ def process_single_file(args, analyzer, path, format):
                 args.outdir, f"{dirname}_{filename}_p{page+1}_layout.jpg"
             )
 
-            cv2.imwrite(out_path, layout)
+            save_image(layout, out_path)
             logger.info(f"Output file: {out_path}")
 
         out_path = os.path.join(args.outdir, f"{dirname}_{filename}_p{page+1}.{format}")

yomitoku/cli/mcp.py

@@ -0,0 +1,165 @@
+import json
+import io
+import csv
+import os
+from pathlib import Path
+
+from mcp.server.fastmcp import Context, FastMCP
+
+from yomitoku import DocumentAnalyzer
+from yomitoku.data.functions import load_image, load_pdf
+from yomitoku.export import convert_json, convert_markdown, convert_csv, convert_html
+
+try:
+    RESOURCE_DIR = os.environ["RESOURCE_DIR"]
+except KeyError:
+    raise ValueError("Environment variable 'RESOURCE_DIR' is not set.")
+
+
+analyzer = None
+
+
+async def load_analyzer(ctx: Context) -> DocumentAnalyzer:
+    """
+    Load the DocumentAnalyzer instance if not already loaded.
+
+    Args:
+        ctx (Context): The context in which the analyzer is being loaded.
+
+    Returns:
+        DocumentAnalyzer: The loaded document analyzer instance.
+    """
+    global analyzer
+    if analyzer is None:
+        await ctx.info("Load document analyzer")
+        analyzer = DocumentAnalyzer(visualize=False, device="cuda")
+    return analyzer
+
+
+mcp = FastMCP("yomitoku")
+
+
+@mcp.tool()
+async def process_ocr(ctx: Context, filename: str, output_format: str) -> str:
+    """
+    Perform OCR on the specified file in the resource direcory and convert
+    the results to the desired format.
+
+    Args:
+        ctx (Context): The context in which the OCR processing is executed.
+        filename (str): The name of the file to process in the resource directory.
+        output_format (str): The desired format for the output. The available options are:
+            - json: Outputs the text as structured data along with positional information.
+            - markdown: Outputs texts and tables in Markdown format.
+            - html: Outputs texts and tables in HTML format.
+            - csv: Outputs texts and tables in CSV format.
+
+    Returns:
+        str: The OCR results converted to the specified format.
+    """
+    analyzer = await load_analyzer(ctx)
+
+    await ctx.info("Start ocr processing")
+
+    file_path = os.path.join(RESOURCE_DIR, filename)
+    if Path(file_path).suffix[1:].lower() in ["pdf"]:
+        imgs = load_pdf(file_path)
+    else:
+        imgs = load_image(file_path)
+
+    results = []
+    for page, img in enumerate(imgs):
+        analyzer.img = img
+        result, _, _ = await analyzer.run(img)
+        results.append(result)
+        await ctx.report_progress(page + 1, len(imgs))
+
+    if output_format == "json":
+        return json.dumps(
+            [
+                convert_json(
+                    result,
+                    out_path=None,
+                    ignore_line_break=True,
+                    img=img,
+                    export_figure=False,
+                    figure_dir=None,
+                ).model_dump()
+                for img, result in zip(imgs, results)
+            ],
+            ensure_ascii=False,
+            sort_keys=True,
+            separators=(",", ": "),
+        )
+    elif output_format == "markdown":
+        return "\n".join(
+            [
+                convert_markdown(
+                    result,
+                    out_path=None,
+                    ignore_line_break=True,
+                    img=img,
+                    export_figure=False,
+                )[0]
+                for img, result in zip(imgs, results)
+            ]
+        )
+    elif output_format == "html":
+        return "\n".join(
+            [
+                convert_html(
+                    result,
+                    out_path=None,
+                    ignore_line_break=True,
+                    img=img,
+                    export_figure=False,
+                    export_figure_letter="",
+                )[0]
+                for img, result in zip(imgs, results)
+            ]
+        )
+    elif output_format == "csv":
+        output = io.StringIO()
+        writer = csv.writer(output, quoting=csv.QUOTE_MINIMAL)
+        for img, result in zip(imgs, results):
+            elements = convert_csv(
+                result,
+                out_path=None,
+                ignore_line_break=True,
+                img=img,
+                export_figure=False,
+            )
+            for element in elements:
+                if element["type"] == "table":
+                    writer.writerows(element["element"])
+                else:
+                    writer.writerow([element["element"]])
+                writer.writerow([""])
+        return output.getvalue()
+    else:
+        raise ValueError(
+            f"Unsupported output format: {output_format}."
+            " Supported formats are json, markdown, html or csv."
+        )
+
+
+@mcp.resource("file://list")
+async def get_file_list() -> list[str]:
+    """
+    Retrieve a list of files in the resource directory.
+
+    Returns:
+        list[str]: A list of filenames in the resource directory.
+    """
+    return os.listdir(RESOURCE_DIR)
+
+
+def run_mcp_server():
+    """
+    Run the MCP server.
+    """
+    mcp.run(transport="stdio")
+
+
+if __name__ == "__main__":
+    run_mcp_server()

yomitoku/export/export_csv.py

@@ -1,7 +1,7 @@
 import csv
 import os
 
-import cv2
+from ..utils.misc import save_image
 
 
 def table_to_csv(table, ignore_line_break):
@@ -54,7 +54,7 @@ def save_figure(
         filename = os.path.splitext(os.path.basename(out_path))[0]
         figure_name = f"{filename}_figure_{i}.png"
         figure_path = os.path.join(save_dir, figure_name)
-        cv2.imwrite(figure_path, figure_img)
+        save_image(figure_img, figure_path)
 
 
 def convert_csv(

yomitoku/export/export_html.py

@@ -1,10 +1,10 @@
 import os
 import re
 from html import escape
-
-import cv2
 from lxml import etree, html
 
+from ..utils.misc import save_image
+
 
 def convert_text_to_html(text):
     """
@@ -122,7 +122,7 @@ def figure_to_html(
         filename = os.path.splitext(os.path.basename(out_path))[0]
         figure_name = f"{filename}_figure_{i}.png"
         figure_path = os.path.join(save_dir, figure_name)
-        cv2.imwrite(figure_path, figure_img)
+        save_image(figure_img, figure_path)
 
         elements.append(
             {
@@ -180,8 +180,13 @@ def convert_html(
     elements = sorted(elements, key=lambda x: x["order"])
 
     html_string = "".join([element["html"] for element in elements])
-    parsed_html = html.fromstring(html_string)
-    formatted_html = etree.tostring(parsed_html, pretty_print=True, encoding="unicode")
+    if not len(html_string) == 0:
+        parsed_html = html.fromstring(html_string)
+        formatted_html = etree.tostring(
+            parsed_html, pretty_print=True, encoding="unicode"
+        )
+    else:
+        formatted_html = ""
 
     return formatted_html, elements
 

yomitoku/export/export_json.py

@@ -1,7 +1,7 @@
 import json
 import os
 
-import cv2
+from ..utils.misc import save_image
 
 
 def paragraph_to_json(paragraph, ignore_line_break):
@@ -33,7 +33,7 @@ def save_figure(
         filename = os.path.splitext(os.path.basename(out_path))[0]
         figure_name = f"{filename}_figure_{i}.png"
         figure_path = os.path.join(save_dir, figure_name)
-        cv2.imwrite(figure_path, figure_img)
+        save_image(figure_img, figure_path)
 
 
 def convert_json(inputs, out_path, ignore_line_break, img, export_figure, figure_dir):

yomitoku/export/export_markdown.py

@@ -1,7 +1,7 @@
 import os
 import re
 
-import cv2
+from ..utils.misc import save_image
 
 
 def escape_markdown_special_chars(text):
@@ -89,7 +89,7 @@ def figure_to_md(
         filename = os.path.splitext(os.path.basename(out_path))[0]
         figure_name = f"{filename}_figure_{i}.png"
         figure_path = os.path.join(save_dir, figure_name)
-        cv2.imwrite(figure_path, figure_img)
+        save_image(figure_img, figure_path)
 
         elements.append(
             {

yomitoku/utils/misc.py

@@ -1,3 +1,6 @@
+import cv2
+
+
 def load_charset(charset_path):
     with open(charset_path, "r", encoding="utf-8") as f:
         charset = f.read()
@@ -9,6 +12,15 @@ def filter_by_flag(elements, flags):
     return [element for element, flag in zip(elements, flags) if flag]
 
 
+def save_image(img, path):
+    success, buffer = cv2.imencode(".jpg", img)
+    if not success:
+        raise ValueError("Failed to encode image")
+
+    with open(path, "wb") as f:
+        f.write(buffer.tobytes())
+
+
 def calc_overlap_ratio(rect_a, rect_b):
     intersection = calc_intersection(rect_a, rect_b)
     if intersection is None:

{yomitoku-0.8.1.dist-info → yomitoku-0.9.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: yomitoku
-Version: 0.8.1
+Version: 0.9.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
@@ -19,6 +19,8 @@ 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: mcp
+Requires-Dist: mcp[cli]>=1.6.0; extra == 'mcp'
 Description-Content-Type: text/markdown
 
 日本語版 | [English](README_EN.md)

{yomitoku-0.8.1.dist-info → yomitoku-0.9.0.dist-info}/RECORD

@@ -10,7 +10,8 @@ yomitoku/table_structure_recognizer.py,sha256=tHjex6deT_FjRK5ePz9bUXA_QIhgv_vYtK
 yomitoku/text_detector.py,sha256=6IwEJJKp_F8YH0Oki0QV-Mqi--P2LGbNKo-_kxBB_eo,4383
 yomitoku/text_recognizer.py,sha256=eaxozNu-Ms6iv8efbKZzn8pJNW1Wo4f86bGhzSMtv3s,5992
 yomitoku/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-yomitoku/cli/main.py,sha256=jQCSwHw4oOwLQjARvaIO1yoSjz-2Rdb9c3DNShLS5OE,12038
+yomitoku/cli/main.py,sha256=9X8QWwsTAv82uNn5Ns9T_laGJPKHDyeEwenaQxnAmn4,12062
+yomitoku/cli/mcp.py,sha256=5h704SsUGNAqVnoO_5S-HY2-bApy_Rf8ajDxl1pkT2k,4888
 yomitoku/configs/__init__.py,sha256=x5-ccjGiP6xxRtDPT7f1Enl7SsE0hSk0G8f7eF9V85I,886
 yomitoku/configs/cfg_layout_parser_rtdtrv2.py,sha256=8PRxB2Ar9UF7-DLtbgSokhrzdXb0veWI6Wc-X8qigRw,2329
 yomitoku/configs/cfg_layout_parser_rtdtrv2_v2.py,sha256=nMrL3uvoVmyzZ909Bz2zmfp9b6AEBLKhIprOvQ5yiQE,2324
@@ -24,10 +25,10 @@ yomitoku/data/__init__.py,sha256=KAofFc9rk9ZdTKBjemu9RM8Vj9XnKbWC2MPZ2RWtOdE,82
 yomitoku/data/dataset.py,sha256=-I4f-FDtgsPnJ2MnXB7FtwihMW3koDaSI1OEoqKneIg,1014
 yomitoku/data/functions.py,sha256=HIrffs0zCJOq8IvQiI_z-b4MwTb-H2wmZjEE_5VpxFs,8040
 yomitoku/export/__init__.py,sha256=gmlikMHRXfzfJ_8q4fyDlnpGms-x1oggQOwJEWHMgBU,508
-yomitoku/export/export_csv.py,sha256=4DT5Nf4FdeGP0olIzv1ypBlswkZSdMB4MeQOgYWe8uk,3375
-yomitoku/export/export_html.py,sha256=syzAapHcUHcUlabmZcQdWiNy2NrRs7LPzA_x39pFtfQ,5494
-yomitoku/export/export_json.py,sha256=6cSXSsyEVJ5Rw2nKSUOcW8_XlGmSLWlWQWCBNmRKsps,2386
-yomitoku/export/export_markdown.py,sha256=7Jib-YXOw70H46kvNc6z0_3LFwX9iwp1eXxsGeylF0I,4681
+yomitoku/export/export_csv.py,sha256=VY8mntUCPDbDco_dyvq5O0_Q4wga9_GTyjHCS-y4UiQ,3399
+yomitoku/export/export_html.py,sha256=LQDyZgbzmI0qJ0-FEK-54r9816H3L9hD10ChMcw0KyA,5620
+yomitoku/export/export_json.py,sha256=iNG37tdIuYG2x3NiiZemKaB6-X45WrhVPZhbX7RUzRI,2410
+yomitoku/export/export_markdown.py,sha256=KrdxDmKzVP_LbTKuDNGGsT31QOPKVsNNlb6wtLEW-1Q,4705
 yomitoku/models/__init__.py,sha256=Enxq9sjJWusZuxecTori8IQa8NEYKaiiptDluHX1avg,144
 yomitoku/models/dbnet_plus.py,sha256=jeWJZm0ihbxoJeAXBFK7uVIwoosx2IUNk7Ut5wRH0vA,7998
 yomitoku/models/parseq.py,sha256=psCPjP3eKjOFAUZJPQQhbD0nWEV5FeOZ0tTK27Rvvbw,8748
@@ -49,9 +50,9 @@ yomitoku/resource/charset.txt,sha256=sU91kSi-9Wk4733bCXy4j_UDmvcsj96sHOq1ppUJlOY
 yomitoku/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 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/misc.py,sha256=cIUrvSJwfWwTui7ueYistf9XPapPR3XgqD2wQjWit40,2901
 yomitoku/utils/visualizer.py,sha256=DjDwHiAu1iFRKh96H3Egq4vuI2s_-9dLCDeykhKi8jo,5251
-yomitoku-0.8.1.dist-info/METADATA,sha256=Ds7gHmT1DxOJROrfpGaJGQKZ1qZ8ur78c6oxdld9kG4,8555
-yomitoku-0.8.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-yomitoku-0.8.1.dist-info/entry_points.txt,sha256=nFV3S11zgBNW0Qq_D0XQNg2R4lNXU_9XUFr6rdJoyF8,52
-yomitoku-0.8.1.dist-info/RECORD,,
+yomitoku-0.9.0.dist-info/METADATA,sha256=vUbrNm2w-7OIqEEXNzFQBDm8y57mTuh1UeJYHBGRo9U,8622
+yomitoku-0.9.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+yomitoku-0.9.0.dist-info/entry_points.txt,sha256=N3PzzSo-fdgri5liPpZ3ItMmRH6oVX14pIU_5pUJiAs,99
+yomitoku-0.9.0.dist-info/RECORD,,

{yomitoku-0.8.1.dist-info → yomitoku-0.9.0.dist-info}/entry_points.txt

@@ -1,2 +1,3 @@
 [console_scripts]
 yomitoku = yomitoku.cli.main:main
+yomitoku_mcp = yomitoku.cli.mcp:run_mcp_server