yomitoku 0.4.1__py3-none-any.whl → 0.7.4__py3-none-any.whl

yomitoku/document_analyzer.py

@@ -2,17 +2,20 @@ import asyncio
 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 OCR, WordPrediction
-from .table_structure_recognizer import TableStructureRecognizerSchema
-from .utils.misc import is_contained, quad_to_xyxy
+from .ocr import OCRSchema, WordPrediction, ocr_aggregate
 from .reading_order import prediction_reading_order
-
-from .utils.visualizer import reading_order_visualizer
+from .table_structure_recognizer import TableStructureRecognizerSchema
+from .utils.misc import calc_overlap_ratio, is_contained, quad_to_xyxy
+from .utils.visualizer import det_visualizer, reading_order_visualizer
 
 
 class ParagraphSchema(BaseSchema):
@@ -38,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):
@@ -98,47 +101,56 @@ def extract_words_within_element(pred_words, element):
     word_sum_width = 0
     word_sum_height = 0
     check_list = [False] * len(pred_words)
+
     for i, word in enumerate(pred_words):
         word_box = quad_to_xyxy(word.points)
         if is_contained(element.box, word_box, threshold=0.5):
-            contained_words.append(word)
             word_sum_width += word_box[2] - word_box[0]
             word_sum_height += word_box[3] - word_box[1]
             check_list[i] = True
 
+            word_element = ParagraphSchema(
+                box=word_box,
+                contents=word.content,
+                direction=word.direction,
+                order=0,
+                role=None,
+            )
+            contained_words.append(word_element)
+
     if len(contained_words) == 0:
         return None, None, check_list
 
-    mean_width = word_sum_width / len(contained_words)
-    mean_height = word_sum_height / len(contained_words)
-
     word_direction = [word.direction for word in contained_words]
     cnt_horizontal = word_direction.count("horizontal")
     cnt_vertical = word_direction.count("vertical")
 
     element_direction = "horizontal" if cnt_horizontal > cnt_vertical else "vertical"
-    if element_direction == "horizontal":
-        contained_words = sorted(
-            contained_words,
-            key=lambda x: (
-                x.points[0][1] // int(mean_height),
-                x.points[0][0],
-            ),
-        )
-    else:
-        contained_words = sorted(
-            contained_words,
-            key=lambda x: (
-                x.points[1][0] // int(mean_width),
-                x.points[1][1],
-            ),
-            reverse=True,
-        )
 
-    contained_words = "\n".join([content.content for content in contained_words])
+    prediction_reading_order(contained_words, element_direction)
+    contained_words = sorted(contained_words, key=lambda x: x.order)
+
+    contained_words = "\n".join([content.contents for content in contained_words])
+
     return (contained_words, element_direction, check_list)
 
 
+def is_vertical(quad, thresh_aspect=2):
+    quad = np.array(quad)
+    width = np.linalg.norm(quad[0] - quad[1])
+    height = np.linalg.norm(quad[1] - quad[2])
+
+    return height > width * thresh_aspect
+
+
+def is_noise(quad, thresh=15):
+    quad = np.array(quad)
+    width = np.linalg.norm(quad[0] - quad[1])
+    height = np.linalg.norm(quad[1] - quad[2])
+
+    return width < thresh or height < thresh
+
+
 def recursive_update(original, new_data):
     for key, value in new_data.items():
         # `value`が辞書の場合、再帰的に更新
@@ -154,8 +166,169 @@ def recursive_update(original, new_data):
     return original
 
 
