yomitoku 0.7.3__py3-none-any.whl → 0.8.0__py3-none-any.whl

yomitoku/cli/main.py

@@ -12,6 +12,7 @@ from ..document_analyzer import DocumentAnalyzer
 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
 
 logger = set_logger(__name__, "INFO")
 
@@ -51,13 +52,13 @@ def merge_all_pages(results):
 
 def save_merged_file(out_path, args, out):
     if args.format == "json":
-        save_json(out_path, args.encoding, out)
+        save_json(out, out_path, args.encoding)
     elif args.format == "csv":
-        save_csv(out_path, args.encoding, out)
+        save_csv(out, out_path, args.encoding)
     elif args.format == "html":
-        save_html(out_path, args.encoding, out)
+        save_html(out, out_path, args.encoding)
     elif args.format == "md":
-        save_markdown(out_path, args.encoding, out)
+        save_markdown(out, out_path, args.encoding)
 
 
 def validate_encoding(encoding):
@@ -76,7 +77,7 @@ def process_single_file(args, analyzer, path, format):
     if path.suffix[1:].lower() in ["pdf"]:
         imgs = load_pdf(path)
     else:
-        imgs = [load_image(path)]
+        imgs = load_image(path)
 
     results = []
     for page, img in enumerate(imgs):
@@ -84,6 +85,10 @@ def process_single_file(args, analyzer, path, format):
         dirname = path.parent.name
         filename = path.stem
 
+        # cv2.imwrite(
+        #    os.path.join(args.outdir, f"{dirname}_{filename}_p{page+1}.jpg"), img
+        # )
+
         if ocr is not None:
             out_path = os.path.join(
                 args.outdir, f"{dirname}_{filename}_p{page+1}_ocr.jpg"
@@ -103,34 +108,51 @@ def process_single_file(args, analyzer, path, format):
         out_path = os.path.join(args.outdir, f"{dirname}_{filename}_p{page+1}.{format}")
 
         if format == "json":
-            json = result.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,
-            )
+            if args.combine:
+                json = convert_json(
+                    result,
+                    out_path,
+                    args.ignore_line_break,
+                    img,
+                    args.figure,
+                    args.figure_dir,
+                )
+            else:
+                json = result.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,
+                )
 
             results.append(
                 {
                     "format": format,
-                    "data": json,
+                    "data": json.model_dump(),
                 }
             )
 
-            if not args.combine:
-                save_json(out_path, args.encoding, json)
-
         elif format == "csv":
-            csv = result.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,
-            )
+            if args.combine:
+                csv = convert_csv(
+                    result,
+                    out_path,
+                    args.ignore_line_break,
+                    img,
+                    args.figure,
+                    args.figure_dir,
+                )
+            else:
+                csv = result.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,
+                )
 
             results.append(
                 {
@@ -139,20 +161,29 @@ def process_single_file(args, analyzer, path, format):
                 }
             )
 
-            if not args.combine:
-                save_csv(out_path, args.encoding, csv)
-
         elif format == "html":
-            html = result.to_html(
-                out_path,
-                ignore_line_break=args.ignore_line_break,
-                img=img,
-                export_figure=args.figure,
-                export_figure_letter=args.figure_letter,
-                figure_width=args.figure_width,
-                figure_dir=args.figure_dir,
-                encoding=args.encoding,
-            )
+            if args.combine:
+                html, _ = convert_html(
+                    result,
+                    out_path,
+                    ignore_line_break=args.ignore_line_break,
+                    img=img,
+                    export_figure=args.figure,
+                    export_figure_letter=args.figure_letter,
+                    figure_width=args.figure_width,
+                    figure_dir=args.figure_dir,
+                )
+            else:
+                html = result.to_html(
+                    out_path,
+                    ignore_line_break=args.ignore_line_break,
+                    img=img,
+                    export_figure=args.figure,
+                    export_figure_letter=args.figure_letter,
+                    figure_width=args.figure_width,
+                    figure_dir=args.figure_dir,
+                    encoding=args.encoding,
+                )
 
             results.append(
                 {
@@ -161,20 +192,29 @@ def process_single_file(args, analyzer, path, format):
                 }
             )
 
-            if not args.combine:
-                save_html(out_path, args.encoding, html)
-
         elif format == "md":
