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

yomitoku/__init__.py

@@ -0,0 +1,20 @@
+from importlib.metadata import version
+
+from .document_analyzer import DocumentAnalyzer
+from .layout_analyzer import LayoutAnalyzer
+from .layout_parser import LayoutParser
+from .ocr import OCR
+from .table_structure_recognizer import TableStructureRecognizer
+from .text_detector import TextDetector
+from .text_recognizer import TextRecognizer
+
+__all__ = [
+    "OCR",
+    "LayoutParser",
+    "TableStructureRecognizer",
+    "TextDetector",
+    "TextRecognizer",
+    "LayoutAnalyzer",
+    "DocumentAnalyzer",
+]
+__version__ = version(__package__)

yomitoku/base.py

@@ -0,0 +1,136 @@
+import time
+from pathlib import Path
+from typing import Union
+
+import torch
+from omegaconf import OmegaConf
+from pydantic import BaseModel, Extra
+
+from .export import export_json
+from .utils.logger import set_logger
+
+logger = set_logger(__name__, "INFO")
+
+
+def load_yaml_config(path_config: str):
+    path_config = Path(path_config)
+    if not path_config.exists():
+        raise FileNotFoundError(f"Config file not found: {path_config}")
+
+    with open(path_config, "r") as file:
+        yaml_config = OmegaConf.load(file)
+    return yaml_config
+
+
+def load_config(
+    default_config,
+    path_config: Union[str, None] = None,
+):
+    cfg = OmegaConf.structured(default_config)
+    if path_config is not None:
+        yaml_config = load_yaml_config(path_config)
+        cfg = OmegaConf.merge(cfg, yaml_config)
+    return cfg
+
+
+def observer(cls, func):
+    def wrapper(*args, **kwargs):
+        try:
+            start = time.time()
+            result = func(*args, **kwargs)
+            elapsed = time.time() - start
+            logger.info(f"{cls.__name__} {func.__name__} elapsed_time: {elapsed}")
+        except Exception as e:
+            logger.error(f"Error occurred in {cls.__name__} {func.__name__}: {e}")
+            raise e
+        return result
+
+    return wrapper
+
+
+class BaseSchema(BaseModel):
+    class Config:
+        extra = Extra.forbid
+        validate_assignment = True
+
+    def to_json(self, out_path: str, **kwargs):
+        export_json(self, out_path, **kwargs)
+
+
+class BaseModule:
+    model_catalog = None
+
+    def __init__(self):
+        if self.model_catalog is None:
+            raise NotImplementedError
+
+        if not issubclass(self.model_catalog.__class__, BaseModelCatalog):
+            raise ValueError(
+                f"{self.model_catalog.__class__} is not SubClass BaseModelCatalog."
+            )
+
+        if len(self.model_catalog.list_model()) == 0:
+            raise ValueError("No model is registered.")
+
+    def __new__(cls, *args, **kwds):
+        logger.info(f"Initialize {cls.__name__}")
+        cls.__call__ = observer(cls, cls.__call__)
+        return super().__new__(cls)
+
+    def load_model(self, name, path_cfg, from_pretrained=True):
+        default_cfg, Net = self.model_catalog.get(name)
+        self._cfg = load_config(default_cfg, path_cfg)
+        if from_pretrained:
+            self.model = Net.from_pretrained(self._cfg.hf_hub_repo, cfg=self._cfg)
+        else:
+            self.model = Net(cfg=self._cfg)
+
+    def save_config(self, path_cfg):
+        OmegaConf.save(self._cfg, path_cfg)
+
+    def log_config(self):
+        logger.info(OmegaConf.to_yaml(self._cfg))
+
+    @classmethod
+    def catalog(cls):
+        display = ""
+        for model in cls.model_catalog.list_model():
+            display += f"{model} "
+        logger.info(f"{cls.__name__} Implemented Models")
+        logger.info(display)
+
+    @property
+    def device(self):
+        return self._device
+
+    @device.setter
+    def device(self, device):
+        if "cuda" in device:
+            if torch.cuda.is_available():
+                self._device = torch.device(device)
+            else:
+                self._device = torch.device("cpu")
+                logger.warning("CUDA is not available. Use CPU instead.")
+        else:
+            self._device = torch.device("cpu")
+
+
+class BaseModelCatalog:
+    def __init__(self):
+        self.catalog = {}
+
+    def get(self, model_name):
+        model_name = model_name.lower()
+        if model_name in self.catalog:
+            return self.catalog[model_name]
+
+        raise ValueError(f"Unknown model: {model_name}")
+
+    def register(self, model_name, config, model):
+        if model_name in self.catalog:
+            raise ValueError(f"{model_name} is already registered.")
+
+        self.catalog[model_name] = (config, model)
+
+    def list_model(self):
+        return list(self.catalog.keys())

