yomitoku 0.4.0.post1.dev0__py3-none-any.whl

yomitoku/text_detector.py

@@ -0,0 +1,103 @@
+from typing import List
+
+import numpy as np
+import torch
+from pydantic import conlist
+
+from .base import BaseModelCatalog, BaseModule, BaseSchema
+from .configs import TextDetectorDBNetConfig
+from .data.functions import (
+    array_to_tensor,
+    resize_shortest_edge,
+    standardization_image,
+)
+from .models import DBNet
+from .postprocessor import DBnetPostProcessor
+from .utils.visualizer import det_visualizer
+
+
+class TextDetectorModelCatalog(BaseModelCatalog):
+    def __init__(self):
+        super().__init__()
+        self.register("dbnet", TextDetectorDBNetConfig, DBNet)
+
+
+class TextDetectorSchema(BaseSchema):
+    points: List[
+        conlist(
+            conlist(int, min_length=2, max_length=2),
+            min_length=4,
+            max_length=4,
+        )
+    ]
+    scores: List[float]
+
+
+class TextDetector(BaseModule):
+    model_catalog = TextDetectorModelCatalog()
+
+    def __init__(
+        self,
+        model_name="dbnet",
+        path_cfg=None,
+        device="cuda",
+        visualize=False,
+        from_pretrained=True,
+    ):
+        super().__init__()
+        self.load_model(
+            model_name,
+            path_cfg,
+            from_pretrained=True,
+        )
+
+        self.device = device
+        self.visualize = visualize
+
+        self.model.eval()
+        self.model.to(self.device)
+
+        self.post_processor = DBnetPostProcessor(**self._cfg.post_process)
+
+    def preprocess(self, img):
+        img = img.copy()
+        img = img[:, :, ::-1].astype(np.float32)
+        resized = resize_shortest_edge(
+            img, self._cfg.data.shortest_size, self._cfg.data.limit_size
+        )
+        normalized = standardization_image(resized)
+        tensor = array_to_tensor(normalized)
+        return tensor
+
+    def postprocess(self, preds, image_size):
+        return self.post_processor(preds, image_size)
+
+    def __call__(self, img):
+        """apply the detection model to the input image.
+
+        Args:
+            img (np.ndarray): target image(BGR)
+        """
+
+        ori_h, ori_w = img.shape[:2]
+        tensor = self.preprocess(img)
+        tensor = tensor.to(self.device)
+        with torch.inference_mode():
+            preds = self.model(tensor)
+
+        quads, scores = self.postprocess(preds, (ori_h, ori_w))
+        outputs = {"points": quads, "scores": scores}
+
+        results = TextDetectorSchema(**outputs)
+
+        vis = None
+        if self.visualize:
+            vis = det_visualizer(
+                preds,
+                img,
+                quads,
+                vis_heatmap=self._cfg.visualize.heatmap,
+                line_color=tuple(self._cfg.visualize.color[::-1]),
+            )
+
+        return results, vis

yomitoku/text_recognizer.py

@@ -0,0 +1,128 @@
+from typing import List
+
+import numpy as np
+import torch
+from pydantic import conlist
+
+from .base import BaseModelCatalog, BaseModule, BaseSchema
+from .configs import TextRecognizerPARSeqConfig
+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
+
+
+class TextRecognizerModelCatalog(BaseModelCatalog):
+    def __init__(self):
+        super().__init__()
+        self.register("parseq", TextRecognizerPARSeqConfig, PARSeq)
+
+
+class TextRecognizerSchema(BaseSchema):
+    contents: List[str]
+    directions: List[str]
+    scores: List[float]
+    points: List[
+        conlist(
+            conlist(int, min_length=2, max_length=2),
+            min_length=4,
+            max_length=4,
+        )
+    ]
+
+
+class TextRecognizer(BaseModule):
+    model_catalog = TextRecognizerModelCatalog()
+
+    def __init__(
+        self,
+        model_name="parseq",
+        path_cfg=None,
+        device="cuda",
+        visualize=False,
+        from_pretrained=True,
+    ):
+        super().__init__()
+        self.load_model(
+            model_name,
+            path_cfg,
+            from_pretrained=True,
+        )
+        self.charset = load_charset(self._cfg.charset)
+        self.tokenizer = Tokenizer(self.charset)
+
+        self.device = device
+
+        self.model.eval()
+        self.model.to(self.device)
+
+        self.visualize = visualize
+
+    def preprocess(self, img, polygons):
+        dataset = ParseqDataset(self._cfg, img, polygons)
+        dataloader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=self._cfg.data.batch_size,
+            shuffle=False,
+            num_workers=self._cfg.data.num_workers,
+        )
+
+        return dataloader
+
+    def postprocess(self, p, points):
+        pred, score = self.tokenizer.decode(p)
+        directions = []
+        for point in points:
+            point = np.array(point)
+            w = np.linalg.norm(point[0] - point[1])
+            h = np.linalg.norm(point[1] - point[2])
+
+            direction = "vertical" if h > w * 2 else "horizontal"
+            directions.append(direction)
+
+        return pred, score, directions
+
+    def __call__(self, img, points, vis=None):
+        """
+        Apply the recognition model to the input image.
+
+        Args:
+            img (np.ndarray): target image(BGR)
+            points (list): list of quadrilaterals. Each quadrilateral is represented as a list of 4 points sorted clockwise.
+            vis (np.ndarray, optional): rendering image. Defaults to None.
+        """
+
+        dataloader = self.preprocess(img, points)
+        preds = []
+        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)
+
+        outputs = {
+            "contents": preds,
+            "scores": scores,
+            "points": points,
+            "directions": directions,
+        }
+        results = TextRecognizerSchema(**outputs)
+
+        if self.visualize:
+            if vis is None:
+                vis = img.copy()
+            vis = rec_visualizer(
+                vis,
+                results,
+                font_size=self._cfg.visualize.font_size,
+                font_color=tuple(self._cfg.visualize.color[::-1]),
+                font_path=self._cfg.visualize.font,
+            )
+
+        return results, vis

