onnxtr 0.1.0__py3-none-any.whl

onnxtr/utils/geometry.py

@@ -0,0 +1,498 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+from copy import deepcopy
+from math import ceil
+from typing import List, Optional, Tuple, Union
+
+import cv2
+import numpy as np
+
+from .common_types import BoundingBox, Polygon4P
+
+__all__ = [
+    "bbox_to_polygon",
+    "polygon_to_bbox",
+    "resolve_enclosing_bbox",
+    "resolve_enclosing_rbbox",
+    "rotate_boxes",
+    "compute_expanded_shape",
+    "rotate_image",
+    "estimate_page_angle",
+    "convert_to_relative_coords",
+    "rotate_abs_geoms",
+    "extract_crops",
+    "extract_rcrops",
+    "shape_translate",
+]
+
+
+def bbox_to_polygon(bbox: BoundingBox) -> Polygon4P:
+    """Convert a bounding box to a polygon
+
+    Args:
+    ----
+        bbox: a bounding box
+
+    Returns:
+    -------
+        a polygon
+    """
+    return bbox[0], (bbox[1][0], bbox[0][1]), (bbox[0][0], bbox[1][1]), bbox[1]
+
+
+def polygon_to_bbox(polygon: Polygon4P) -> BoundingBox:
+    """Convert a polygon to a bounding box
+
+    Args:
+    ----
+        polygon: a polygon
+
+    Returns:
+    -------
+        a bounding box
+    """
+    x, y = zip(*polygon)
+    return (min(x), min(y)), (max(x), max(y))
+
+
+def shape_translate(data: np.ndarray, format: str) -> np.ndarray:
+    """Translate the shape of the input data to the desired format
+
+    Args:
+    ----
+        data: input data in shape (B, C, H, W) or (B, H, W, C) or (C, H, W) or (H, W, C)
+        format: target format ('BCHW', 'BHWC', 'CHW', or 'HWC')
+
+    Returns:
+    -------
+        the reshaped data
+    """
+    # Get the current shape
+    current_shape = data.shape
+
+    # Check the number of dimensions
+    num_dims = len(current_shape)
+
+    if num_dims != len(format):
+        return data
+
+    if format == "BCHW" and data.shape[1] in [1, 3]:
+        return data
+    elif format == "BHWC" and data.shape[-1] in [1, 3]:
+        return data
+    elif format == "CHW" and data.shape[0] in [1, 3]:
+        return data
+    elif format == "HWC" and data.shape[-1] in [1, 3]:
+        return data
+    elif format == "BCHW" and data.shape[1] not in [1, 3]:
+        return np.moveaxis(data, -1, 1)
+    elif format == "BHWC" and data.shape[-1] not in [1, 3]:
+        return np.moveaxis(data, 1, -1)
+    elif format == "CHW" and data.shape[0] not in [1, 3]:
+        return np.moveaxis(data, -1, 0)
+    elif format == "HWC" and data.shape[-1] not in [1, 3]:
+        return np.moveaxis(data, 0, -1)
+    else:
+        return data
+
+
+def resolve_enclosing_bbox(bboxes: Union[List[BoundingBox], np.ndarray]) -> Union[BoundingBox, np.ndarray]:
+    """Compute enclosing bbox either from:
+
+    Args:
+    ----
+        bboxes: boxes in one of the following formats:
+
+            - an array of boxes: (*, 5), where boxes have this shape:
+            (xmin, ymin, xmax, ymax, score)
+
+            - a list of BoundingBox
+
+    Returns:
+    -------
+        a (1, 5) array (enclosing boxarray), or a BoundingBox
+    """
+    if isinstance(bboxes, np.ndarray):
+        xmin, ymin, xmax, ymax, score = np.split(bboxes, 5, axis=1)
+        return np.array([xmin.min(), ymin.min(), xmax.max(), ymax.max(), score.mean()])
+    else:
+        x, y = zip(*[point for box in bboxes for point in box])
+        return (min(x), min(y)), (max(x), max(y))
+
+
+def resolve_enclosing_rbbox(rbboxes: List[np.ndarray], intermed_size: int = 1024) -> np.ndarray:
+    """Compute enclosing rotated bbox either from:
+
+    Args:
+    ----
+        rbboxes: boxes in one of the following formats:
+
+            - an array of boxes: (*, 5), where boxes have this shape:
+            (xmin, ymin, xmax, ymax, score)
+
+            - a list of BoundingBox
+        intermed_size: size of the intermediate image
+
+    Returns:
+    -------
+        a (1, 5) array (enclosing boxarray), or a BoundingBox
+    """
+    cloud: np.ndarray = np.concatenate(rbboxes, axis=0)
+    # Convert to absolute for minAreaRect
+    cloud *= intermed_size
+    rect = cv2.minAreaRect(cloud.astype(np.int32))
+    return cv2.boxPoints(rect) / intermed_size  # type: ignore[operator]
+
+
+def rotate_abs_points(points: np.ndarray, angle: float = 0.0) -> np.ndarray:
+    """Rotate points counter-clockwise.
+
+    Args:
+    ----
+        points: array of size (N, 2)
+        angle: angle between -90 and +90 degrees
+
+    Returns:
+    -------
+        Rotated points
+    """
+    angle_rad = angle * np.pi / 180.0  # compute radian angle for np functions
+    rotation_mat = np.array(
+        [[np.cos(angle_rad), -np.sin(angle_rad)], [np.sin(angle_rad), np.cos(angle_rad)]], dtype=points.dtype
+    )
+    return np.matmul(points, rotation_mat.T)
+
+
+def compute_expanded_shape(img_shape: Tuple[int, int], angle: float) -> Tuple[int, int]:
+    """Compute the shape of an expanded rotated image
+
+    Args:
+    ----
+        img_shape: the height and width of the image
+        angle: angle between -90 and +90 degrees
+
+    Returns:
+    -------
+        the height and width of the rotated image
+    """
+    points: np.ndarray = np.array([
+        [img_shape[1] / 2, img_shape[0] / 2],
+        [-img_shape[1] / 2, img_shape[0] / 2],
+    ])
+
+    rotated_points = rotate_abs_points(points, angle)
+
+    wh_shape = 2 * np.abs(rotated_points).max(axis=0)
+    return wh_shape[1], wh_shape[0]
+
+
+def rotate_abs_geoms(
+    geoms: np.ndarray,
+    angle: float,
+    img_shape: Tuple[int, int],
+    expand: bool = True,
+) -> np.ndarray:
+    """Rotate a batch of bounding boxes or polygons by an angle around the
+    image center.
+
+    Args:
+    ----
+        geoms: (N, 4) or (N, 4, 2) array of ABSOLUTE coordinate boxes
+        angle: anti-clockwise rotation angle in degrees
+        img_shape: the height and width of the image
+        expand: whether the image should be padded to avoid information loss
+
+    Returns:
+    -------
+        A batch of rotated polygons (N, 4, 2)
+    """
+    # Switch to polygons
+    polys = (
+        np.stack([geoms[:, [0, 1]], geoms[:, [2, 1]], geoms[:, [2, 3]], geoms[:, [0, 3]]], axis=1)
+        if geoms.ndim == 2
+        else geoms
+    )
+    polys = polys.astype(np.float32)
+
+    # Switch to image center as referential
+    polys[..., 0] -= img_shape[1] / 2
+    polys[..., 1] = img_shape[0] / 2 - polys[..., 1]
+
+    # Rotated them around image center
+    rotated_polys = rotate_abs_points(polys.reshape(-1, 2), angle).reshape(-1, 4, 2)
+    # Switch back to top-left corner as referential
+    target_shape = compute_expanded_shape(img_shape, angle) if expand else img_shape
+    # Clip coords to fit since there is no expansion
+    rotated_polys[..., 0] = (rotated_polys[..., 0] + target_shape[1] / 2).clip(0, target_shape[1])
+    rotated_polys[..., 1] = (target_shape[0] / 2 - rotated_polys[..., 1]).clip(0, target_shape[0])
+
+    return rotated_polys
+
+
+def remap_boxes(loc_preds: np.ndarray, orig_shape: Tuple[int, int], dest_shape: Tuple[int, int]) -> np.ndarray:
+    """Remaps a batch of rotated locpred (N, 4, 2) expressed for an origin_shape to a destination_shape.
+    This does not impact the absolute shape of the boxes, but allow to calculate the new relative RotatedBbox
+    coordinates after a resizing of the image.
+
+    Args:
+    ----
+        loc_preds: (N, 4, 2) array of RELATIVE loc_preds
+        orig_shape: shape of the origin image
+        dest_shape: shape of the destination image
+
+    Returns:
+    -------
+        A batch of rotated loc_preds (N, 4, 2) expressed in the destination referencial
+    """
+    if len(dest_shape) != 2:
+        raise ValueError(f"Mask length should be 2, was found at: {len(dest_shape)}")
+    if len(orig_shape) != 2:
+        raise ValueError(f"Image_shape length should be 2, was found at: {len(orig_shape)}")
+    orig_height, orig_width = orig_shape
+    dest_height, dest_width = dest_shape
+    mboxes = loc_preds.copy()
+    mboxes[:, :, 0] = ((loc_preds[:, :, 0] * orig_width) + (dest_width - orig_width) / 2) / dest_width
+    mboxes[:, :, 1] = ((loc_preds[:, :, 1] * orig_height) + (dest_height - orig_height) / 2) / dest_height
+
+    return mboxes
+
+
+def rotate_boxes(
+    loc_preds: np.ndarray,
+    angle: float,
+    orig_shape: Tuple[int, int],
+    min_angle: float = 1.0,
+    target_shape: Optional[Tuple[int, int]] = None,
+) -> np.ndarray:
+    """Rotate a batch of straight bounding boxes (xmin, ymin, xmax, ymax, c) or rotated bounding boxes
+    (4, 2) of an angle, if angle > min_angle, around the center of the page.
+    If target_shape is specified, the boxes are remapped to the target shape after the rotation. This
+    is done to remove the padding that is created by rotate_page(expand=True)
+
+    Args:
+    ----
+        loc_preds: (N, 5) or (N, 4, 2) array of RELATIVE boxes
+        angle: angle between -90 and +90 degrees
+        orig_shape: shape of the origin image
+        min_angle: minimum angle to rotate boxes
+        target_shape: shape of the destination image
+
+    Returns:
+    -------
+        A batch of rotated boxes (N, 4, 2): or a batch of straight bounding boxes
+    """
+    # Change format of the boxes to rotated boxes
+    _boxes = loc_preds.copy()
+    if _boxes.ndim == 2:
+        _boxes = np.stack(
+            [
+                _boxes[:, [0, 1]],
+                _boxes[:, [2, 1]],
+                _boxes[:, [2, 3]],
+                _boxes[:, [0, 3]],
+            ],
+            axis=1,
+        )
+    # If small angle, return boxes (no rotation)
+    if abs(angle) < min_angle or abs(angle) > 90 - min_angle:
+        return _boxes
+    # Compute rotation matrix
+    angle_rad = angle * np.pi / 180.0  # compute radian angle for np functions
+    rotation_mat = np.array(
+        [[np.cos(angle_rad), -np.sin(angle_rad)], [np.sin(angle_rad), np.cos(angle_rad)]], dtype=_boxes.dtype
+    )
+    # Rotate absolute points
+    points: np.ndarray = np.stack((_boxes[:, :, 0] * orig_shape[1], _boxes[:, :, 1] * orig_shape[0]), axis=-1)
+    image_center = (orig_shape[1] / 2, orig_shape[0] / 2)
+    rotated_points = image_center + np.matmul(points - image_center, rotation_mat)
+    rotated_boxes: np.ndarray = np.stack(
+        (rotated_points[:, :, 0] / orig_shape[1], rotated_points[:, :, 1] / orig_shape[0]), axis=-1
+    )
+
+    # Apply a mask if requested
+    if target_shape is not None:
+        rotated_boxes = remap_boxes(rotated_boxes, orig_shape=orig_shape, dest_shape=target_shape)
+
+    return rotated_boxes
+
+
+def rotate_image(
+    image: np.ndarray,
+    angle: float,
+    expand: bool = False,
+    preserve_origin_shape: bool = False,
+) -> np.ndarray:
+    """Rotate an image counterclockwise by an given angle.
+
+    Args:
+    ----
+        image: numpy tensor to rotate
+        angle: rotation angle in degrees, between -90 and +90
+        expand: whether the image should be padded before the rotation
+        preserve_origin_shape: if expand is set to True, resizes the final output to the original image size
+
+    Returns:
+    -------
+        Rotated array, padded by 0 by default.
+    """
+    # Compute the expanded padding
+    exp_img: np.ndarray
+    if expand:
+        exp_shape = compute_expanded_shape(image.shape[:2], angle)  # type: ignore[arg-type]
+        h_pad, w_pad = (
+            int(max(0, ceil(exp_shape[0] - image.shape[0]))),
+            int(max(0, ceil(exp_shape[1] - image.shape[1]))),
+        )
+        exp_img = np.pad(image, ((h_pad // 2, h_pad - h_pad // 2), (w_pad // 2, w_pad - w_pad // 2), (0, 0)))
+    else:
+        exp_img = image
+
+    height, width = exp_img.shape[:2]
+    rot_mat = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1.0)
+    rot_img = cv2.warpAffine(exp_img, rot_mat, (width, height))
+    if expand:
+        # Pad to get the same aspect ratio
+        if (image.shape[0] / image.shape[1]) != (rot_img.shape[0] / rot_img.shape[1]):
+            # Pad width
+            if (rot_img.shape[0] / rot_img.shape[1]) > (image.shape[0] / image.shape[1]):
+                h_pad, w_pad = 0, int(rot_img.shape[0] * image.shape[1] / image.shape[0] - rot_img.shape[1])
+            # Pad height
+            else:
+                h_pad, w_pad = int(rot_img.shape[1] * image.shape[0] / image.shape[1] - rot_img.shape[0]), 0
+            rot_img = np.pad(rot_img, ((h_pad // 2, h_pad - h_pad // 2), (w_pad // 2, w_pad - w_pad // 2), (0, 0)))
+        if preserve_origin_shape:
+            # rescale
+            rot_img = cv2.resize(rot_img, image.shape[:-1][::-1], interpolation=cv2.INTER_LINEAR)
+
+    return rot_img
+
+
+def estimate_page_angle(polys: np.ndarray) -> float:
+    """Takes a batch of rotated previously ORIENTED polys (N, 4, 2) (rectified by the classifier) and return the
+    estimated angle ccw in degrees
+    """
+    # Compute mean left points and mean right point with respect to the reading direction (oriented polygon)
+    xleft = polys[:, 0, 0] + polys[:, 3, 0]
+    yleft = polys[:, 0, 1] + polys[:, 3, 1]
+    xright = polys[:, 1, 0] + polys[:, 2, 0]
+    yright = polys[:, 1, 1] + polys[:, 2, 1]
+    with np.errstate(divide="raise", invalid="raise"):
+        try:
+            return float(
+                np.median(np.arctan((yleft - yright) / (xright - xleft)) * 180 / np.pi)  # Y axis from top to bottom!
+            )
+        except FloatingPointError:
+            return 0.0
+
+
+def convert_to_relative_coords(geoms: np.ndarray, img_shape: Tuple[int, int]) -> np.ndarray:
+    """Convert a geometry to relative coordinates
+
+    Args:
+    ----
+        geoms: a set of polygons of shape (N, 4, 2) or of straight boxes of shape (N, 4)
+        img_shape: the height and width of the image
+
+    Returns:
+    -------
+        the updated geometry
+    """
+    # Polygon
+    if geoms.ndim == 3 and geoms.shape[1:] == (4, 2):
+        polygons: np.ndarray = np.empty(geoms.shape, dtype=np.float32)
+        polygons[..., 0] = geoms[..., 0] / img_shape[1]
+        polygons[..., 1] = geoms[..., 1] / img_shape[0]
+        return polygons.clip(0, 1)
+    if geoms.ndim == 2 and geoms.shape[1] == 4:
+        boxes: np.ndarray = np.empty(geoms.shape, dtype=np.float32)
+        boxes[:, ::2] = geoms[:, ::2] / img_shape[1]
+        boxes[:, 1::2] = geoms[:, 1::2] / img_shape[0]
+        return boxes.clip(0, 1)
+
+    raise ValueError(f"invalid format for arg `geoms`: {geoms.shape}")
+
+
+def extract_crops(img: np.ndarray, boxes: np.ndarray, channels_last: bool = True) -> List[np.ndarray]:
+    """Created cropped images from list of bounding boxes
+
+    Args:
+    ----
+        img: input image
+        boxes: bounding boxes of shape (N, 4) where N is the number of boxes, and the relative
+            coordinates (xmin, ymin, xmax, ymax)
+        channels_last: whether the channel dimensions is the last one instead of the last one
+
+    Returns:
+    -------
+        list of cropped images
+    """
+    if boxes.shape[0] == 0:
+        return []
+    if boxes.shape[1] != 4:
+        raise AssertionError("boxes are expected to be relative and in order (xmin, ymin, xmax, ymax)")
+
+    # Project relative coordinates
+    _boxes = boxes.copy()
+    h, w = img.shape[:2] if channels_last else img.shape[-2:]
+    if not np.issubdtype(_boxes.dtype, np.integer):
+        _boxes[:, [0, 2]] *= w
+        _boxes[:, [1, 3]] *= h
+        _boxes = _boxes.round().astype(int)
+        # Add last index
+        _boxes[2:] += 1
+    if channels_last:
+        return deepcopy([img[box[1] : box[3], box[0] : box[2]] for box in _boxes])
+
+    return deepcopy([img[:, box[1] : box[3], box[0] : box[2]] for box in _boxes])
+
+
+def extract_rcrops(
+    img: np.ndarray, polys: np.ndarray, dtype=np.float32, channels_last: bool = True
+) -> List[np.ndarray]:
+    """Created cropped images from list of rotated bounding boxes
+
+    Args:
+    ----
+        img: input image
+        polys: bounding boxes of shape (N, 4, 2)
+        dtype: target data type of bounding boxes
+        channels_last: whether the channel dimensions is the last one instead of the last one
+
+    Returns:
+    -------
+        list of cropped images
+    """
+    if polys.shape[0] == 0:
+        return []
+    if polys.shape[1:] != (4, 2):
+        raise AssertionError("polys are expected to be quadrilateral, of shape (N, 4, 2)")
+
+    # Project relative coordinates
+    _boxes = polys.copy()
+    height, width = img.shape[:2] if channels_last else img.shape[-2:]
+    if not np.issubdtype(_boxes.dtype, np.integer):
+        _boxes[:, :, 0] *= width
+        _boxes[:, :, 1] *= height
+
+    src_pts = _boxes[:, :3].astype(np.float32)
+    # Preserve size
+    d1 = np.linalg.norm(src_pts[:, 0] - src_pts[:, 1], axis=-1)
+    d2 = np.linalg.norm(src_pts[:, 1] - src_pts[:, 2], axis=-1)
+    # (N, 3, 2)
+    dst_pts = np.zeros((_boxes.shape[0], 3, 2), dtype=dtype)
+    dst_pts[:, 1, 0] = dst_pts[:, 2, 0] = d1 - 1
+    dst_pts[:, 2, 1] = d2 - 1
+    # Use a warp transformation to extract the crop
+    crops = [
+        cv2.warpAffine(
+            img if channels_last else img.transpose(1, 2, 0),
+            # Transformation matrix
+            cv2.getAffineTransform(src_pts[idx], dst_pts[idx]),
+            (int(d1[idx]), int(d2[idx])),
+        )
+        for idx in range(_boxes.shape[0])
+    ]
+    return crops

onnxtr/utils/multithreading.py

@@ -0,0 +1,50 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+
+import multiprocessing as mp
+import os
+from multiprocessing.pool import ThreadPool
+from typing import Any, Callable, Iterable, Iterator, Optional
+
+from onnxtr.file_utils import ENV_VARS_TRUE_VALUES
+
+__all__ = ["multithread_exec"]
+
+
+def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: Optional[int] = None) -> Iterator[Any]:
+    """Execute a given function in parallel for each element of a given sequence
+
+    >>> from onnxtr.utils.multithreading import multithread_exec
+    >>> entries = [1, 4, 8]
+    >>> results = multithread_exec(lambda x: x ** 2, entries)
+
+    Args:
+    ----
+        func: function to be executed on each element of the iterable
+        seq: iterable
+        threads: number of workers to be used for multiprocessing
+
+    Returns:
+    -------
+        iterator of the function's results using the iterable as inputs
+
+    Notes:
+    -----
+        This function uses ThreadPool from multiprocessing package, which uses `/dev/shm` directory for shared memory.
+        If you do not have write permissions for this directory (if you run `onnxtr` on AWS Lambda for instance),
+        you might want to disable multiprocessing. To achieve that, set 'ONNXTR_MULTIPROCESSING_DISABLE' to 'TRUE'.
+    """
+    threads = threads if isinstance(threads, int) else min(16, mp.cpu_count())
+    # Single-thread
+    if threads < 2 or os.environ.get("ONNXTR_MULTIPROCESSING_DISABLE", "").upper() in ENV_VARS_TRUE_VALUES:
+        results = map(func, seq)
+    # Multi-threading
+    else:
+        with ThreadPool(threads) as tp:
+            # ThreadPool's map function returns a list, but seq could be of a different type
+            # That's why wrapping result in map to return iterator
+            results = map(lambda x: x, tp.map(func, seq))  # noqa: C417
+    return results

onnxtr/utils/reconstitution.py

@@ -0,0 +1,70 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+from typing import Any, Dict, Optional
+
+import numpy as np
+from anyascii import anyascii
+from PIL import Image, ImageDraw
+
+from .fonts import get_font
+
+__all__ = ["synthesize_page"]
+
+
+def synthesize_page(
+    page: Dict[str, Any],
+    draw_proba: bool = False,
+    font_family: Optional[str] = None,
+) -> np.ndarray:
+    """Draw a the content of the element page (OCR response) on a blank page.
+
+    Args:
+    ----
+        page: exported Page object to represent
+        draw_proba: if True, draw words in colors to represent confidence. Blue: p=1, red: p=0
+        font_size: size of the font, default font = 13
+        font_family: family of the font
+
+    Returns:
+    -------
+        the synthesized page
+    """
+    # Draw template
+    h, w = page["dimensions"]
+    response = 255 * np.ones((h, w, 3), dtype=np.int32)
+
+    # Draw each word
+    for block in page["blocks"]:
+        for line in block["lines"]:
+            for word in line["words"]:
+                # Get absolute word geometry
+                (xmin, ymin), (xmax, ymax) = word["geometry"]
+                xmin, xmax = int(round(w * xmin)), int(round(w * xmax))
+                ymin, ymax = int(round(h * ymin)), int(round(h * ymax))
+
+                # White drawing context adapted to font size, 0.75 factor to convert pts --> pix
+                font = get_font(font_family, int(0.75 * (ymax - ymin)))
+                img = Image.new("RGB", (xmax - xmin, ymax - ymin), color=(255, 255, 255))
+                d = ImageDraw.Draw(img)
+                # Draw in black the value of the word
+                try:
+                    d.text((0, 0), word["value"], font=font, fill=(0, 0, 0))
+                except UnicodeEncodeError:
+                    # When character cannot be encoded, use its anyascii version
+                    d.text((0, 0), anyascii(word["value"]), font=font, fill=(0, 0, 0))
+
+                # Colorize if draw_proba
+                if draw_proba:
+                    p = int(255 * word["confidence"])
+                    mask = np.where(np.array(img) == 0, 1, 0)
+                    proba: np.ndarray = np.array([255 - p, 0, p])
+                    color = mask * proba[np.newaxis, np.newaxis, :]
+                    white_mask = 255 * (1 - mask)
+                    img = color + white_mask
+
+                # Write to response page
+                response[ymin:ymax, xmin:xmax, :] = np.array(img)
+
+    return response

onnxtr/utils/repr.py

@@ -0,0 +1,64 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+# Adapted from https://github.com/pytorch/torch/blob/master/torch/nn/modules/module.py
+
+from typing import List
+
+__all__ = ["NestedObject"]
+
+
+def _addindent(s_, num_spaces):
+    s = s_.split("\n")
+    # don't do anything for single-line stuff
+    if len(s) == 1:
+        return s_
+    first = s.pop(0)
+    s = [(num_spaces * " ") + line for line in s]
+    s = "\n".join(s)
+    s = first + "\n" + s
+    return s
+
+
+class NestedObject:
+    """Base class for all nested objects in onnxtr"""
+
+    _children_names: List[str]
+
+    def extra_repr(self) -> str:
+        return ""
+
+    def __repr__(self):
+        # We treat the extra repr like the sub-object, one item per line
+        extra_lines = []
+        extra_repr = self.extra_repr()
+        # empty string will be split into list ['']
+        if extra_repr:
+            extra_lines = extra_repr.split("\n")
+        child_lines = []
+        if hasattr(self, "_children_names"):
+            for key in self._children_names:
+                child = getattr(self, key)
+                if isinstance(child, list) and len(child) > 0:
+                    child_str = ",\n".join([repr(subchild) for subchild in child])
+                    if len(child) > 1:
+                        child_str = _addindent(f"\n{child_str},", 2) + "\n"
+                    child_str = f"[{child_str}]"
+                else:
+                    child_str = repr(child)
+                child_str = _addindent(child_str, 2)
+                child_lines.append("(" + key + "): " + child_str)
+        lines = extra_lines + child_lines
+
+        main_str = self.__class__.__name__ + "("
+        if lines:
+            # simple one-liner info, which most builtin Modules will use
+            if len(extra_lines) == 1 and not child_lines:
+                main_str += extra_lines[0]
+            else:
+                main_str += "\n  " + "\n  ".join(lines) + "\n"
+
+        main_str += ")"
+        return main_str