-            md = result.to_markdown(
-                out_path,
-                ignore_line_break=args.ignore_line_break,
-                img=img,
-                export_figure=args.figure,
-                export_figure_letter=args.figure_letter,
-                figure_width=args.figure_width,
-                figure_dir=args.figure_dir,
-                encoding=args.encoding,
-            )
+            if args.combine:
+                md, _ = convert_markdown(
+                    result,
+                    out_path,
+                    ignore_line_break=args.ignore_line_break,
+                    img=img,
+                    export_figure=args.figure,
+                    export_figure_letter=args.figure_letter,
+                    figure_width=args.figure_width,
+                    figure_dir=args.figure_dir,
+                )
+            else:
+                md = result.to_markdown(
+                    out_path,
+                    ignore_line_break=args.ignore_line_break,
+                    img=img,
+                    export_figure=args.figure,
+                    export_figure_letter=args.figure_letter,
+                    figure_width=args.figure_width,
+                    figure_dir=args.figure_dir,
+                    encoding=args.encoding,
+                )
 
             results.append(
                 {
@@ -183,9 +223,6 @@ def process_single_file(args, analyzer, path, format):
                 }
             )
 
-            if not args.combine:
-                save_markdown(out_path, args.encoding, md)
-
     out = merge_all_pages(results)
     if args.combine:
         out_path = os.path.join(args.outdir, f"{dirname}_{filename}.{format}")

yomitoku/configs/__init__.py

@@ -1,10 +1,14 @@
 from .cfg_layout_parser_rtdtrv2 import LayoutParserRTDETRv2Config
+from .cfg_layout_parser_rtdtrv2_v2 import LayoutParserRTDETRv2V2Config
 from .cfg_table_structure_recognizer_rtdtrv2 import (
     TableStructureRecognizerRTDETRv2Config,
 )
 from .cfg_text_detector_dbnet import TextDetectorDBNetConfig
+from .cfg_text_detector_dbnet_v2 import TextDetectorDBNetV2Config
 from .cfg_text_recognizer_parseq import TextRecognizerPARSeqConfig
 from .cfg_text_recognizer_parseq_small import TextRecognizerPARSeqSmallConfig