yomitoku/utils/graph.py

@@ -0,0 +1,20 @@
+class Node:
+    def __init__(self, id, prop):
+        self.id = id
+        self.prop = prop
+        self.parents = []
+        self.children = []
+
+        self.is_locked = False
+
+    def add_link(self, node):
+        if node in self.children:
+            return
+
+        self.children.append(node)
+        node.parents.append(self)
+
+    def __repr__(self):
+        if "contents" in self.prop:
+            return self.prop["contents"]
+        return "table"

yomitoku/utils/logger.py

@@ -0,0 +1,15 @@
+import warnings
+from logging import Formatter, StreamHandler, getLogger
+
+
+def set_logger(name, level="INFO"):
+    logger = getLogger(name)
+    logger.setLevel(level)
+    handler = StreamHandler()
+    handler.setLevel(level)
+    format = Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    handler.setFormatter(format)
+    logger.addHandler(handler)
+
+    warnings.filterwarnings("ignore")
+    return logger

yomitoku/utils/misc.py

@@ -0,0 +1,102 @@
+def load_charset(charset_path):
+    with open(charset_path, "r") as f:
+        charset = f.read()
+    return charset
+
+
+def filter_by_flag(elements, flags):
+    assert len(elements) == len(flags)
+    return [element for element, flag in zip(elements, flags) if flag]
+
+
+def is_contained(rect_a, rect_b, threshold=0.8):
+    """二つの矩形A, Bが与えられたとき、矩形Bが矩形Aに含まれるかどうかを判定する。
+    ずれを許容するため、重複率求め、thresholdを超える場合にTrueを返す。
+
+
+    Args:
+        rect_a (np.array): x1, y1, x2, y2
+        rect_b (np.array): x1, y1, x2, y2
+        threshold (float, optional): 判定の閾値. Defaults to 0.9.
+
+    Returns:
+        bool: 矩形Bが矩形Aに含まれる場合True
+    """
+
+    intersection = calc_intersection(rect_a, rect_b)
+    if intersection is None:
+        return False
+
+    ix1, iy1, ix2, iy2 = intersection
+
+    overlap_width = ix2 - ix1
+    overlap_height = iy2 - iy1
+    bx1, by1, bx2, by2 = rect_b
+
+    b_area = (bx2 - bx1) * (by2 - by1)
+    overlap_area = overlap_width * overlap_height
+
+    if overlap_area / b_area > threshold:
+        return True
+
+    return False
+
+
+def calc_intersection(rect_a, rect_b):
+    ax1, ay1, ax2, ay2 = map(int, rect_a)
+    bx1, by1, bx2, by2 = map(int, rect_b)
+
+    # 交差領域の左上と右下の座標
+    ix1 = max(ax1, bx1)
+    iy1 = max(ay1, by1)
+    ix2 = min(ax2, bx2)
+    iy2 = min(ay2, by2)
+
+    overlap_width = max(0, ix2 - ix1)
+    overlap_height = max(0, iy2 - iy1)
+
+    if overlap_width == 0 or overlap_height == 0:
+        return None
+
+    return [ix1, iy1, ix2, iy2]
+
+
+def is_intersected_horizontal(rect_a, rect_b):
+    _, ay1, _, ay2 = map(int, rect_a)
+    _, by1, _, by2 = map(int, rect_b)
+
+    # 交差領域の左上と右下の座標
+    iy1 = max(ay1, by1)
+    iy2 = min(ay2, by2)
+
+    overlap_height = max(0, iy2 - iy1)
+
+    if overlap_height == 0:
+        return False
+
+    return True
+
+
+def is_intersected_vertical(rect_a, rect_b):
+    ax1, _, ax2, _ = map(int, rect_a)
+    bx1, _, bx2, _ = map(int, rect_b)
+
+    # 交差領域の左上と右下の座標
+    ix1 = max(ax1, bx1)
+    ix2 = min(ax2, bx2)
+
+    overlap_width = max(0, ix2 - ix1)
+
+    if overlap_width == 0:
+        return False
+
+    return True
+
+
+def quad_to_xyxy(quad):
+    x1 = min([x for x, _ in quad])
+    y1 = min([y for _, y in quad])
+    x2 = max([x for x, _ in quad])
+    y2 = max([y for _, y in quad])
+
+    return x1, y1, x2, y2

