docling-ibm-models 0.1.0__py3-none-any.whl

docling_ibm_models/layoutmodel/layout_predictor.py

@@ -0,0 +1,171 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+import os
+from collections.abc import Iterable
+from typing import Union
+
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+
+MODEL_CHECKPOINT_FN = "model.pt"
+DEFAULT_NUM_THREADS = 4
+
+
+# Classes:
+CLASSES_MAP = {
+    0: "background",
+    1: "Caption",
+    2: "Footnote",
+    3: "Formula",
+    4: "List-item",
+    5: "Page-footer",
+    6: "Page-header",
+    7: "Picture",
+    8: "Section-header",
+    9: "Table",
+    10: "Text",
+    11: "Title",
+    12: "Document Index",
+    13: "Code",
+    14: "Checkbox-Selected",
+    15: "Checkbox-Unselected",
+    16: "Form",
+    17: "Key-Value Region",
+}
+
+
+class LayoutPredictor:
+    r"""
+    Document layout prediction using ONNX
+    """
+
+    def __init__(
+        self, artifact_path: str, num_threads: int = None, use_cpu_only: bool = False
+    ):
+        r"""
+        Provide the artifact path that contains the LayoutModel ONNX file
+
+        The number of threads is decided, in the following order, by:
+        1. The init method parameter `num_threads`, if it is set.
+        2. The envvar "OMP_NUM_THREADS", if it is set.
+        3. The default value DEFAULT_NUM_THREADS.
+
+        The execution provided is decided, in the following order:
+        1. If the init method parameter `cpu_only` is True or the envvar "USE_CPU_ONLY" is set,
+           it uses the "CPUExecutionProvider".
+        3. Otherwise if the "CUDAExecutionProvider" is present, use:
+            ["CUDAExecutionProvider", "CPUExecutionProvider"]:
+
+        Parameters
+        ----------
+        artifact_path: Path for the model ONNX file.
+        num_threads: (Optional) Number of threads to run the inference.
+        use_cpu_only: (Optional) If True, it forces CPU as the execution provider.
+
+        Raises
+        ------
+        FileNotFoundError when the model's ONNX file is missing
+        """
+        # Set basic params
+        self._threshold = 0.6  # Score threshold
+        self._image_size = 640
+        self._size = np.asarray([[self._image_size, self._image_size]], dtype=np.int64)
+
+        # Get env vars
+        self._use_cpu_only = use_cpu_only or ("USE_CPU_ONLY" in os.environ)
+        if num_threads is None:
+            num_threads = int(os.environ.get("OMP_NUM_THREADS", DEFAULT_NUM_THREADS))
+        self._num_threads = num_threads
+
+        # Decide the execution providers
+        if (
+            not self._use_cpu_only
+            and "CUDAExecutionProvider" in ort.get_available_providers()
+        ):
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        else:
+            providers = ["CPUExecutionProvider"]
+        self._providers = providers
+
+        # Model ONNX file
+        self._onnx_fn = os.path.join(artifact_path, MODEL_CHECKPOINT_FN)
+        if not os.path.isfile(self._onnx_fn):
+            raise FileNotFoundError("Missing ONNX file: {}".format(self._onnx_fn))
+
+        # ONNX options
+        self._options = ort.SessionOptions()
+        self._options.intra_op_num_threads = self._num_threads
+        self.sess = ort.InferenceSession(
+            self._onnx_fn,
+            sess_options=self._options,
+            providers=self._providers,
+        )
+
+    def info(self) -> dict:
+        r"""
+        Get information about the configuration of LayoutPredictor
+        """
+        info = {
+            "onnx_file": self._onnx_fn,
+            "intra_op_num_threads": self._num_threads,
+            "use_cpu_only": self._use_cpu_only,
+            "providers": self._providers,
+            "image_size": self._image_size,
+            "threshold": self._threshold,
+        }
+        return info
+
+    def predict(self, orig_img: Union[Image, np.array]) -> Iterable[dict]:
+        r"""
+        Predict bounding boxes for a given image.
+        The origin (0, 0) is the top-left corner and the predicted bbox coords are provided as:
+        [left, top, right, bottom]
+
+        Parameter
+        ---------
+        origin_img: Image to be predicted as a PIL Image object or numpy array.
+
+        Yield
+        -----
+        Bounding box as a dict with the keys: "label", "confidence", "l", "t", "r", "b"
+
+        Raises
+        ------
+        TypeError when the input image is not supported
+        """
+        # Convert image format
+        if isinstance(orig_img, Image.Image):
+            page_img = orig_img.convert("RGB")
+        elif isinstance(orig_img, np.ndarray):
+            page_img = Image.fromarray(orig_img).convert("RGB")
+        else:
+            raise TypeError("Not supported input image format")
+
+        w, h = page_img.size
+        page_img = page_img.resize((self._image_size, self._image_size))
+        page_data = np.array(page_img, dtype=np.uint8) / np.float32(255.0)
+        page_data = np.expand_dims(np.transpose(page_data, axes=[2, 0, 1]), axis=0)
+
+        # Predict
+        labels, boxes, scores = self.sess.run(
+            output_names=None,
+            input_feed={
+                "images": page_data,
+                "orig_target_sizes": self._size,
+            },
+        )
+
+        # Yield output
+        for label, box, score in zip(labels[0], boxes[0], scores[0]):
+            if score > self._threshold:
+                yield {
+                    "l": box[0] / self._image_size * w,
+                    "t": box[1] / self._image_size * h,
+                    "r": box[2] / self._image_size * w,
+                    "b": box[3] / self._image_size * h,
+                    "label": CLASSES_MAP[label],
+                    "confidence": score,
+                }

