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

yomitoku/cli/main.py

@@ -1,5 +1,6 @@
 import argparse
 import os
+import torch
 from pathlib import Path
 
 import cv2
@@ -13,6 +14,18 @@ from ..utils.logger import set_logger
 logger = set_logger(__name__, "INFO")
 
 
+def validate_encoding(encoding):
+    if encoding not in [
+        "utf-8",
+        "utf-8-sig",
+        "shift-jis",
+        "euc-jp",
+        "cp932",
+    ]:
+        raise ValueError(f"Invalid encoding: {encoding}")
+    return True
+
+
 def process_single_file(args, analyzer, path, format):
     if path.suffix[1:].lower() in ["pdf"]:
         imgs = load_pdf(path)
@@ -21,7 +34,6 @@ def process_single_file(args, analyzer, path, format):
 
     for page, img in enumerate(imgs):
         results, ocr, layout = analyzer(img)
-
         dirname = path.parent.name
         filename = path.stem
 
@@ -47,11 +59,19 @@ def process_single_file(args, analyzer, path, format):
             results.to_json(
                 out_path,
                 ignore_line_break=args.ignore_line_break,
+                encoding=args.encoding,
+                img=img,
+                export_figure=args.figure,
+                figure_dir=args.figure_dir,
             )
         elif format == "csv":
             results.to_csv(
                 out_path,
                 ignore_line_break=args.ignore_line_break,
+                encoding=args.encoding,
+                img=img,
+                export_figure=args.figure,
+                figure_dir=args.figure_dir,
             )
         elif format == "html":
             results.to_html(
@@ -62,6 +82,7 @@ def process_single_file(args, analyzer, path, format):
                 export_figure_letter=args.figure_letter,
                 figure_width=args.figure_width,
                 figure_dir=args.figure_dir,
+                encoding=args.encoding,
             )
         elif format == "md":
             results.to_markdown(
@@ -72,6 +93,7 @@ def process_single_file(args, analyzer, path, format):
                 export_figure_letter=args.figure_letter,
                 figure_width=args.figure_width,
                 figure_dir=args.figure_dir,
+                encoding=args.encoding,
             )
 
         logger.info(f"Output file: {out_path}")
@@ -168,6 +190,12 @@ def main():
         default="figures",
         help="directory to save figure images",
     )
+    parser.add_argument(
+        "--encoding",
+        type=str,
+        default="utf-8",
+        help="Specifies the character encoding for the output file to be exported. If unsupported characters are included, they will be ignored.",
+    )
 
     args = parser.parse_args()
 
@@ -181,6 +209,8 @@ def main():
             f"Invalid output format: {args.format}. Supported formats are {SUPPORT_OUTPUT_FORMAT}"
         )
 
+    validate_encoding(args.encoding)
+
     if format == "markdown":
         format = "md"
 
@@ -205,7 +235,9 @@ def main():
 
     if args.lite:
         configs["ocr"]["text_recognizer"]["model_name"] = "parseq-small"
-        configs["ocr"]["text_detector"]["infer_onnx"] = True
+
+        if args.device == "cpu" or not torch.cuda.is_available():
+            configs["ocr"]["text_detector"]["infer_onnx"] = True
 
         # Note: Text Detector以外はONNX推論よりもPyTorch推論の方が速いため、ONNX推論は行わない
         # configs["ocr"]["text_recognizer"]["infer_onnx"] = True

yomitoku/document_analyzer.py

@@ -2,17 +2,26 @@ import asyncio
 from concurrent.futures import ThreadPoolExecutor
 from typing import List, Union
 
+import numpy as np
+
 from pydantic import conlist
 
 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 .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
 
 
 class ParagraphSchema(BaseSchema):