yomitoku/utils/visualizer.py

@@ -0,0 +1,179 @@
+import cv2
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+
+from ..constants import PALETTE
+
+
+def _reading_order_visualizer(img, elements, line_color, tip_size):
+    out = img.copy()
+    for i, element in enumerate(elements):
+        if i == 0:
+            continue
+
+        prev_element = elements[i - 1]
+        cur_x1, cur_y1, cur_x2, cur_y2 = element.box
+        prev_x1, prev_y1, prev_x2, prev_y2 = prev_element.box
+
+        cur_center = (
+            cur_x1 + (cur_x2 - cur_x1) / 2,
+            cur_y1 + (cur_y2 - cur_y1) / 2,
+        )
+        prev_center = (
+            prev_x1 + (prev_x2 - prev_x1) / 2,
+            prev_y1 + (prev_y2 - prev_y1) / 2,
+        )
+
+        arrow_length = np.linalg.norm(np.array(cur_center) - np.array(prev_center))
+
+        # tipLength を計算（矢印長さに対する固定サイズの割合）
+        if arrow_length > 0:
+            tip_length = tip_size / arrow_length
+        else:
+            tip_length = 0  # 長さが0なら矢じりもゼロ
+
+        cv2.arrowedLine(
+            out,
+            (int(prev_center[0]), int(prev_center[1])),
+            (int(cur_center[0]), int(cur_center[1])),
+            line_color,
+            2,
+            tipLength=tip_length,
+        )
+    return out
+
+
+def reading_order_visualizer(
+    img,
+    results,
+    line_color=(0, 0, 255),
+    tip_size=10,
+    visualize_figure_letter=False,
+):
+    elements = results.paragraphs + results.tables + results.figures
+    elements = sorted(elements, key=lambda x: x.order)
+
+    out = _reading_order_visualizer(img, elements, line_color, tip_size)
+
+    if visualize_figure_letter:
+        for figure in results.figures:
+            out = _reading_order_visualizer(
+                out, figure.paragraphs, line_color=(0, 255, 0), tip_size=5
+            )
+
+    return out
+
+
+def det_visualizer(preds, img, quads, vis_heatmap=False, line_color=(0, 255, 0)):
+    preds = preds["binary"][0]
+    binary = preds.detach().cpu().numpy()
+    out = img.copy()
+    h, w = out.shape[:2]
+    binary = binary.squeeze(0)
+    binary = (binary * 255).astype(np.uint8)
+    if vis_heatmap:
+        binary = cv2.resize(binary, (w, h), interpolation=cv2.INTER_LINEAR)
+        heatmap = cv2.applyColorMap(binary, cv2.COLORMAP_JET)
+        out = cv2.addWeighted(out, 0.5, heatmap, 0.5, 0)
+
+    for quad in quads:
+        quad = np.array(quad).astype(np.int32)
+        out = cv2.polylines(out, [quad], True, line_color, 1)
+    return out
+
+
+def layout_visualizer(results, img):
+    out = img.copy()
+    results_dict = results.dict()
+    for id, (category, preds) in enumerate(results_dict.items()):
+        for element in preds:
+            box = element["box"]
+            role = element["role"]
+
+            if role is None:
+                role = ""
+            else:
+                role = f"({role})"
+
+            color = PALETTE[id % len(PALETTE)]
+            x1, y1, x2, y2 = tuple(map(int, box))
+            out = cv2.rectangle(out, (x1, y1), (x2, y2), color, 2)
+            out = cv2.putText(
+                out,
+                category + role,
+                (x1, y1),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.5,
+                color,
+                2,
+            )
+
+    return out
+
+
+def table_visualizer(img, table):
+    out = img.copy()
+    cells = table.cells
+    for cell in cells:
+        box = cell.box
+        row = cell.row
+        col = cell.col
+        row_span = cell.row_span
+        col_span = cell.col_span
+
+        text = f"[{row}, {col}] ({row_span}x{col_span})"
+
+        x1, y1, x2, y2 = map(int, box)
+        out = cv2.rectangle(out, (x1, y1), (x2, y2), (255, 0, 255), 2)
+        out = cv2.putText(
+            out,
+            text,
+            (x1, y1),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.5,
+            (255, 0, 0),
+            2,
+        )
+
+    return out
+
+
+def rec_visualizer(
+    img,
+    outputs,
+    font_path,
+    font_size=12,
+    font_color=(255, 0, 0),
+):
+    out = img.copy()
+    pillow_img = Image.fromarray(out)
+    draw = ImageDraw.Draw(pillow_img)
+
+    for pred, quad, direction in zip(
+        outputs.contents, outputs.points, outputs.directions
+    ):
+        quad = np.array(quad).astype(np.int32)
+        font = ImageFont.truetype(font_path, font_size)
+        if direction == "horizontal":
+            x_offset = 0
+            y_offset = -font_size
+
+            pos_x = quad[0][0] + x_offset
+            pox_y = quad[0][1] + y_offset
+            draw.text((pos_x, pox_y), pred, font=font, fill=font_color)
+        else:
+            x_offset = -font_size
+            y_offset = 0
+
+            pos_x = quad[0][0] + x_offset
+            pox_y = quad[0][1] + y_offset
+            draw.text(
+                (pos_x, pox_y),
+                pred,
+                font=font,
+                fill=font_color,
+                direction="ttb",
+            )
+
+    out = np.array(pillow_img)
+    return out