docling_ibm_models/tableformer/common.py

@@ -0,0 +1,200 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+import argparse
+import json
+import logging
+import os
+
+import torch
+
+import docling_ibm_models.tableformer.settings as s
+from docling_ibm_models.tableformer.models.common.base_model import BaseModel
+
+LOG_LEVEL = logging.DEBUG
+logger = s.get_custom_logger("common", LOG_LEVEL)
+
+
+def validate_config(config):
+    r"""
+    Validate the provided configuration file.
+    A ValueError exception will be thrown in case the config file is invalid
+
+    Parameters
+    ----------
+    config : dictionary
+        Configuration for the tablemodel
+
+    Returns
+    -------
+    bool : True on success
+    """
+    if "model" not in config:
+        return True
+    if "preparation" not in config:
+        return True
+    assert (
+        "max_tag_len" in config["preparation"]
+    ), "Config error: 'preparation.max_tag_len' parameter is missing"
+    if "seq_len" in config["model"]:
+        assert (
+            config["model"]["seq_len"] > 0
+        ), "Config error: 'model.seq_len' should be positive"
+        assert config["model"]["seq_len"] <= (
+            config["preparation"]["max_tag_len"] + 2
+        ), "Config error: 'model.seq_len' should be up to 'preparation.max_tag_len' + 2"
+
+    return True
+
+
+def parse_arguments():
+    r"""
+    Parse the input arguments
+    A ValueError exception will be thrown in case the config file is invalid
+    """
+    parser = argparse.ArgumentParser(description="Train the TableModel")
+    parser.add_argument(
+        "-c", "--config", required=True, default=None, help="configuration file (JSON)"
+    )
+    args = parser.parse_args()
+    config_filename = args.config
+
+    assert os.path.isfile(config_filename), "FAILURE: Config file not found."
+    return read_config(config_filename)
+
+
+def read_config(config_filename):
+    with open(config_filename, "r") as fd:
+        config = json.load(fd)
+
+    # Validate the config file
+    validate_config(config)
+
+    return config
+
+
+def safe_get_parameter(input_dict, index_path, default=None, required=False):
+    r"""
+    Safe get parameter from a nested dictionary.
+
+    Provide a nested dictionary (dictionary of dictionaries) and a list of indices:
+    - If the whole index path exists the value pointed by it is returned
+    - Otherwise the default value is returned.
+
+    Input:
+        input_dict: Data structure of nested dictionaries.
+        index_path: List with the indices path to follow inside the input_dict.
+        default: Default value to return if the indices path is broken.
+        required: If true a ValueError exception will be raised in case the parameter does not exist
+    Output:
+        The value pointed by the index path or "default".
+    """
+    if input_dict is None or index_path is None:
+        return default
+
+    d = input_dict
+    for i in index_path[:-1]:
+        if i not in d:
+            if required:
+                raise ValueError("Missing parameter: {}".format(i))
+            return default
+        d = d[i]
+
+    last_index = index_path[-1]
+    if last_index not in d:
+        if required:
+            raise ValueError("Missing parameter: {}".format(last_index))
+        return default
+
+    return d[last_index]
+
+
+def get_prepared_data_filename(prepared_data_part, dataset_name):
+    r"""
+    Build the full filename of the prepared data part
+
+    Parameters
+    ----------
+    prepared_data_part : string
+        Part of the prepared data
+    dataset_name : string
+        Name of the dataset
+
+    Returns
+    -------
+    string
+        The full filename for the prepared file
+    """
+    template = s.PREPARED_DATA_PARTS[prepared_data_part]
+    if "<POSTFIX>" in template:
+        template = template.replace("<POSTFIX>", dataset_name)
+    return template
+
+
+def create_dataset_and_model(config, purpose, fixed_padding=False):
+    r"""
+    Gets a model from configuration
+
+    Parameters
+    ---------
+    config : Dictionary
+        The configuration of the model
+    purpose : string
+        One of "train", "eval", "predict"
+    fixed_padding : bool
+        Parameter passed to the constructor of the DataLoader
+
+    Returns
+    -------
+    In case a Model cannot be initialized return None, None, None. Otherwise:
+
+    device : selected device
+    dataset : Instance of the DataLoader
+    model : Instance of the model
+    """
+    from docling_ibm_models.tableformer.data_management.tf_dataset import TFDataset
+
+    model_type = config["model"]["type"]
+    model = None
+
+    # Get env vars:
+    use_cpu_only = os.environ.get("USE_CPU_ONLY", False)
+    use_cuda_only = not use_cpu_only
+
+    # Use the cpu for the evaluation
+    device = "cpu"  # Default, run on CPU
+    num_gpus = torch.cuda.device_count()  # Check if GPU is available
+    if use_cuda_only:
+        device = "cuda:0" if num_gpus > 0 else "cpu"  # Run on first available GPU
+    else:
+        device = "cpu"
+
+    # Create the DataLoader
+    # loader = DataLoader(config, purpose, fixed_padding=fixed_padding)
+    dataset = TFDataset(config, purpose, fixed_padding=fixed_padding)
+    dataset.set_device(device)
+    dataset_val = None
+    if config["train"]["validation"] and purpose == "train":
+        dataset_val = TFDataset(config, "val", fixed_padding=fixed_padding)
+        dataset_val.set_device(device)
+    if model_type == "TableModel04_rs":
+        from docling_ibm_models.tableformer.models.table04_rs.tablemodel04_rs import (  # noqa: F401
+            TableModel04_rs,
+        )
+    # Find the model class and create an instance of it
+    for candidate in BaseModel.__subclasses__():
+        if candidate.__name__ == model_type:
+            init_data = dataset.get_init_data()
+            model = candidate(config, init_data, purpose, device)
+
+    if model is None:
+        logger.warn("Not found model: " + str(model_type))
+        return None, None, None
+
+    logger.info("Found model: " + str(model_type))
+
+    if purpose == s.PREDICT_PURPOSE:
+        return device, dataset, model
+    else:
+        return device, dataset, dataset_val, model