@@ -98,41 +107,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: (sum([p[1] for p in x.points]) / 4),
-        )
-    else:
-        contained_words = sorted(
-            contained_words,
-            key=lambda x: (sum([p[0] for p in x.points]) / 4),
-            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`が辞書の場合、再帰的に更新
@@ -148,8 +172,163 @@ 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):
         default_configs = {
             "ocr": {
                 "text_detector": {
@@ -180,8 +359,16 @@ 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
 
     def aggregate(self, ocr_res, layout_res):
@@ -286,16 +473,31 @@ 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
 

yomitoku/export/export_csv.py

@@ -1,4 +1,6 @@
 import csv
+import cv2
+import os
 
 
 def table_to_csv(table, ignore_line_break):
@@ -33,7 +35,36 @@ 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 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 = []
     for table in inputs.tables:
         table_csv = table_to_csv(table, ignore_line_break)
@@ -58,9 +89,17 @@ def export_csv(inputs, out_path: str, ignore_line_break: bool = False):
             }
         )
 
+    if export_figure:
+        save_figure(
+            inputs.figures,
+            img,
+            out_path,
+            figure_dir=figure_dir,
+        )
+
     elements = sorted(elements, key=lambda x: x["order"])
 
-    with open(out_path, "w", newline="", encoding="utf-8") as f:
+    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

@@ -110,6 +110,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)
@@ -154,6 +156,7 @@ def export_html(
     img=None,
     figure_width=200,
     figure_dir="figures",
+    encoding: str = "utf-8",
 ):
     html_string = ""
     elements = []
@@ -184,5 +187,5 @@ def export_html(
     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:
+    with open(out_path, "w", encoding=encoding, errors="ignore") as f:
         f.write(formatted_html)

yomitoku/export/export_json.py

@@ -1,5 +1,8 @@
 import json
 
+import cv2
+import os
+
 
 def paragraph_to_json(paragraph, ignore_line_break):
     if ignore_line_break:
@@ -12,7 +15,36 @@ 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 export_json(
+    inputs,
+    out_path,
+    ignore_line_break=False,
+    encoding: str = "utf-8",
+    img=None,
+    export_figure=False,
+    figure_dir="figures",
+):
     from yomitoku.document_analyzer import DocumentAnalyzerSchema
 
     if isinstance(inputs, DocumentAnalyzerSchema):
@@ -23,7 +55,15 @@ 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 export_figure:
+            save_figure(
+                inputs.figures,
+                img,
+                out_path,
+                figure_dir=figure_dir,
+            )
+
+    with open(out_path, "w", encoding=encoding, errors="ignore") as f:
         json.dump(
             inputs.model_dump(),
             f,

yomitoku/export/export_markdown.py

@@ -75,6 +75,8 @@ def figure_to_md(
     width=200,
     figure_dir="figures",
 ):
+    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)
@@ -117,6 +119,7 @@ def export_markdown(
     export_figure=True,
     figure_width=200,
     figure_dir="figures",
+    encoding: str = "utf-8",
 ):
     elements = []
     for table in inputs.tables:
@@ -141,5 +144,5 @@ def export_markdown(
     elements = sorted(elements, key=lambda x: x["order"])
     markdown = "\n".join([element["md"] for element in elements])
 
-    with open(out_path, "w", encoding="utf-8") as f:
+    with open(out_path, "w", encoding=encoding, errors="ignore") as f:
         f.write(markdown)

yomitoku/layout_analyzer.py

@@ -15,7 +15,7 @@ class LayoutAnalyzerSchema(BaseSchema):
 
 
 class LayoutAnalyzer:
-    def __init__(self, configs=None, device="cuda", visualize=False):
+    def __init__(self, configs={}, device="cuda", visualize=False):
         layout_parser_kwargs = {
             "device": device,
             "visualize": visualize,
@@ -26,10 +26,6 @@ class LayoutAnalyzer:
         }
 
         if isinstance(configs, dict):
-            assert (
-                "layout_parser" in configs or "table_structure_recognizer" in configs
-            ), "Invalid config key. Please check the config keys."
-
             if "layout_parser" in configs:
                 layout_parser_kwargs.update(configs["layout_parser"])
 

yomitoku/layout_parser.py

@@ -104,7 +104,6 @@ class LayoutParser(BaseModule):
         self.visualize = visualize
 
         self.model.eval()
-        self.model.to(self.device)
 
         self.postprocessor = RTDETRPostProcessor(
             num_classes=self._cfg.RTDETRTransformerv2.num_classes,
@@ -132,6 +131,8 @@ class LayoutParser(BaseModule):
             if not os.path.exists(path_onnx):
                 self.convert_onnx(path_onnx)
 
+            self.model = None
+
             model = onnx.load(path_onnx)
             if torch.cuda.is_available() and device == "cuda":
                 self.sess = onnxruntime.InferenceSession(
@@ -140,6 +141,9 @@ class LayoutParser(BaseModule):
             else:
                 self.sess = onnxruntime.InferenceSession(model.SerializeToString())
 
+        if self.model is not None:
+            self.model.to(self.device)
+
     def convert_onnx(self, path_onnx):
         dynamic_axes = {
             "input": {0: "batch_size"},

yomitoku/ocr.py

@@ -16,16 +16,37 @@ class WordPrediction(BaseSchema):
     )
     content: str
     direction: str
-    det_score: float
     rec_score: float
+    det_score: float
 
 
 class OCRSchema(BaseSchema):
     words: List[WordPrediction]
 
 
+def ocr_aggregate(det_outputs, rec_outputs):
+    words = []
+    for points, det_score, pred, rec_score, direction in zip(
+        det_outputs.points,
+        det_outputs.scores,
+        rec_outputs.contents,
+        rec_outputs.scores,
+        rec_outputs.directions,
+    ):
+        words.append(
+            {
+                "points": points,
+                "content": pred,
+                "direction": direction,
+                "det_score": det_score,
+                "rec_score": rec_score,
+            }
+        )
+    return words
+
+
 class OCR:
-    def __init__(self, configs=None, device="cuda", visualize=False):
+    def __init__(self, configs={}, device="cuda", visualize=False):
         text_detector_kwargs = {
             "device": device,
             "visualize": visualize,
@@ -36,10 +57,6 @@ class OCR:
         }
 
         if isinstance(configs, dict):
-            assert (
-                "text_detector" in configs or "text_recognizer" in configs
-            ), "Invalid config key. Please check the config keys."
-
             if "text_detector" in configs:
                 text_detector_kwargs.update(configs["text_detector"])
             if "text_recognizer" in configs:
@@ -52,26 +69,6 @@ class OCR:
         self.detector = TextDetector(**text_detector_kwargs)
         self.recognizer = TextRecognizer(**text_recognizer_kwargs)
 
-    def aggregate(self, det_outputs, rec_outputs):
-        words = []
-        for points, det_score, pred, rec_score, direction in zip(
-            det_outputs.points,
-            det_outputs.scores,
-            rec_outputs.contents,
-            rec_outputs.scores,
-            rec_outputs.directions,
-        ):
-            words.append(
-                {
-                    "points": points,
-                    "content": pred,
-                    "direction": direction,
-                    "det_score": det_score,
-                    "rec_score": rec_score,
-                }
-            )
-        return words
-
     def __call__(self, img):
         """_summary_
 
@@ -82,6 +79,6 @@ class OCR:
         det_outputs, vis = self.detector(img)
         rec_outputs, vis = self.recognizer(img, det_outputs.points, vis=vis)
 
-        outputs = {"words": self.aggregate(det_outputs, rec_outputs)}
+        outputs = {"words": ocr_aggregate(det_outputs, rec_outputs)}
         results = OCRSchema(**outputs)
         return results, vis

yomitoku/table_structure_recognizer.py

@@ -35,10 +35,17 @@ class TableCellSchema(BaseSchema):
     contents: Union[str, None]
 
 
+class TableLineSchema(BaseSchema):
+    box: conlist(int, min_length=4, max_length=4)
+    score: float
+
+
 class TableStructureRecognizerSchema(BaseSchema):
     box: conlist(int, min_length=4, max_length=4)
     n_row: int
     n_col: int
+    rows: List[TableLineSchema]
+    cols: List[TableLineSchema]
     cells: List[TableCellSchema]
     order: int
 
@@ -133,8 +140,6 @@ class TableStructureRecognizer(BaseModule):
             num_top_queries=self._cfg.RTDETRTransformerv2.num_queries,
         )
 
-        self.save_config("table_structure_recognitizer.yaml")
-
         self.transforms = T.Compose(
             [
                 T.Resize(self._cfg.data.img_size),
@@ -155,6 +160,8 @@ class TableStructureRecognizer(BaseModule):
             if not os.path.exists(path_onnx):
                 self.convert_onnx(path_onnx)
 
+            self.model = None
+
             model = onnx.load(path_onnx)
             if torch.cuda.is_available() and device == "cuda":
                 self.sess = onnxruntime.InferenceSession(
@@ -163,6 +170,9 @@ class TableStructureRecognizer(BaseModule):
             else:
                 self.sess = onnxruntime.InferenceSession(model.SerializeToString())
 
+        if self.model is not None:
+            self.model.to(self.device)
+
     def convert_onnx(self, path_onnx):
         dynamic_axes = {
             "input": {0: "batch_size"},
@@ -232,7 +242,7 @@ class TableStructureRecognizer(BaseModule):
             category_elements
         )
 
-        cells, n_row, n_col = self.extract_cell_elements(category_elements)
+        cells, rows, cols = self.extract_cell_elements(category_elements)
 
         table_x, table_y = data["offset"]
         table_x2 = table_x + data["size"][1]
@@ -241,8 +251,10 @@ class TableStructureRecognizer(BaseModule):
 
         table = {
             "box": table_box,
-            "n_row": n_row,
-            "n_col": n_col,
+            "n_row": len(rows),
+            "n_col": len(cols),
+            "rows": rows,
+            "cols": cols,
             "cells": cells,
             "order": 0,
         }
@@ -262,7 +274,10 @@ class TableStructureRecognizer(BaseModule):
         cells = extract_cells(row_boxes, col_boxes)
         cells = filter_contained_cells_within_spancell(cells, span_boxes)
 
-        return cells, len(row_boxes), len(col_boxes)
+        rows = sorted(elements["row"], key=lambda x: x["box"][1])
+        cols = sorted(elements["col"], key=lambda x: x["box"][0])
+
+        return cells, rows, cols
 
     def __call__(self, img, table_boxes, vis=None):
         img_tensors = self.preprocess(img, table_boxes)
@@ -282,7 +297,9 @@ class TableStructureRecognizer(BaseModule):
                     pred = self.model(data["tensor"])
 
             table = self.postprocess(pred, data)
-            outputs.append(table)
+
+            if table.n_row > 0 and table.n_col > 0:
+                outputs.append(table)
 
         if vis is None and self.visualize:
             vis = img.copy()

yomitoku/text_detector.py

@@ -61,8 +61,6 @@ class TextDetector(BaseModule):
         self.visualize = visualize
 
         self.model.eval()
-        self.model.to(self.device)
-
         self.post_processor = DBnetPostProcessor(**self._cfg.post_process)
         self.infer_onnx = infer_onnx
 
@@ -72,6 +70,8 @@ class TextDetector(BaseModule):
             if not os.path.exists(path_onnx):
                 self.convert_onnx(path_onnx)
 
+            self.model = None
+
             model = onnx.load(path_onnx)
             if torch.cuda.is_available() and device == "cuda":
                 self.sess = onnxruntime.InferenceSession(
@@ -80,6 +80,11 @@ class TextDetector(BaseModule):
             else:
                 self.sess = onnxruntime.InferenceSession(model.SerializeToString())
 
+            self.model = None
+
+        if self.model is not None:
+            self.model.to(self.device)
+
     def convert_onnx(self, path_onnx):
         dynamic_axes = {
             "input": {0: "batch_size", 2: "height", 3: "width"},
@@ -138,9 +143,9 @@ class TextDetector(BaseModule):
         vis = None
         if self.visualize:
             vis = det_visualizer(
-                preds,
                 img,
                 quads,
+                preds=preds,
                 vis_heatmap=self._cfg.visualize.heatmap,
                 line_color=tuple(self._cfg.visualize.color[::-1]),
             )

yomitoku/text_recognizer.py

@@ -64,7 +64,6 @@ class TextRecognizer(BaseModule):
 
         self.model.tokenizer = self.tokenizer
         self.model.eval()
-        self.model.to(self.device)
 
         self.visualize = visualize
 
@@ -76,6 +75,8 @@ class TextRecognizer(BaseModule):
             if not os.path.exists(path_onnx):
                 self.convert_onnx(path_onnx)
 
+            self.model = None
+
             model = onnx.load(path_onnx)
             if torch.cuda.is_available() and device == "cuda":
                 self.sess = onnxruntime.InferenceSession(
@@ -84,17 +85,31 @@ class TextRecognizer(BaseModule):
             else:
                 self.sess = onnxruntime.InferenceSession(model.SerializeToString())
 
+        if self.model is not None:
+            self.model.to(self.device)
+
     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,
-        )
+        dataloader = self._make_mini_batch(dataset)
 
         return dataloader
 
+    def _make_mini_batch(self, dataset):
+        mini_batches = []
+        mini_batch = []
+        for data in dataset:
+            data = torch.unsqueeze(data, 0)
+            mini_batch.append(data)
+
+            if len(mini_batch) == self._cfg.data.batch_size:
+                mini_batches.append(torch.cat(mini_batch, 0))
+                mini_batch = []
+        else:
+            if len(mini_batch) > 0:
+                mini_batches.append(torch.cat(mini_batch, 0))
+
+        return mini_batches
+
     def convert_onnx(self, path_onnx):
         img_size = self._cfg.data.img_size
         input = torch.randn(1, 3, *img_size, requires_grad=True)

yomitoku/utils/misc.py

@@ -9,6 +9,24 @@ def filter_by_flag(elements, flags):
     return [element for element, flag in zip(elements, flags) if flag]
 
 
+def calc_overlap_ratio(rect_a, rect_b):
+    intersection = calc_intersection(rect_a, rect_b)
+    if intersection is None:
+        return 0, None
+
+    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
+
+    overlap_ratio = overlap_area / b_area
+    return overlap_ratio, intersection
+
+
 def is_contained(rect_a, rect_b, threshold=0.8):
     """二つの矩形A, Bが与えられたとき、矩形Bが矩形Aに含まれるかどうかを判定する。
     ずれを許容するため、重複率求め、thresholdを超える場合にTrueを返す。
@@ -23,20 +41,9 @@ def is_contained(rect_a, rect_b, threshold=0.8):
         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
+    overlap_ratio, _ = calc_overlap_ratio(rect_a, rect_b)
 
-    if overlap_area / b_area > threshold:
+    if overlap_ratio > threshold:
         return True
 
     return False

yomitoku/utils/visualizer.py

@@ -66,14 +66,14 @@ def reading_order_visualizer(
     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()
+def det_visualizer(img, quads, preds=None, vis_heatmap=False, line_color=(0, 255, 0)):
     out = img.copy()
     h, w = out.shape[:2]
-    binary = binary.squeeze(0)
-    binary = (binary * 255).astype(np.uint8)
     if vis_heatmap:
+        preds = preds["binary"][0]
+        binary = preds.detach().cpu().numpy()
+        binary = binary.squeeze(0)
+        binary = (binary * 255).astype(np.uint8)
         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)

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: yomitoku
-Version: 0.6.0
+Version: 0.7.1
 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
@@ -10,7 +10,6 @@ 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
@@ -20,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: gpu
+Requires-Dist: onnxruntime-gpu>=1.20.1; extra == 'gpu'
 Description-Content-Type: text/markdown
 
 日本語版 | [English](README_EN.md)
@@ -72,6 +73,11 @@ Markdown でエクスポートした結果は関してはリポジトリ内の[s
 pip install yomitoku
 ```
 
+onnxruntimeの実行にGPUを使用する場合
+```
+pip install yomitoku[gpu]
+```
+
 - pytorch はご自身の CUDA のバージョンにあったものをインストールしてください。デフォルトでは CUDA12.4 以上に対応したものがインストールされます。
 - pytorch は 2.5 以上のバージョンに対応しています。その関係で CUDA11.8 以上のバージョンが必要になります。対応できない場合は、リポジトリ内の Dockerfile を利用してください。
 
@@ -89,7 +95,8 @@ yomitoku ${path_data} -f md -o results -v --figure　--lite
 - `-d`, `--device` モデルを実行するためのデバイスを指定します。gpu が利用できない場合は cpu で推論が実行されます。(デフォルト: cuda)
 - `--ignore_line_break` 画像の改行位置を無視して、段落内の文章を連結して返します。（デフォルト：画像通りの改行位置位置で改行します。）
 - `--figure_letter` 検出した図表に含まれる文字も出力ファイルにエクスポートします。
-- `--figure` 検出した図、画像を出力ファイルにエクスポートします。(html と markdown のみ)
+- `--figure` 検出した図、画像を出力ファイルにエクスポートします。
+- `--encoding` エクスポートする出力ファイルの文字エンコーディングを指定します。サポートされていない文字コードが含まれる場合は、その文字を無視します。(utf-8, utf-8-sig, shift-jis, enc-jp, cp932)
 
 その他のオプションに関しては、ヘルプを参照
 

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

@@ -1,16 +1,16 @@
 yomitoku/__init__.py,sha256=kXOM8RbpwwLABG3p3vPT3dJWBk4JX2MFGrOeBEW0hKM,543
 yomitoku/base.py,sha256=lzR_V8t87aRasmFdFwD-8KAeSahSTI3AZaEn6g8sOv8,3871
 yomitoku/constants.py,sha256=zlW5QRc_u_F3C2RAgBFWyHJZexBnJT5N15GC-9d3iLo,686
-yomitoku/document_analyzer.py,sha256=HIg-nVzDhJIP-h-tn4uU86KakgHdlAhosEqK_i-SWe4,9906
-yomitoku/layout_analyzer.py,sha256=QTeRcVd8aySz8u6dg2ikET77ar3sqlukRLBwYfTyMPM,2033
-yomitoku/layout_parser.py,sha256=Yni1C_7j4fzHcdmBNNGRZPc23W_6J6HwPPQVjYvaztM,7539
-yomitoku/ocr.py,sha256=Rcojw0aGA6yDF2RjqfK23_rMw-xm61KGd8JmTCTOOVU,2516
+yomitoku/document_analyzer.py,sha256=B2F_MXFKbq58ePDCgcZKk_bgQUkno1ehYb6CZmAekCk,16234
+yomitoku/layout_analyzer.py,sha256=VhNf1ZQFoozj6WUGk5ll1p2p1jk5X3j-JPcDbTAoSl4,1856
+yomitoku/layout_parser.py,sha256=V_mAkZxke1gwHfnxBFMTOJ8hnz2X_kfZu2lLiMd8cAs,7610
+yomitoku/ocr.py,sha256=JSTjkupcxHITQm6ERnzU7As0c3KWf8-oxc0AqNoWHXo,2272
 yomitoku/reading_order.py,sha256=OfhOS9ttPDoPSuHrIRKyOzG19GGeRufbuSKDqhsohh4,6404
-yomitoku/table_structure_recognizer.py,sha256=Wf_Ehmf6V27iVLmw2o9i7kJnbwEOhuExI-ljIO3a8NE,9043
-yomitoku/text_detector.py,sha256=fbwKelsVfwCt5YL4h-WEf4qkniv5cXmyaLR6oSYz0eA,4167
-yomitoku/text_recognizer.py,sha256=Iu-IzwaziNjmrTeSw9aoN9BDTHkNOzsZhViCv45yiN8,5422
+yomitoku/table_structure_recognizer.py,sha256=Eam9t7OjW4a-UWk_dl-ylbOcinN_Te_ovuri2naldL0,9482
+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=qDB_YNK7abstIr9tYLiJjNU3xLSCd5x1UNDKqwUi2Rk,6885
+yomitoku/cli/main.py,sha256=Dcu-jBql5dP4DiN0QWbPAQscnJ7tTtbWUnaohz3FFbs,7868
 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
@@ -21,10 +21,10 @@ 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=-n8eYPIzDQuiixeqpTbWaN9aQ5oFyl7XRfpv51oKPTI,1979
-yomitoku/export/export_html.py,sha256=X3H_orkS1BRlQo8Z1NzgrFwsIboDzRAx9etmqj90k2Y,4866
-yomitoku/export/export_json.py,sha256=1ChvCAHfCmMQvCfcAb1p3fSpr4elNAs3xBSIbpfn3bc,998
-yomitoku/export/export_markdown.py,sha256=mCcsXUWBLrYc1NcRSBFfBT28d6eCddAF1oHp0qdBEnE,3986
+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/models/__init__.py,sha256=Enxq9sjJWusZuxecTori8IQa8NEYKaiiptDluHX1avg,144
 yomitoku/models/dbnet_plus.py,sha256=jeWJZm0ihbxoJeAXBFK7uVIwoosx2IUNk7Ut5wRH0vA,7998
 yomitoku/models/parseq.py,sha256=-DQMQuON2jwtb4Ib2V0O19un9w-WG4rXS0SiscydrXU,8593
@@ -46,9 +46,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=2Eyy7-9K_h4Mal1zGXq6OlxubfNzhS0mEYwn_xt7xl8,2497
-yomitoku/utils/visualizer.py,sha256=2pSmbhUPylzVVJ0bXtGDoNmMdArAByab4Py7Xavvs_A,5230
-yomitoku-0.6.0.dist-info/METADATA,sha256=XDmMBtDx9MjXPuzcARwOwJXRN7PMCsQDwc38jDSwX5g,8134
-yomitoku-0.6.0.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
-yomitoku-0.6.0.dist-info/entry_points.txt,sha256=nFV3S11zgBNW0Qq_D0XQNg2R4lNXU_9XUFr6rdJoyF8,52
-yomitoku-0.6.0.dist-info/RECORD,,
+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.6.0.dist-info → yomitoku-0.7.1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.26.3
+Generator: hatchling 1.27.0
 Root-Is-Purelib: true
 Tag: py3-none-any