yomitoku-0.4.0.post1.dev0.dist-info/METADATA

@@ -0,0 +1,127 @@
+Metadata-Version: 2.3
+Name: yomitoku
+Version: 0.4.0.post1.dev0
+Summary: Yomitoku is a document image analysis package powered by AI technology 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.9
+Requires-Dist: huggingface-hub>=0.26.1
+Requires-Dist: lxml>=5.3.0
+Requires-Dist: omegaconf>=2.3.0
+Requires-Dist: opencv-python>=4.10.0.84
+Requires-Dist: pdf2image>=1.17.0
+Requires-Dist: pyclipper>=1.3.0.post6
+Requires-Dist: pydantic>=2.9.2
+Requires-Dist: shapely>=2.0.6
+Requires-Dist: timm>=1.0.11
+Requires-Dist: torch>=2.5.0
+Requires-Dist: torchvision>=0.20.0
+Description-Content-Type: text/markdown
+
+# YomiToku
+
+![Python](https://img.shields.io/badge/Python-3.9|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)
+![OS](https://img.shields.io/badge/OS-Linux|MacOS-1793D1.svg?&style=fla)
+[![Document](https://img.shields.io/badge/docs-live-brightgreen)](https://kotaro-kinoshita.github.io/yomitoku-dev/)
+
+<img src="static/logo/horizontal.png" width="800px">
+
+## 🌟 概要
+
+YomiToku は日本語に特化した AI 文章画像解析エンジン(Document AI)です。画像内の文字の全文 OCR およびレイアウト解析機能を有しており、画像内の文字情報や図表を認識、抽出、変換します。
+
+- 🤖 日本語データセットで学習した 4 種類(文字位置の検知、文字列認識、レイアウト解析、表の構造認識)の AI モデルを搭載しています。4 種類のモデルはすべて独自に学習されたモデルで日本語文書に対して、高精度に推論可能です。
+- 🇯🇵 各モデルは日本語の文書画像に特化して学習されており、7000 文字を超える日本語文字の認識をサーポート、縦書きなど日本語特有のレイアウト構造の文書画像の解析も可能です。（日本語以外にも英語の文書に対しても対応しています）。
+- 📈 レイアウト解析、表の構造解析, 読み順推定機能により、文書画像のレイアウトの意味的構造を壊さずに情報を抽出することが可能です。
+- 📄 多様な出力形式をサポートしています。html やマークダウン、json、csv のいずれかのフォーマットに変換可能です。また、文書内に含まれる図表、画像の抽出の出力も可能です。
+- ⚡ GPU 環境で高速に動作し、効率的に文書の文字起こし解析が可能です。また、VRAM も 8GB 以内で動作し、ハイエンドな GPU を用意する必要はありません。
+
+## 🖼️ デモ
+
+[gallery.md](gallery.md)にも複数種類の画像の検証結果を掲載しています。
+
+|                          入力画像                          |                       OCR の結果                        |
+| :--------------------------------------------------------: | :-----------------------------------------------------: |
+|        <img src="static/in/demo.jpg" width="400px">        | <img src="static/out/in_demo_p1_ocr.jpg" width="400px"> |
+|                    レイアウト解析の結果                    |     エクスポート<br>(HTML で出力したものをスクショ)     |
+| <img src="static/out/in_demo_p1_layout.jpg" width="400px"> |   <img src="static/out/demo_html.png" width="400px">    |
+
+Markdown でエクスポートした結果は関してはリポジトリ内の[static/out/in_demo_p1.md](static/out/in_demo_p1.md)を参照
+
+- `赤枠` : 図、画像等の位置
+- `緑枠` : 表領域全体の位置
+- `ピンク枠` : 表のセル構造(セル上の文字は [行番号, 列番号] (rowspan x colspan)を表します)
+- `青枠` : 段落、テキストグループ領域
+- `赤矢印` : 読み順推定の結果
+
+画像の出典:[「令和 6 年版情報通信白書 3 章 2 節 AI の進化に伴い発展するテクノロジー」](https://www.soumu.go.jp/johotsusintokei/whitepaper/ja/r06/pdf/n1410000.pdf)：（総務省） を加工して作成
+
+## 📣 リリース情報
+
+- 2024 年 12 月 XX YomiToku vX.X.X を公開
+
+## 💡 インストールの方法
+
+```
+pip install git+https://github.com/kotaro-kinoshita/yomitoku-dev.git@main
+```
+
+- pytorch がご自身の GPU の環境にあったものをインストールしてください
+
+### 依存ライブラリ
+
+pdf ファイルの解析を行うためには、別途、[poppler](https://poppler.freedesktop.org/)のインストールが必要です。
+
+**Mac**
+
+```
+brew install poppler
+```
+
+**Linux**
+
+```
+apt install poppler-utils -y
+```
+
+## 🚀 実行方法
+
+```
+yomitoku ${path_data} -f md -o results -v --figure
+```
+
+- `${path_data}` 解析対象の画像が含まれたディレクトリか画像ファイルのパスを直接して指定してください。ディレクトリを対象とした場合はディレクトリのサブディレクトリ内の画像も含めて処理を実行します。
+- `-f`, `--format` 出力形式のファイルフォーマットを指定します。(json, csv, html, md をサポート)
+- `-o`, `--outdir` 出力先のディレクトリ名を指定します。存在しない場合は新規で作成されます。
+- `-v`, `--vis` を指定すると解析結果を可視化した画像を出力します。
+- `-d`, `--device` モデルを実行するためのデバイスを指定します。gpu が利用できない場合は cpu で推論が実行されます。(デフォルト: cuda)
+- `--ignore_line_break` 画像の改行位置を無視して、段落内の文章を連結して返します。（デフォルト：画像通りの改行位置位置で改行します。）
+- `figure_letter` 検出した図表に含まれる文字も出力ファイルにエクスポートします。
+- `figure` 検出した図、画像を出力ファイルにエクスポートします。(html と markdown のみ)
+
+その他のオプションに関しては、ヘルプを参照
+
+```
+yomitoku --help
+```
+
+### Note
+
+- CPU を用いての推論向けに最適化されておらず、処理時間が長くなりますので、GPU での実行を推奨します。
+- 活字のみ識別をサポートしております。手書き文字に関しては、読み取れる場合もありますが、公式にはサポートしておりません。
+- OCR は文書 OCR と情景 OCR(看板など紙以外にプリントされた文字)に大別されますが、Yomitoku は文書 OCR 向けに最適化されています。
+- AI-OCR の識別精度を高めるために、入力画像の解像度が重要です。低解像度画像では識別精度が低下します。最低でも画像の短辺を 720px 以上の画像で推論することをお勧めします。
+
+## 📝 ドキュメント
+
+パッケージの詳細は[ドキュメント](https://kotaro-kinoshita.github.io/yomitoku-dev/)を確認してください。
+
+## LICENSE
+
+本リポジトリ内に格納されているリソースのライセンスは YomiToku は CC BY-NC-SA 4.0 に従います。
+非商用での個人利用、研究目的での利用は自由に利用できます。
+商用目的での利用に関しては、別途、商用ライセンスを提供しますので、開発者にお問い合わせください。
+
+YomiToku © 2024 by MLism Inc. 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/