yomitoku/cli/main.py

@@ -0,0 +1,230 @@
+import argparse
+import os
+from pathlib import Path
+
+import cv2
+import time
+
+from ..constants import SUPPORT_OUTPUT_FORMAT
+from ..data.functions import load_image, load_pdf
+from ..document_analyzer import DocumentAnalyzer
+from ..utils.logger import set_logger
+
+logger = set_logger(__name__, "INFO")
+
+
+def process_single_file(args, analyzer, path, format):
+    if path.suffix[1:].lower() in ["pdf"]:
+        imgs = load_pdf(path)
+    else:
+        imgs = [load_image(path)]
+
+    for page, img in enumerate(imgs):
+        results, ocr, layout = analyzer(img)
+
+        dirname = path.parent.name
+        filename = path.stem
+
+        if ocr is not None:
+            out_path = os.path.join(
+                args.outdir, f"{dirname}_{filename}_p{page+1}_ocr.jpg"
+            )
+
+            cv2.imwrite(out_path, ocr)
+            logger.info(f"Output file: {out_path}")
+
+        if layout is not None:
+            out_path = os.path.join(
+                args.outdir, f"{dirname}_{filename}_p{page+1}_layout.jpg"
+            )
+
+            cv2.imwrite(out_path, layout)
+            logger.info(f"Output file: {out_path}")
+
+        out_path = os.path.join(args.outdir, f"{dirname}_{filename}_p{page+1}.{format}")
+
+        if format == "json":
+            results.to_json(
+                out_path,
+                ignore_line_break=args.ignore_line_break,
+            )
+        elif format == "csv":
+            results.to_csv(
+                out_path,
+                ignore_line_break=args.ignore_line_break,
+            )
+        elif format == "html":
+            results.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,
+            )
+        elif format == "md":
+            results.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,
+            )
+
+        logger.info(f"Output file: {out_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "arg1",
+        type=str,
+        help="path of target image file or directory",
+    )
+    parser.add_argument(
+        "-f",
+        "--format",
+        type=str,
+        default="json",
+        help="output format type (json or csv or html or md)",
+    )
+    parser.add_argument(
+        "-v",
+        "--vis",
+        action="store_true",
+        help="if set, visualize the result",
+    )
+    parser.add_argument(
+        "-o",
+        "--outdir",
+        type=str,
+        default="results",
+        help="output directory",
+    )
+    parser.add_argument(
+        "-d",
+        "--device",
+        type=str,
+        default="cuda",
+        help="device to use",
+    )
+    parser.add_argument(
+        "--td_cfg",
+        type=str,
+        default=None,
+        help="path of text detector config file",
+    )
+    parser.add_argument(
+        "--tr_cfg",
+        type=str,
+        default=None,
+        help="path of text recognizer config file",
+    )
+    parser.add_argument(
+        "--lp_cfg",
+        type=str,
+        default=None,
+        help="path of layout parser config file",
+    )
+    parser.add_argument(
+        "--tsr_cfg",
+        type=str,
+        default=None,
+        help="path of table structure recognizer config file",
+    )
+    parser.add_argument(
+        "--ignore_line_break",
+        action="store_true",
+        help="if set, ignore line break in the output",
+    )
+    parser.add_argument(
+        "--figure",
+        action="store_true",
+        help="if set, export figure in the output",
+    )
+    parser.add_argument(
+        "--figure_letter",
+        action="store_true",
+        help="if set, export letter within figure in the output",
+    )
+    parser.add_argument(
+        "--figure_width",
+        type=int,
+        default=200,
+        help="width of figure image in the output",
+    )
+    parser.add_argument(
+        "--figure_dir",
+        type=str,
+        default="figures",
+        help="directory to save figure images",
+    )
+
+    args = parser.parse_args()
+
+    path = Path(args.arg1)
+    if not path.exists():
+        raise FileNotFoundError(f"File not found: {args.arg1}")
+
+    format = args.format.lower()
+    if format not in SUPPORT_OUTPUT_FORMAT:
+        raise ValueError(
+            f"Invalid output format: {args.format}. Supported formats are {SUPPORT_OUTPUT_FORMAT}"
+        )
+
+    if format == "markdown":
+        format = "md"
+
+    configs = {
+        "ocr": {
+            "text_detector": {
+                "path_cfg": args.td_cfg,
+            },
+            "text_recognizer": {
+                "path_cfg": args.tr_cfg,
+            },
+        },
+        "layout_analyzer": {
+            "layout_parser": {
+                "path_cfg": args.lp_cfg,
+            },
+            "table_structure_recognizer": {
+                "path_cfg": args.tsr_cfg,
+            },
+        },
+    }
+
+    analyzer = DocumentAnalyzer(
+        configs=configs,
+        visualize=args.vis,
+        device=args.device,
+    )
+
+    os.makedirs(args.outdir, exist_ok=True)
+    logger.info(f"Output directory: {args.outdir}")
+
+    if path.is_dir():
+        all_files = [f for f in path.rglob("*") if f.is_file()]
+        for f in all_files:
+            try:
+                start = time.time()
+                file_path = Path(f)
+                logger.info(f"Processing file: {file_path}")
+                process_single_file(args, analyzer, file_path, format)
+                end = time.time()
+                logger.info(f"Total Processing time: {end-start:.2f} sec")
+            except Exception:
+                continue
+    else:
+        start = time.time()
+        logger.info(f"Processing file: {path}")
+        process_single_file(args, analyzer, path, format)
+        end = time.time()
+        logger.info(f"Total Processing time: {end-start:.2f} sec")
+
+
+if __name__ == "__main__":
+    main()