+def _extract_words_within_table(words, table, check_list):
+    horizontal_words = []
+    vertical_words = []
+
+    for i, (points, score) in enumerate(zip(words.points, words.scores)):
+        word_box = quad_to_xyxy(points)
+        if is_contained(table.box, word_box, threshold=0.5):
+            if is_vertical(points):
+                vertical_words.append({"points": points, "score": score})
+            else:
+                horizontal_words.append({"points": points, "score": score})
+
+            check_list[i] = True
+
+    return (horizontal_words, vertical_words, check_list)
+
+
+def _calc_overlap_words_on_lines(lines, words):
+    overlap_ratios = [[0 for _ in lines] for _ in words]
+
+    for i, word in enumerate(words):
+        word_box = quad_to_xyxy(word["points"])
+        for j, row in enumerate(lines):
+            overlap_ratio, _ = calc_overlap_ratio(
+                row.box,
+                word_box,
+            )
+            overlap_ratios[i][j] = overlap_ratio
+
+    return overlap_ratios
+
+
+def _correct_vertical_word_boxes(overlap_ratios_vertical, table, table_words_vertical):
+    allocated_cols = [cols.index(max(cols)) for cols in overlap_ratios_vertical]
+
+    new_points = []
+    new_scores = []
+    for i, col_index in enumerate(allocated_cols):
+        col_cells = []
+        for cell in table.cells:
+            if cell.col <= (col_index + 1) < (cell.col + cell.col_span):
+                col_cells.append(cell)
+
+        word_point = table_words_vertical[i]["points"]
+        word_score = table_words_vertical[i]["score"]
+
+        for cell in col_cells:
+            word_box = quad_to_xyxy(word_point)
+
+            _, intersection = calc_overlap_ratio(
+                cell.box,
+                word_box,
+            )
+
+            if intersection is not None:
+                _, y1, _, y2 = intersection
+
+                new_point = [
+                    [word_point[0][0], max(word_point[0][1], y1)],
+                    [word_point[1][0], max(word_point[1][1], y1)],
+                    [word_point[2][0], min(word_point[2][1], y2)],
+                    [word_point[3][0], min(word_point[3][1], y2)],
+                ]
+
+                if not is_noise(new_point):
+                    new_points.append(new_point)
+                    new_scores.append(word_score)
+
+    return new_points, new_scores
+
+
+def _correct_horizontal_word_boxes(
+    overlap_ratios_horizontal, table, table_words_horizontal
+):
+    allocated_rows = [rows.index(max(rows)) for rows in overlap_ratios_horizontal]
+
+    new_points = []
+    new_scores = []
+    for i, row_index in enumerate(allocated_rows):
+        row_cells = []
+        for cell in table.cells:
+            if cell.row <= (row_index + 1) < (cell.row + cell.row_span):
+                row_cells.append(cell)
+
+        word_point = table_words_horizontal[i]["points"]
+        word_score = table_words_horizontal[i]["score"]
+
+        for cell in row_cells:
+            word_box = quad_to_xyxy(word_point)
+
+            _, intersection = calc_overlap_ratio(
+                cell.box,
+                word_box,
+            )
+
+            if intersection is not None:
+                x1, _, x2, _ = intersection
+
+                new_point = [
+                    [max(word_point[0][0], x1), word_point[0][1]],
+                    [min(word_point[1][0], x2), word_point[1][1]],
+                    [min(word_point[2][0], x2), word_point[2][1]],
+                    [max(word_point[3][0], x1), word_point[3][1]],
+                ]
+
+                if not is_noise(new_point):
+                    new_points.append(new_point)
+                    new_scores.append(word_score)
+
+    return new_points, new_scores
+
+
+def _split_text_across_cells(results_det, results_layout):
+    check_list = [False] * len(results_det.points)
+    new_points = []
+    new_scores = []
+    for table in results_layout.tables:
+        table_words_horizontal, table_words_vertical, check_list = (
+            _extract_words_within_table(results_det, table, check_list)
+        )
+
+        overlap_ratios_horizontal = _calc_overlap_words_on_lines(
+            table.rows,
+            table_words_horizontal,
+        )
+
+        overlap_ratios_vertical = _calc_overlap_words_on_lines(
+            table.cols,
+            table_words_vertical,
+        )
+
+        new_points_horizontal, new_scores_horizontal = _correct_horizontal_word_boxes(
+            overlap_ratios_horizontal, table, table_words_horizontal
+        )
+
+        new_points_vertical, new_scores_vertical = _correct_vertical_word_boxes(
+            overlap_ratios_vertical, table, table_words_vertical
+        )
+
+        new_points.extend(new_points_horizontal)
+        new_scores.extend(new_scores_horizontal)
+        new_points.extend(new_points_vertical)
+        new_scores.extend(new_scores_vertical)
+
+    for i, flag in enumerate(check_list):
+        if not flag:
+            new_points.append(results_det.points[i])
+            new_scores.append(results_det.scores[i])
+
+    results_det.points = new_points
+    results_det.scores = new_scores
+
+    return results_det
+
+
 class DocumentAnalyzer:
-    def __init__(self, configs=None, device="cuda", visualize=False):
+    def __init__(
+        self,
+        configs={},
+        device="cuda",
+        visualize=False,
+        ignore_meta=False,
+    ):
         default_configs = {
             "ocr": {
                 "text_detector": {
@@ -186,10 +359,20 @@ class DocumentAnalyzer:
                 "configs must be a dict. See the https://kotaro-kinoshita.github.io/yomitoku-dev/usage/"
             )
 
-        self.ocr = OCR(configs=default_configs["ocr"])
-        self.layout = LayoutAnalyzer(configs=default_configs["layout_analyzer"])
+        self.text_detector = TextDetector(
+            **default_configs["ocr"]["text_detector"],
+        )
+        self.text_recognizer = TextRecognizer(
+            **default_configs["ocr"]["text_recognizer"]
+        )
+
+        self.layout = LayoutAnalyzer(
+            configs=default_configs["layout_analyzer"],
+        )
         self.visualize = visualize
 
+        self.ignore_meta = ignore_meta
+
     def aggregate(self, ocr_res, layout_res):
         paragraphs = []
         check_list = [False] * len(ocr_res.words)
@@ -250,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 = [
@@ -292,24 +479,39 @@ class DocumentAnalyzer:
         with ThreadPoolExecutor(max_workers=2) as executor:
             loop = asyncio.get_running_loop()
             tasks = [
-                loop.run_in_executor(executor, self.ocr, img),
+                # loop.run_in_executor(executor, self.ocr, img),
+                loop.run_in_executor(executor, self.text_detector, img),
                 loop.run_in_executor(executor, self.layout, img),
             ]
 
             results = await asyncio.gather(*tasks)
 
-            results_ocr, ocr = results[0]
+            results_det, _ = results[0]
             results_layout, layout = results[1]
 
-        outputs = self.aggregate(results_ocr, results_layout)
+            results_det = _split_text_across_cells(results_det, results_layout)
+
+            vis_det = None
+            if self.visualize:
+                vis_det = det_visualizer(
+                    img,
+                    results_det.points,
+                )
+
+            results_rec, ocr = self.text_recognizer(img, results_det.points, vis_det)
+
+            outputs = {"words": ocr_aggregate(results_det, results_rec)}
+            results_ocr = OCRSchema(**outputs)
+            outputs = self.aggregate(results_ocr, results_layout)
+
         results = DocumentAnalyzerSchema(**outputs)
         return results, ocr, layout
 
     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,19 @@
-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, convert_csv
+from .export_html import export_html, save_html, convert_html
+from .export_json import export_json, save_json, convert_json
+from .export_markdown import export_markdown, save_markdown, convert_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",
+    "convert_html",
+    "convert_markdown",
+    "convert_csv",
+    "convert_json",
+]

yomitoku/export/export_csv.py

@@ -1,4 +1,7 @@
 import csv
+import os
+
+import cv2
 
 
 def table_to_csv(table, ignore_line_break):
@@ -33,7 +36,35 @@ def paragraph_to_csv(paragraph, ignore_line_break):
     return contents
 
 
-def export_csv(inputs, out_path: str, ignore_line_break: bool = False):
+def save_figure(
+    figures,
+    img,
+    out_path,
+    figure_dir="figures",
+):
+    assert img is not None, "img is required for saving figures"
+
+    for i, figure in enumerate(figures):
+        x1, y1, x2, y2 = map(int, figure.box)
+        figure_img = img[y1:y2, x1:x2, :]
+        save_dir = os.path.dirname(out_path)
+        save_dir = os.path.join(save_dir, figure_dir)
+        os.makedirs(save_dir, exist_ok=True)
+
+        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)
+
+
+def convert_csv(
+    inputs,
+    out_path,
+    ignore_line_break,
+    img=None,
+    export_figure: bool = True,
+    figure_dir="figures",
+):
     elements = []
     for table in inputs.tables:
         table_csv = table_to_csv(table, ignore_line_break)
@@ -60,7 +91,45 @@ def export_csv(inputs, out_path: str, ignore_line_break: bool = False):
 
     elements = sorted(elements, key=lambda x: x["order"])
 
-    with open(out_path, "w", newline="", encoding="utf-8") as f:
+    if export_figure:
+        save_figure(
+            inputs.figures,
+            img,
+            out_path,
+            figure_dir=figure_dir,
+        )
+
+    return elements
+
+
+def export_csv(
+    inputs,
+    out_path: str,
+    ignore_line_break: bool = False,
+    encoding: str = "utf-8",
+    img=None,
+    export_figure: bool = True,
+    figure_dir="figures",
+):
+    elements = convert_csv(
+        inputs,
+        out_path,
+        ignore_line_break,
+        img,
+        export_figure,
+        figure_dir,
+    )
+
+    save_csv(elements, out_path, encoding)
+    return elements
+
+
+def save_csv(
+    elements,
+    out_path,
+    encoding,
+):
+    with open(out_path, "w", newline="", encoding=encoding, errors="ignore") as f:
         writer = csv.writer(f, quoting=csv.QUOTE_MINIMAL)
         for element in elements:
             if element["type"] == "table":

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
 
 
@@ -110,6 +109,8 @@ def figure_to_html(
     figure_dir="figures",
     width=200,
 ):
+    assert img is not None, "img is required for saving figures"
+
     elements = []
     for i, figure in enumerate(figures):
         x1, y1, x2, y2 = map(int, figure.box)
@@ -145,12 +146,12 @@ def figure_to_html(
     return elements
 
 
-def export_html(
+def convert_html(
     inputs,
-    out_path: str,
-    ignore_line_break: bool = False,
-    export_figure: bool = True,
-    export_figure_letter: bool = False,
+    out_path,
+    ignore_line_break,
+    export_figure,
+    export_figure_letter,
     img=None,
     figure_width=200,
     figure_dir="figures",
@@ -179,10 +180,43 @@ 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)
-
     parsed_html = html.fromstring(html_string)
     formatted_html = etree.tostring(parsed_html, pretty_print=True, encoding="unicode")
 
-    with open(out_path, "w", encoding="utf-8") as f:
-        f.write(formatted_html)
+    return formatted_html, elements
+
+
+def export_html(
+    inputs,
+    out_path: str,
+    ignore_line_break: bool = False,
+    export_figure: bool = True,
+    export_figure_letter: bool = False,
+    img=None,
+    figure_width=200,
+    figure_dir="figures",
+    encoding: str = "utf-8",
+):
+    formatted_html, elements = convert_html(
+        inputs,
+        out_path,
+        ignore_line_break,
+        export_figure,
+        export_figure_letter,
+        img,
+        figure_width,
+        figure_dir,
+    )
+
+    save_html(formatted_html, out_path, encoding)
+
+    return formatted_html
+
+
+def save_html(
+    html,
+    out_path,
+    encoding,
+):
+    with open(out_path, "w", encoding=encoding, errors="ignore") as f:
+        f.write(html)

yomitoku/export/export_json.py

@@ -1,4 +1,7 @@
 import json
+import os
+
+import cv2
 
 
 def paragraph_to_json(paragraph, ignore_line_break):
@@ -12,7 +15,28 @@ def table_to_json(table, ignore_line_break):
             cell.contents = cell.contents.replace("\n", "")
 
 
-def export_json(inputs, out_path, ignore_line_break=False):
+def save_figure(
+    figures,
+    img,
+    out_path,
+    figure_dir="figures",
+):
+    assert img is not None, "img is required for saving figures"
+
+    for i, figure in enumerate(figures):
+        x1, y1, x2, y2 = map(int, figure.box)
+        figure_img = img[y1:y2, x1:x2, :]
+        save_dir = os.path.dirname(out_path)
+        save_dir = os.path.join(save_dir, figure_dir)
+        os.makedirs(save_dir, exist_ok=True)
+
+        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)
+
+
+def convert_json(inputs, out_path, ignore_line_break, img, export_figure, figure_dir):
     from yomitoku.document_analyzer import DocumentAnalyzerSchema
 
     if isinstance(inputs, DocumentAnalyzerSchema):
@@ -23,9 +47,48 @@ def export_json(inputs, out_path, ignore_line_break=False):
         for paragraph in inputs.paragraphs:
             paragraph_to_json(paragraph, ignore_line_break)
 
-    with open(out_path, "w", encoding="utf-8") as f:
+    if isinstance(inputs, DocumentAnalyzerSchema) and export_figure:
+        save_figure(
+            inputs.figures,
+            img,
+            out_path,
+            figure_dir=figure_dir,
+        )
+
+    return inputs
+
+
+def export_json(
+    inputs,
+    out_path,
+    ignore_line_break=False,
+    encoding: str = "utf-8",
+    img=None,
+    export_figure=False,
+    figure_dir="figures",
+):
+    inputs = convert_json(
+        inputs,
+        out_path,
+        ignore_line_break,
+        img,
+        export_figure,
+        figure_dir,
+    )
+
+    save_json(
+        inputs.model_dump(),
+        out_path,
+        encoding,
+    )
+
+    return inputs
+
+
+def save_json(data, out_path, encoding):
+    with open(out_path, "w", encoding=encoding, errors="ignore") as f:
         json.dump(
-            inputs.model_dump(),
+            data,
             f,
             ensure_ascii=False,
             indent=4,