+from .cfg_text_recognizer_parseq_v2 import TextRecognizerPARSeqV2Config
+
 
 __all__ = [
     "TextDetectorDBNetConfig",
@@ -12,4 +16,7 @@ __all__ = [
     "LayoutParserRTDETRv2Config",
     "TableStructureRecognizerRTDETRv2Config",
     "TextRecognizerPARSeqSmallConfig",
+    "LayoutParserRTDETRv2V2Config",
+    "TextDetectorDBNetV2Config",
+    "TextRecognizerPARSeqV2Config",
 ]

yomitoku/configs/cfg_layout_parser_rtdtrv2_v2.py

@@ -0,0 +1,89 @@
+from dataclasses import dataclass, field
+from typing import List
+
+
+@dataclass
+class Data:
+    img_size: List[int] = field(default_factory=lambda: [640, 640])
+
+
+@dataclass
+class BackBone:
+    depth: int = 50
+    variant: str = "d"
+    freeze_at: int = 0
+    return_idx: List[int] = field(default_factory=lambda: [1, 2, 3])
+    num_stages: int = 4
+    freeze_norm: bool = True
+
+
+@dataclass
+class Encoder:
+    in_channels: List[int] = field(default_factory=lambda: [512, 1024, 2048])
+    feat_strides: List[int] = field(default_factory=lambda: [8, 16, 32])
+
+    # intra
+    hidden_dim: int = 256
+    use_encoder_idx: List[int] = field(default_factory=lambda: [2])
+    num_encoder_layers: int = 1
+    nhead: int = 8
+    dim_feedforward: int = 1024
+    dropout: float = 0.0
+    enc_act: str = "gelu"
+
+    # cross
+    expansion: float = 1.0
+    depth_mult: int = 1
+    act: str = "silu"
+
+
+@dataclass
+class Decoder:
+    num_classes: int = 6
+    feat_channels: List[int] = field(default_factory=lambda: [256, 256, 256])
+    feat_strides: List[int] = field(default_factory=lambda: [8, 16, 32])
+    hidden_dim: int = 256
+    num_levels: int = 3
+
+    num_layers: int = 6
+    num_queries: int = 300
+
+    num_denoising: int = 100
+    label_noise_ratio: float = 0.5
+    box_noise_scale: float = 1.0
+    eval_spatial_size: List[int] = field(default_factory=lambda: [640, 640])
+
+    eval_idx: int = -1
+
+    num_points: List[int] = field(default_factory=lambda: [4, 4, 4])
+    cross_attn_method: str = "default"
+    query_select_method: str = "default"
+
+
+@dataclass
+class LayoutParserRTDETRv2V2Config:
+    hf_hub_repo: str = "KotaroKinoshita/yomitoku-layout-parser-rtdtrv2-v2"
+    thresh_score: float = 0.5
+    data: Data = field(default_factory=Data)
+    PResNet: BackBone = field(default_factory=BackBone)
+    HybridEncoder: Encoder = field(default_factory=Encoder)
+    RTDETRTransformerv2: Decoder = field(default_factory=Decoder)
+
+    category: List[str] = field(
+        default_factory=lambda: [
+            "tables",
+            "figures",
+            "paragraphs",
+            "section_headings",
+            "page_header",
+            "page_footer",
+        ]
+    )
+
+    role: List[str] = field(
+        default_factory=lambda: [
+            "section_headings",
+            "page_header",
+            "page_footer",
+        ]
+    )

yomitoku/configs/cfg_text_detector_dbnet_v2.py

@@ -0,0 +1,49 @@
+from dataclasses import dataclass, field
+from typing import List
+
+
+@dataclass
+class BackBone:
+    name: str = "resnet50"
+    dilation: bool = True
+
+
+@dataclass
+class Decoder:
+    in_channels: list[int] = field(default_factory=lambda: [256, 512, 1024, 2048])
+    hidden_dim: int = 256
+    adaptive: bool = True
+    serial: bool = True
+    smooth: bool = False
+    k: int = 50
+
+
+@dataclass
+class Data:
+    shortest_size: int = 1280
+    limit_size: int = 1600
+
+
+@dataclass
+class PostProcess:
+    min_size: int = 2
+    thresh: float = 0.4
+    box_thresh: float = 0.5
+    max_candidates: int = 1500
+    unclip_ratio: float = 6.0
+
+
+@dataclass
+class Visualize:
+    color: List[int] = field(default_factory=lambda: [0, 255, 0])
+    heatmap: bool = False
+
+
+@dataclass
+class TextDetectorDBNetV2Config:
+    hf_hub_repo: str = "KotaroKinoshita/yomitoku-text-detector-dbnet-v2"
+    backbone: BackBone = field(default_factory=BackBone)
+    decoder: Decoder = field(default_factory=Decoder)
+    data: Data = field(default_factory=Data)
+    post_process: PostProcess = field(default_factory=PostProcess)
+    visualize: Visualize = field(default_factory=Visualize)

yomitoku/configs/cfg_text_recognizer_parseq_v2.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass, field
+from typing import List
+
+from ..constants import ROOT_DIR
+
+
+@dataclass
+class Data:
+    num_workers: int = 4
+    batch_size: int = 128
+    img_size: List[int] = field(default_factory=lambda: [32, 800])
+
+
+@dataclass
+class Encoder:
+    patch_size: List[int] = field(default_factory=lambda: [8, 8])
+    num_heads: int = 8
+    embed_dim: int = 512
+    mlp_ratio: int = 4
+    depth: int = 12
+
+
+@dataclass
+class Decoder:
+    embed_dim: int = 512
+    num_heads: int = 8
+    mlp_ratio: int = 4
+    depth: int = 1
+
+
+@dataclass
+class Visualize:
+    font: str = str(ROOT_DIR + "/resource/MPLUS1p-Medium.ttf")
+    color: List[int] = field(default_factory=lambda: [0, 0, 255])  # RGB
+    font_size: int = 18
+
+
+@dataclass
+class TextRecognizerPARSeqV2Config:
+    hf_hub_repo: str = "KotaroKinoshita/yomitoku-text-recognizer-parseq-middle-v2"
+    charset: str = str(ROOT_DIR + "/resource/charset.txt")
+    num_tokens: int = 7312
+    max_label_length: int = 100
+    decode_ar: int = 1
+    refine_iters: int = 1
+
+    data: Data = field(default_factory=Data)
+    encoder: Encoder = field(default_factory=Encoder)
+    decoder: Decoder = field(default_factory=Decoder)
+
+    visualize: Visualize = field(default_factory=Visualize)

yomitoku/data/functions.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
 import cv2
+from PIL import Image
 import numpy as np
 import torch
 import pypdfium2
@@ -15,6 +16,20 @@ from ..utils.logger import set_logger
 logger = set_logger(__name__)
 
 
+def validate_image(img: np.ndarray):
+    h, w = img.shape[:2]
+    if h < MIN_IMAGE_SIZE or w < MIN_IMAGE_SIZE:
+        raise ValueError("Image size is too small.")
+
+    if min(h, w) < WARNING_IMAGE_SIZE:
+        logger.warning(
+            """
+            The image size is small, which may result in reduced OCR accuracy. 
+            The process will continue, but it is recommended to input images with a minimum size of 720 pixels on the shorter side.
+            """
+        )
+
+
 def load_image(image_path: str) -> np.ndarray:
     """
     Open an image file.
@@ -40,24 +55,27 @@ def load_image(image_path: str) -> np.ndarray:
             "PDF file is not supported by load_image(). Use load_pdf() instead."
         )
 
-    img = cv2.imread(image_path, cv2.IMREAD_COLOR)
-
-    if img is None:
+    try:
+        img = Image.open(image_path)
+    except Exception:
         raise ValueError("Invalid image data.")
 
-    h, w = img.shape[:2]
-    if h < MIN_IMAGE_SIZE or w < MIN_IMAGE_SIZE:
-        raise ValueError("Image size is too small.")
-
-    if min(h, w) < WARNING_IMAGE_SIZE:
-        logger.warning(
-            """
-            The image size is small, which may result in reduced OCR accuracy. 
-            The process will continue, but it is recommended to input images with a minimum size of 720 pixels on the shorter side.
-            """
-        )
+    pages = []
+    if ext in ["tif", "tiff"]:
+        try:
+            while True:
+                img_arr = np.array(img.copy().convert("RGB"))
+                validate_image(img_arr)
+                pages.append(img_arr[:, :, ::-1])
+                img.seek(img.tell() + 1)
+        except EOFError:
+            pass
+    else:
+        img_arr = np.array(img.convert("RGB"))
+        validate_image(img_arr)
+        pages.append(img_arr[:, :, ::-1])
 
-    return img
+    return pages
 
 
 def load_pdf(pdf_path: str, dpi=200) -> list[np.ndarray]:

yomitoku/export/__init__.py

@@ -1,7 +1,7 @@
-from .export_csv import export_csv, save_csv
-from .export_html import export_html, save_html
-from .export_json import export_json, save_json
-from .export_markdown import export_markdown, save_markdown
+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",
@@ -12,4 +12,8 @@ __all__ = [
     "save_markdown",
     "save_csv",
     "save_json",
+    "convert_html",
+    "convert_markdown",
+    "convert_csv",
+    "convert_json",
 ]

yomitoku/export/export_csv.py

@@ -57,11 +57,10 @@ def save_figure(
         cv2.imwrite(figure_path, figure_img)
 
 
-def export_csv(
+def convert_csv(
     inputs,
-    out_path: str,
-    ignore_line_break: bool = False,
-    encoding: str = "utf-8",
+    out_path,
+    ignore_line_break,
     img=None,
     export_figure: bool = True,
     figure_dir="figures",
@@ -90,6 +89,8 @@ def export_csv(
             }
         )
 
+    elements = sorted(elements, key=lambda x: x["order"])
+
     if export_figure:
         save_figure(
             inputs.figures,
@@ -98,11 +99,36 @@ def export_csv(
             figure_dir=figure_dir,
         )
 
-    elements = sorted(elements, key=lambda x: x["order"])
     return elements
 
 
-def save_csv(out_path, encoding, 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:

yomitoku/export/export_html.py

@@ -146,16 +146,15 @@ 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",
-    encoding: str = "utf-8",
 ):
     html_string = ""
     elements = []
@@ -181,13 +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")
+
+    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(out_path, encoding, 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

@@ -36,15 +36,7 @@ def save_figure(
         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",
-):
+def convert_json(inputs, out_path, ignore_line_break, img, export_figure, figure_dir):
     from yomitoku.document_analyzer import DocumentAnalyzerSchema
 
     if isinstance(inputs, DocumentAnalyzerSchema):
@@ -55,18 +47,45 @@ def export_json(
         for paragraph in inputs.paragraphs:
             paragraph_to_json(paragraph, ignore_line_break)
 
-        if export_figure:
-            save_figure(
-                inputs.figures,
-                img,
-                out_path,
-                figure_dir=figure_dir,
-            )
+    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.model_dump()
+    return inputs
 
 
-def save_json(out_path, encoding, data):
+def save_json(data, out_path, encoding):
     with open(out_path, "w", encoding=encoding, errors="ignore") as f:
         json.dump(
             data,

yomitoku/export/export_markdown.py

@@ -111,16 +111,15 @@ def figure_to_md(
     return elements
 
 
-def export_markdown(
+def convert_markdown(
     inputs,
-    out_path: str,
+    out_path,
+    ignore_line_break=False,
     img=None,
-    ignore_line_break: bool = False,
     export_figure_letter=False,
     export_figure=True,
     figure_width=200,
     figure_dir="figures",
-    encoding: str = "utf-8",
 ):
     elements = []
     for table in inputs.tables:
@@ -144,10 +143,39 @@ def export_markdown(
 
     elements = sorted(elements, key=lambda x: x["order"])
     markdown = "\n".join([element["md"] for element in elements])
+    return markdown, elements
+
 
+def export_markdown(
+    inputs,
+    out_path: str,
+    ignore_line_break: bool = False,
+    img=None,
+    export_figure_letter=False,
+    export_figure=True,
+    figure_width=200,
+    figure_dir="figures",
+    encoding: str = "utf-8",
+):
+    markdown, elements = convert_markdown(
+        inputs,
+        out_path,
+        ignore_line_break,
+        img,
+        export_figure_letter,
+        export_figure,
+        figure_width,
+        figure_dir,
+    )
+
+    save_markdown(markdown, out_path, encoding)
     return markdown
 
 
-def save_markdown(out_path, encoding, markdown):
+def save_markdown(
+    markdown,
+    out_path,
+    encoding,
+):
     with open(out_path, "w", encoding=encoding, errors="ignore") as f:
         f.write(markdown)

yomitoku/layout_parser.py

@@ -12,7 +12,7 @@ from pydantic import conlist
 from .constants import ROOT_DIR
 
 from .base import BaseModelCatalog, BaseModule, BaseSchema
-from .configs import LayoutParserRTDETRv2Config
+from .configs import LayoutParserRTDETRv2Config, LayoutParserRTDETRv2V2Config
 from .models import RTDETRv2
 from .postprocessor import RTDETRPostProcessor
 from .utils.misc import filter_by_flag, is_contained
@@ -35,6 +35,7 @@ class LayoutParserModelCatalog(BaseModelCatalog):
     def __init__(self):
         super().__init__()
         self.register("rtdetrv2", LayoutParserRTDETRv2Config, RTDETRv2)
+        self.register("rtdetrv2v2", LayoutParserRTDETRv2V2Config, RTDETRv2)
 
 
 def filter_contained_rectangles_within_category(category_elements):
@@ -91,7 +92,7 @@ class LayoutParser(BaseModule):
 
     def __init__(
         self,
-        model_name="rtdetrv2",
+        model_name="rtdetrv2v2",
         path_cfg=None,
         device="cuda",
         visualize=False,

yomitoku/postprocessor/rtdetr_postprocessor.py

@@ -49,6 +49,13 @@ class RTDETRPostProcessor(nn.Module):
     def extra_repr(self) -> str:
         return f"use_focal_loss={self.use_focal_loss}, num_classes={self.num_classes}, num_top_queries={self.num_top_queries}"
 
+    def clamp(self, boxes, h, w):
+        boxes[:, 0] = torch.clamp(boxes[:, 0], min=torch.Tensor([0]), max=None)
+        boxes[:, 1] = torch.clamp(boxes[:, 1], min=torch.Tensor([0]), max=None)
+        boxes[:, 2] = torch.clamp(boxes[:, 2], min=torch.Tensor([0]), max=w)
+        boxes[:, 3] = torch.clamp(boxes[:, 3], min=torch.Tensor([0]), max=h)
+        return boxes
+
     # def forward(self, outputs, orig_target_sizes):
     def forward(self, outputs, orig_target_sizes: torch.Tensor, threshold):
         logits, boxes = outputs["pred_logits"], outputs["pred_boxes"]
@@ -57,6 +64,8 @@ class RTDETRPostProcessor(nn.Module):
         bbox_pred = torchvision.ops.box_convert(boxes, in_fmt="cxcywh", out_fmt="xyxy")
         bbox_pred *= orig_target_sizes.repeat(1, 2).unsqueeze(1)
 
+        w, h = orig_target_sizes.unbind(1)
+
         if self.use_focal_loss:
             scores = F.sigmoid(logits)
             scores, index = torch.topk(scores.flatten(1), self.num_top_queries, dim=-1)
@@ -104,9 +113,10 @@ class RTDETRPostProcessor(nn.Module):
             sco = sco[sco > threshold]
 
             lab = lab.cpu().numpy()
-            box = box.cpu().numpy()
             sco = sco.cpu().numpy()
 
+            box = self.clamp(box.cpu(), h.cpu(), w.cpu()).numpy()
+
             result = dict(labels=lab, boxes=box, scores=sco)
             results.append(result)
 

yomitoku/text_detector.py

@@ -6,7 +6,10 @@ import os
 from pydantic import conlist
 
 from .base import BaseModelCatalog, BaseModule, BaseSchema
-from .configs import TextDetectorDBNetConfig
+from .configs import (
+    TextDetectorDBNetConfig,
+    TextDetectorDBNetV2Config,
+)
 from .data.functions import (
     array_to_tensor,
     resize_shortest_edge,
@@ -25,6 +28,7 @@ class TextDetectorModelCatalog(BaseModelCatalog):
     def __init__(self):
         super().__init__()
         self.register("dbnet", TextDetectorDBNetConfig, DBNet)
+        self.register("dbnetv2", TextDetectorDBNetV2Config, DBNet)
 
 
 class TextDetectorSchema(BaseSchema):
@@ -43,7 +47,7 @@ class TextDetector(BaseModule):
 
     def __init__(
         self,
-        model_name="dbnet",
+        model_name="dbnetv2",
         path_cfg=None,
         device="cuda",
         visualize=False,

yomitoku/text_recognizer.py

@@ -7,7 +7,11 @@ import unicodedata
 from pydantic import conlist
 
 from .base import BaseModelCatalog, BaseModule, BaseSchema
-from .configs import TextRecognizerPARSeqConfig, TextRecognizerPARSeqSmallConfig
+from .configs import (
+    TextRecognizerPARSeqConfig,
+    TextRecognizerPARSeqSmallConfig,
+    TextRecognizerPARSeqV2Config,
+)
 from .data.dataset import ParseqDataset
 from .models import PARSeq
 from .postprocessor import ParseqTokenizer as Tokenizer
@@ -23,6 +27,7 @@ class TextRecognizerModelCatalog(BaseModelCatalog):
     def __init__(self):
         super().__init__()
         self.register("parseq", TextRecognizerPARSeqConfig, PARSeq)
+        self.register("parseqv2", TextRecognizerPARSeqV2Config, PARSeq)
         self.register("parseq-small", TextRecognizerPARSeqSmallConfig, PARSeq)
 
 
@@ -44,7 +49,7 @@ class TextRecognizer(BaseModule):
 
     def __init__(
         self,
-        model_name="parseq",
+        model_name="parseqv2",
         path_cfg=None,
         device="cuda",
         visualize=False,

{yomitoku-0.7.3.dist-info → yomitoku-0.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: yomitoku
-Version: 0.7.3
+Version: 0.8.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
@@ -37,7 +37,7 @@ Description-Content-Type: text/markdown
 YomiToku は日本語に特化した AI 文章画像解析エンジン(Document AI)です。画像内の文字の全文 OCR およびレイアウト解析機能を有しており、画像内の文字情報や図表を認識、抽出、変換します。
 
 - 🤖 日本語データセットで学習した 4 種類(文字位置の検知、文字列認識、レイアウト解析、表の構造認識)の AI モデルを搭載しています。4 種類のモデルはすべて独自に学習されたモデルで日本語文書に対して、高精度に推論可能です。
-- 🇯🇵 各モデルは日本語の文書画像に特化して学習されており、7000 文字を超える日本語文字の認識をサーポート、縦書きなど日本語特有のレイアウト構造の文書画像の解析も可能です。（日本語以外にも英語の文書に対しても対応しています）。
+- 🇯🇵 各モデルは日本語の文書画像に特化して学習されており、7000 文字を超える日本語文字の認識をサーポート、手書き文字、縦書きなど日本語特有のレイアウト構造の文書画像の解析も可能です。（日本語以外にも英語の文書に対しても対応しています）。
 - 📈 レイアウト解析、表の構造解析, 読み順推定機能により、文書画像のレイアウトの意味的構造を壊さずに情報を抽出することが可能です。
 - 📄 多様な出力形式をサポートしています。html やマークダウン、json、csv のいずれかのフォーマットに変換可能です。また、文書内に含まれる図表、画像の抽出の出力も可能です。
 - ⚡ GPU 環境で高速に動作し、効率的に文書の文字起こし解析が可能です。また、VRAM も 8GB 以内で動作し、ハイエンドな GPU を用意する必要はありません。
@@ -103,7 +103,6 @@ yomitoku --help
 **NOTE**
 
 - GPU での実行を推奨します。CPU を用いての推論向けに最適化されておらず、処理時間が長くなります。
-- 活字のみ識別をサポートしております。手書き文字に関しては、読み取れる場合もありますが、公式にはサポートしておりません。
 - Yomitoku は文書 OCR 向けに最適化されており、情景 OCR(看板など紙以外にプリントされた文字の読み取り)向けには最適化されていません。
 - AI-OCR の識別精度を高めるために、入力画像の解像度が重要です。低解像度画像では識別精度が低下します。最低でも画像の短辺を 720px 以上の画像で推論することをお勧めします。
 

{yomitoku-0.7.3.dist-info → yomitoku-0.8.0.dist-info}/RECORD

@@ -3,28 +3,31 @@ yomitoku/base.py,sha256=9U3sfe69O6vuO430JzzKQQNkgPsLM9WdLfOUUhp3Ljs,3878
 yomitoku/constants.py,sha256=zlW5QRc_u_F3C2RAgBFWyHJZexBnJT5N15GC-9d3iLo,686
 yomitoku/document_analyzer.py,sha256=wQMmXACDsDmyaxg2OnG9Og5Nx53WPUkQdUmgYtljACQ,16412
 yomitoku/layout_analyzer.py,sha256=VhNf1ZQFoozj6WUGk5ll1p2p1jk5X3j-JPcDbTAoSl4,1856
-yomitoku/layout_parser.py,sha256=V_mAkZxke1gwHfnxBFMTOJ8hnz2X_kfZu2lLiMd8cAs,7610
+yomitoku/layout_parser.py,sha256=0MgbCsD90srQdsxkGEL0TgKm4rkmGzsQYx0sjKQ03yc,7718
 yomitoku/ocr.py,sha256=JSTjkupcxHITQm6ERnzU7As0c3KWf8-oxc0AqNoWHXo,2272
 yomitoku/reading_order.py,sha256=OfhOS9ttPDoPSuHrIRKyOzG19GGeRufbuSKDqhsohh4,6404
 yomitoku/table_structure_recognizer.py,sha256=tHjex6deT_FjRK5ePz9bUXA_QIhgv_vYtK-ynm4ALxg,9625
-yomitoku/text_detector.py,sha256=XgqhtbNcJww2x3BrH8EFz45qC6kqPKCX9hsa-dzRoIA,4274
-yomitoku/text_recognizer.py,sha256=t95sbxve-E9VOCaU9CFGZIlk_a4my9KfFfr9tXws9As,5871
+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=WvQO9V5HzxxvRAIsGkrDl9OGrmaKsAbDBrg4ApCSy_c,10527
-yomitoku/configs/__init__.py,sha256=e1Alss5QJLZSNfD6zLEG6xu5vDQDw-4Jayiqq8bq52s,571
+yomitoku/cli/main.py,sha256=jQCSwHw4oOwLQjARvaIO1yoSjz-2Rdb9c3DNShLS5OE,12038
+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
 yomitoku/configs/cfg_table_structure_recognizer_rtdtrv2.py,sha256=o70GMHD8k-zeBeJtuhPS8x7vVB-ffucnJXeSyn-0AXo,2116
 yomitoku/configs/cfg_text_detector_dbnet.py,sha256=U9k48PON7haoOaytiELhbZRpv9RMiUm6nnfHmdxIa9Q,1153
+yomitoku/configs/cfg_text_detector_dbnet_v2.py,sha256=PzdV6-f75ba-KBEBcPxyo9STWQ6m5-1Rl3MFBLl2TSc,1148
 yomitoku/configs/cfg_text_recognizer_parseq.py,sha256=hpFs3nKqh4XdU3BZMTultegtLEGahEsCaZdjfKC_MO8,1247
 yomitoku/configs/cfg_text_recognizer_parseq_small.py,sha256=uCm_VC_G79IbZpOiK8fgYzAJ4b98H5pf328wyQomtfo,1259
+yomitoku/configs/cfg_text_recognizer_parseq_v2.py,sha256=GfHzbByOKjH21PRTxT8x_fU4r4Mda6F750Z8pjNeb8g,1249
 yomitoku/data/__init__.py,sha256=KAofFc9rk9ZdTKBjemu9RM8Vj9XnKbWC2MPZ2RWtOdE,82
 yomitoku/data/dataset.py,sha256=-I4f-FDtgsPnJ2MnXB7FtwihMW3koDaSI1OEoqKneIg,1014
-yomitoku/data/functions.py,sha256=7a_3xDKAQVdWfzQwFcdyJBojoyzUa3ePZOnG4pX1dpI,7532
-yomitoku/export/__init__.py,sha256=fkwOtqH0lh6eZQW5b4EMSjIH1FmWYLKKszahR-jQYSg,366
-yomitoku/export/export_csv.py,sha256=B234jlNeO4n5kQ_lwxxAZe_O2ipTbeDYlWU1zyyaVrw,3001
-yomitoku/export/export_html.py,sha256=pCLoxV10_SzRWmZlDnHuyfPFIuUGB3ZkqSdABVU7DTs,5038
-yomitoku/export/export_json.py,sha256=D6dD04gcPR5lmfHFVX-iGOYapsOVaJ_kH1Qhs6d2O0M,2035
-yomitoku/export/export_markdown.py,sha256=D1kX3X8odWa0pf4AFZ6gik5EKMKK7pgpQXaHHv6pWDI,4170
+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/models/__init__.py,sha256=Enxq9sjJWusZuxecTori8IQa8NEYKaiiptDluHX1avg,144
 yomitoku/models/dbnet_plus.py,sha256=jeWJZm0ihbxoJeAXBFK7uVIwoosx2IUNk7Ut5wRH0vA,7998
 yomitoku/models/parseq.py,sha256=psCPjP3eKjOFAUZJPQQhbD0nWEV5FeOZ0tTK27Rvvbw,8748
@@ -40,7 +43,7 @@ yomitoku/onnx/.gitkeep,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 yomitoku/postprocessor/__init__.py,sha256=W4vUuqBaFtH5dlSBIYgyaCroGLMjpV6RrNGIBQ8NFVw,243
 yomitoku/postprocessor/dbnet_postporcessor.py,sha256=o_y8b5REd2dFEdIpRcr6o-XBfOCHo9rBYGwokP_uhTc,4948
 yomitoku/postprocessor/parseq_tokenizer.py,sha256=e89_g_bc4Au3SchuxoJfJNATJTxFmVYetzXyAzPWm28,4315
-yomitoku/postprocessor/rtdetr_postprocessor.py,sha256=TCv1t1zCxg2rSirsLm4sXlaltGubH-roVdEqnUoRs-8,3905
+yomitoku/postprocessor/rtdetr_postprocessor.py,sha256=zp_PEAIl0-b7EJIWVZFrAaEUBSp9OgBVd1G-mP9R20E,4350
 yomitoku/resource/MPLUS1p-Medium.ttf,sha256=KLL1KkCumIBkgQtx1n4SffdaFuCNffThktEAbkB1OU8,1758908
 yomitoku/resource/charset.txt,sha256=sU91kSi-9Wk4733bCXy4j_UDmvcsj96sHOq1ppUJlOY,21672
 yomitoku/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -48,7 +51,7 @@ yomitoku/utils/graph.py,sha256=LKNB8ZhSQwOZMfeAimPMF5UCVVr2ZaUWoGDkz8z-uGU,456
 yomitoku/utils/logger.py,sha256=uOmtQDr0A0JD7wyFshedL08BiNrQorHnpktRXba8bjU,424
 yomitoku/utils/misc.py,sha256=FbwPLeIYYBvNf9wQh2RoEonTM5BF7_IwaEqmRsYHKA8,2673
 yomitoku/utils/visualizer.py,sha256=DjDwHiAu1iFRKh96H3Egq4vuI2s_-9dLCDeykhKi8jo,5251
-yomitoku-0.7.3.dist-info/METADATA,sha256=pHT4lQyl9cN4KbsOTooiJAaEEJqXhmAl9SVZKVaPkR0,8717
-yomitoku-0.7.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-yomitoku-0.7.3.dist-info/entry_points.txt,sha256=nFV3S11zgBNW0Qq_D0XQNg2R4lNXU_9XUFr6rdJoyF8,52
-yomitoku-0.7.3.dist-info/RECORD,,
+yomitoku-0.8.0.dist-info/METADATA,sha256=CH5KOT64Q8AMOaKkUbbd9rI1Zcd_dBk_OXd2GguC4f0,8555
+yomitoku-0.8.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+yomitoku-0.8.0.dist-info/entry_points.txt,sha256=nFV3S11zgBNW0Qq_D0XQNg2R4lNXU_9XUFr6rdJoyF8,52
+yomitoku-0.8.0.dist-info/RECORD,,