yomitoku/configs/__init__.py

@@ -0,0 +1,13 @@
+from .cfg_layout_parser_rtdtrv2 import LayoutParserRTDETRv2Config
+from .cfg_table_structure_recognizer_rtdtrv2 import (
+    TableStructureRecognizerRTDETRv2Config,
+)
+from .cfg_text_detector_dbnet import TextDetectorDBNetConfig
+from .cfg_text_recognizer_parseq import TextRecognizerPARSeqConfig
+
+__all__ = [
+    "TextDetectorDBNetConfig",
+    "TextRecognizerPARSeqConfig",
+    "LayoutParserRTDETRv2Config",
+    "TableStructureRecognizerRTDETRv2Config",
+]

yomitoku/configs/cfg_layout_parser_rtdtrv2.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 LayoutParserRTDETRv2Config:
+    hf_hub_repo: str = "KotaroKinoshita/yomitoku-layout-parser-rtdtrv2-open-beta"
+    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_table_structure_recognizer_rtdtrv2.py

@@ -0,0 +1,80 @@
+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 = 3
+    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  # 1.0 0.4
+    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 TableStructureRecognizerRTDETRv2Config:
+    hf_hub_repo: str = (
+        "KotaroKinoshita/yomitoku-table-structure-recognizer-rtdtrv2-open-beta"
+    )
+    thresh_score: float = 0.4
+    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: [
+            "row",
+            "col",
+            "span",
+        ]
+    )

yomitoku/configs/cfg_text_detector_dbnet.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.2
+    box_thresh: float = 0.5
+    max_candidates: int = 1500
+    unclip_ratio: float = 2.0
+
+
+@dataclass
+class Visualize:
+    color: List[int] = field(default_factory=lambda: [0, 255, 0])
+    heatmap: bool = False
+
+
+@dataclass
+class TextDetectorDBNetConfig:
+    hf_hub_repo: str = "KotaroKinoshita/yomitoku-text-detector-dbnet-open-beta"
+    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.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 TextRecognizerPARSeqConfig:
+    hf_hub_repo: str = "KotaroKinoshita/yomitoku-text-recognizer-parseq-open-beta"
+    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/constants.py

@@ -0,0 +1,32 @@
+import os
+
+ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+SUPPORT_OUTPUT_FORMAT = ["json", "csv", "html", "markdown", "md"]
+SUPPORT_INPUT_FORMAT = ["jpg", "jpeg", "png", "bmp", "tiff", "tif", "pdf"]
+MIN_IMAGE_SIZE = 32
+WARNING_IMAGE_SIZE = 720
+
+PALETTE = [
+    [255, 0, 0],
+    [0, 255, 0],
+    [0, 0, 255],
+    [255, 255, 0],
+    [0, 255, 255],
+    [255, 0, 255],
+    [128, 0, 0],
+    [0, 128, 0],
+    [0, 0, 128],
+    [255, 128, 0],
+    [0, 255, 128],
+    [128, 0, 255],
+    [128, 255, 0],
+    [0, 128, 255],
+    [255, 0, 128],
+    [255, 128, 128],
+    [128, 255, 128],
+    [128, 128, 255],
+    [255, 255, 128],
+    [255, 128, 255],
+    [128, 255, 255],
+    [128, 128, 128],
+]