onnxtr 0.1.0__py3-none-any.whl

onnxtr/models/_utils.py

@@ -0,0 +1,141 @@
+# 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 math import floor
+from statistics import median_low
+from typing import List, Optional, Tuple
+
+import cv2
+import numpy as np
+from langdetect import LangDetectException, detect_langs
+
+__all__ = ["estimate_orientation", "get_language"]
+
+
+def get_max_width_length_ratio(contour: np.ndarray) -> float:
+    """Get the maximum shape ratio of a contour.
+
+    Args:
+    ----
+        contour: the contour from cv2.findContour
+
+    Returns:
+    -------
+        the maximum shape ratio
+    """
+    _, (w, h), _ = cv2.minAreaRect(contour)
+    return max(w / h, h / w)
+
+
+def estimate_orientation(img: np.ndarray, n_ct: int = 50, ratio_threshold_for_lines: float = 5) -> int:
+    """Estimate the angle of the general document orientation based on the
+     lines of the document and the assumption that they should be horizontal.
+
+    Args:
+    ----
+        img: the img or bitmap to analyze (H, W, C)
+        n_ct: the number of contours used for the orientation estimation
+        ratio_threshold_for_lines: this is the ratio w/h used to discriminates lines
+
+    Returns:
+    -------
+        the angle of the general document orientation
+    """
+    assert len(img.shape) == 3 and img.shape[-1] in [1, 3], f"Image shape {img.shape} not supported"
+    max_value = np.max(img)
+    min_value = np.min(img)
+    if max_value <= 1 and min_value >= 0 or (max_value <= 255 and min_value >= 0 and img.shape[-1] == 1):
+        thresh = img.astype(np.uint8)
+    if max_value <= 255 and min_value >= 0 and img.shape[-1] == 3:
+        gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        gray_img = cv2.medianBlur(gray_img, 5)
+        thresh = cv2.threshold(gray_img, thresh=0, maxval=255, type=cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]  # type: ignore[assignment]
+
+    # try to merge words in lines
+    (h, w) = img.shape[:2]
+    k_x = max(1, (floor(w / 100)))
+    k_y = max(1, (floor(h / 100)))
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (k_x, k_y))
+    thresh = cv2.dilate(thresh, kernel, iterations=1)  # type: ignore[assignment]
+
+    # extract contours
+    contours, _ = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+
+    # Sort contours
+    contours = sorted(contours, key=get_max_width_length_ratio, reverse=True)
+
+    angles = []
+    for contour in contours[:n_ct]:
+        _, (w, h), angle = cv2.minAreaRect(contour)
+        if w / h > ratio_threshold_for_lines:  # select only contours with ratio like lines
+            angles.append(angle)
+        elif w / h < 1 / ratio_threshold_for_lines:  # if lines are vertical, substract 90 degree
+            angles.append(angle - 90)
+
+    if len(angles) == 0:
+        return 0  # in case no angles is found
+    else:
+        median = -median_low(angles)
+        return round(median) if abs(median) != 0 else 0
+
+
+def rectify_crops(
+    crops: List[np.ndarray],
+    orientations: List[int],
+) -> List[np.ndarray]:
+    """Rotate each crop of the list according to the predicted orientation:
+    0: already straight, no rotation
+    1: 90 ccw, rotate 3 times ccw
+    2: 180, rotate 2 times ccw
+    3: 270 ccw, rotate 1 time ccw
+    """
+    # Inverse predictions (if angle of +90 is detected, rotate by -90)
+    orientations = [4 - pred if pred != 0 else 0 for pred in orientations]
+    return (
+        [crop if orientation == 0 else np.rot90(crop, orientation) for orientation, crop in zip(orientations, crops)]
+        if len(orientations) > 0
+        else []
+    )
+
+
+def rectify_loc_preds(
+    page_loc_preds: np.ndarray,
+    orientations: List[int],
+) -> Optional[np.ndarray]:
+    """Orient the quadrangle (Polygon4P) according to the predicted orientation,
+    so that the points are in this order: top L, top R, bot R, bot L if the crop is readable
+    """
+    return (
+        np.stack(
+            [
+                np.roll(page_loc_pred, orientation, axis=0)
+                for orientation, page_loc_pred in zip(orientations, page_loc_preds)
+            ],
+            axis=0,
+        )
+        if len(orientations) > 0
+        else None
+    )
+
+
+def get_language(text: str) -> Tuple[str, float]:
+    """Get languages of a text using langdetect model.
+    Get the language with the highest probability or no language if only a few words or a low probability
+
+    Args:
+    ----
+        text (str): text
+
+    Returns:
+    -------
+        The detected language in ISO 639 code and confidence score
+    """
+    try:
+        lang = detect_langs(text.lower())[0]
+    except LangDetectException:
+        return "unknown", 0.0
+    if len(text) <= 1 or (len(text) <= 5 and lang.prob <= 0.2):
+        return "unknown", 0.0
+    return lang.lang, lang.prob

onnxtr/models/builder.py

@@ -0,0 +1,355 @@
+# 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, List, Optional, Tuple
+
+import numpy as np
+from scipy.cluster.hierarchy import fclusterdata
+
+from onnxtr.io.elements import Block, Document, Line, Page, Word
+from onnxtr.utils.geometry import estimate_page_angle, resolve_enclosing_bbox, resolve_enclosing_rbbox, rotate_boxes
+from onnxtr.utils.repr import NestedObject
+
+__all__ = ["DocumentBuilder"]
+
+
+class DocumentBuilder(NestedObject):
+    """Implements a document builder
+
+    Args:
+    ----
+        resolve_lines: whether words should be automatically grouped into lines
+        resolve_blocks: whether lines should be automatically grouped into blocks
+        paragraph_break: relative length of the minimum space separating paragraphs
+        export_as_straight_boxes: if True, force straight boxes in the export (fit a rectangle
+            box to all rotated boxes). Else, keep the boxes format unchanged, no matter what it is.
+    """
+
+    def __init__(
+        self,
+        resolve_lines: bool = True,
+        resolve_blocks: bool = True,
+        paragraph_break: float = 0.035,
+        export_as_straight_boxes: bool = False,
+    ) -> None:
+        self.resolve_lines = resolve_lines
+        self.resolve_blocks = resolve_blocks
+        self.paragraph_break = paragraph_break
+        self.export_as_straight_boxes = export_as_straight_boxes
+
+    @staticmethod
+    def _sort_boxes(boxes: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """Sort bounding boxes from top to bottom, left to right
+
+        Args:
+        ----
+            boxes: bounding boxes of shape (N, 4) or (N, 4, 2) (in case of rotated bbox)
+
+        Returns:
+        -------
+            tuple: indices of ordered boxes of shape (N,), boxes
+                If straight boxes are passed tpo the function, boxes are unchanged
+                else: boxes returned are straight boxes fitted to the straightened rotated boxes
+                so that we fit the lines afterwards to the straigthened page
+        """
+        if boxes.ndim == 3:
+            boxes = rotate_boxes(
+                loc_preds=boxes,
+                angle=-estimate_page_angle(boxes),
+                orig_shape=(1024, 1024),
+                min_angle=5.0,
+            )
+            boxes = np.concatenate((boxes.min(1), boxes.max(1)), -1)
+        return (boxes[:, 0] + 2 * boxes[:, 3] / np.median(boxes[:, 3] - boxes[:, 1])).argsort(), boxes
+
+    def _resolve_sub_lines(self, boxes: np.ndarray, word_idcs: List[int]) -> List[List[int]]:
+        """Split a line in sub_lines
+
+        Args:
+        ----
+            boxes: bounding boxes of shape (N, 4)
+            word_idcs: list of indexes for the words of the line
+
+        Returns:
+        -------
+            A list of (sub-)lines computed from the original line (words)
+        """
+        lines = []
+        # Sort words horizontally
+        word_idcs = [word_idcs[idx] for idx in boxes[word_idcs, 0].argsort().tolist()]
+
+        # Eventually split line horizontally
+        if len(word_idcs) < 2:
+            lines.append(word_idcs)
+        else:
+            sub_line = [word_idcs[0]]
+            for i in word_idcs[1:]:
+                horiz_break = True
+
+                prev_box = boxes[sub_line[-1]]
+                # Compute distance between boxes
+                dist = boxes[i, 0] - prev_box[2]
+                # If distance between boxes is lower than paragraph break, same sub-line
+                if dist < self.paragraph_break:
+                    horiz_break = False
+
+                if horiz_break:
+                    lines.append(sub_line)
+                    sub_line = []
+
+                sub_line.append(i)
+            lines.append(sub_line)
+
+        return lines
+
+    def _resolve_lines(self, boxes: np.ndarray) -> List[List[int]]:
+        """Order boxes to group them in lines
+
+        Args:
+        ----
+            boxes: bounding boxes of shape (N, 4) or (N, 4, 2) in case of rotated bbox
+
+        Returns:
+        -------
+            nested list of box indices
+        """
+        # Sort boxes, and straighten the boxes if they are rotated
+        idxs, boxes = self._sort_boxes(boxes)
+
+        # Compute median for boxes heights
+        y_med = np.median(boxes[:, 3] - boxes[:, 1])
+
+        lines = []
+        words = [idxs[0]]  # Assign the top-left word to the first line
+        # Define a mean y-center for the line
+        y_center_sum = boxes[idxs[0]][[1, 3]].mean()
+
+        for idx in idxs[1:]:
+            vert_break = True
+
+            # Compute y_dist
+            y_dist = abs(boxes[idx][[1, 3]].mean() - y_center_sum / len(words))
+            # If y-center of the box is close enough to mean y-center of the line, same line
+            if y_dist < y_med / 2:
+                vert_break = False
+
+            if vert_break:
+                # Compute sub-lines (horizontal split)
+                lines.extend(self._resolve_sub_lines(boxes, words))
+                words = []
+                y_center_sum = 0
+
+            words.append(idx)
+            y_center_sum += boxes[idx][[1, 3]].mean()
+
+        # Use the remaining words to form the last(s) line(s)
+        if len(words) > 0:
+            # Compute sub-lines (horizontal split)
+            lines.extend(self._resolve_sub_lines(boxes, words))
+
+        return lines
+
+    @staticmethod
+    def _resolve_blocks(boxes: np.ndarray, lines: List[List[int]]) -> List[List[List[int]]]:
+        """Order lines to group them in blocks
+
+        Args:
+        ----
+            boxes: bounding boxes of shape (N, 4) or (N, 4, 2)
+            lines: list of lines, each line is a list of idx
+
+        Returns:
+        -------
+            nested list of box indices
+        """
+        # Resolve enclosing boxes of lines
+        if boxes.ndim == 3:
+            box_lines: np.ndarray = np.asarray([
+                resolve_enclosing_rbbox([tuple(boxes[idx, :, :]) for idx in line])  # type: ignore[misc]
+                for line in lines
+            ])
+        else:
+            _box_lines = [
+                resolve_enclosing_bbox([(tuple(boxes[idx, :2]), tuple(boxes[idx, 2:])) for idx in line])
+                for line in lines
+            ]
+            box_lines = np.asarray([(x1, y1, x2, y2) for ((x1, y1), (x2, y2)) in _box_lines])
+
+        # Compute geometrical features of lines to clusterize
+        # Clusterizing only with box centers yield to poor results for complex documents
+        if boxes.ndim == 3:
+            box_features: np.ndarray = np.stack(
+                (
+                    (box_lines[:, 0, 0] + box_lines[:, 0, 1]) / 2,
+                    (box_lines[:, 0, 0] + box_lines[:, 2, 0]) / 2,
+                    (box_lines[:, 0, 0] + box_lines[:, 2, 1]) / 2,
+                    (box_lines[:, 0, 1] + box_lines[:, 2, 1]) / 2,
+                    (box_lines[:, 0, 1] + box_lines[:, 2, 0]) / 2,
+                    (box_lines[:, 2, 0] + box_lines[:, 2, 1]) / 2,
+                ),
+                axis=-1,
+            )
+        else:
+            box_features = np.stack(
+                (
+                    (box_lines[:, 0] + box_lines[:, 3]) / 2,
+                    (box_lines[:, 1] + box_lines[:, 2]) / 2,
+                    (box_lines[:, 0] + box_lines[:, 2]) / 2,
+                    (box_lines[:, 1] + box_lines[:, 3]) / 2,
+                    box_lines[:, 0],
+                    box_lines[:, 1],
+                ),
+                axis=-1,
+            )
+        # Compute clusters
+        clusters = fclusterdata(box_features, t=0.1, depth=4, criterion="distance", metric="euclidean")
+
+        _blocks: Dict[int, List[int]] = {}
+        # Form clusters
+        for line_idx, cluster_idx in enumerate(clusters):
+            if cluster_idx in _blocks.keys():
+                _blocks[cluster_idx].append(line_idx)
+            else:
+                _blocks[cluster_idx] = [line_idx]
+
+        # Retrieve word-box level to return a fully nested structure
+        blocks = [[lines[idx] for idx in block] for block in _blocks.values()]
+
+        return blocks
+
+    def _build_blocks(
+        self,
+        boxes: np.ndarray,
+        word_preds: List[Tuple[str, float]],
+        crop_orientations: List[Dict[str, Any]],
+    ) -> List[Block]:
+        """Gather independent words in structured blocks
+
+        Args:
+        ----
+            boxes: bounding boxes of all detected words of the page, of shape (N, 5) or (N, 4, 2)
+            word_preds: list of all detected words of the page, of shape N
+            crop_orientations: list of dictoinaries containing
+                the general orientation (orientations + confidences) of the crops
+
+        Returns:
+        -------
+            list of block elements
+        """
+        if boxes.shape[0] != len(word_preds):
+            raise ValueError(f"Incompatible argument lengths: {boxes.shape[0]}, {len(word_preds)}")
+
+        if boxes.shape[0] == 0:
+            return []
+
+        # Decide whether we try to form lines
+        _boxes = boxes
+        if self.resolve_lines:
+            lines = self._resolve_lines(_boxes if _boxes.ndim == 3 else _boxes[:, :4])
+            # Decide whether we try to form blocks
+            if self.resolve_blocks and len(lines) > 1:
+                _blocks = self._resolve_blocks(_boxes if _boxes.ndim == 3 else _boxes[:, :4], lines)
+            else:
+                _blocks = [lines]
+        else:
+            # Sort bounding boxes, one line for all boxes, one block for the line
+            lines = [self._sort_boxes(_boxes if _boxes.ndim == 3 else _boxes[:, :4])[0]]  # type: ignore[list-item]
+            _blocks = [lines]
+
+        blocks = [
+            Block([
+                Line([
+                    Word(
+                        *word_preds[idx],
+                        tuple([tuple(pt) for pt in boxes[idx].tolist()]),  # type: ignore[arg-type]
+                        crop_orientations[idx],
+                    )
+                    if boxes.ndim == 3
+                    else Word(
+                        *word_preds[idx],
+                        ((boxes[idx, 0], boxes[idx, 1]), (boxes[idx, 2], boxes[idx, 3])),
+                        crop_orientations[idx],
+                    )
+                    for idx in line
+                ])
+                for line in lines
+            ])
+            for lines in _blocks
+        ]
+
+        return blocks
+
+    def extra_repr(self) -> str:
+        return (
+            f"resolve_lines={self.resolve_lines}, resolve_blocks={self.resolve_blocks}, "
+            f"paragraph_break={self.paragraph_break}, "
+            f"export_as_straight_boxes={self.export_as_straight_boxes}"
+        )
+
+    def __call__(
+        self,
+        pages: List[np.ndarray],
+        boxes: List[np.ndarray],
+        text_preds: List[List[Tuple[str, float]]],
+        page_shapes: List[Tuple[int, int]],
+        crop_orientations: List[Dict[str, Any]],
+        orientations: Optional[List[Dict[str, Any]]] = None,
+        languages: Optional[List[Dict[str, Any]]] = None,
+    ) -> Document:
+        """Re-arrange detected words into structured blocks
+
+        Args:
+        ----
+            pages: list of N elements, where each element represents the page image
+            boxes: list of N elements, where each element represents the localization predictions, of shape (*, 5)
+                or (*, 6) for all words for a given page
+            text_preds: list of N elements, where each element is the list of all word prediction (text + confidence)
+            page_shapes: shape of each page, of size N
+            crop_orientations: list of N elements, where each element is
+                a dictionary containing the general orientation (orientations + confidences) of the crops
+            orientations: optional, list of N elements,
+                where each element is a dictionary containing the orientation (orientation + confidence)
+            languages: optional, list of N elements,
+                where each element is a dictionary containing the language (language + confidence)
+
+        Returns:
+        -------
+            document object
+        """
+        if len(boxes) != len(text_preds) != len(crop_orientations) or len(boxes) != len(page_shapes) != len(
+            crop_orientations
+        ):
+            raise ValueError("All arguments are expected to be lists of the same size")
+
+        _orientations = (
+            orientations if isinstance(orientations, list) else [None] * len(boxes)  # type: ignore[list-item]
+        )
+        _languages = languages if isinstance(languages, list) else [None] * len(boxes)  # type: ignore[list-item]
+        if self.export_as_straight_boxes and len(boxes) > 0:
+            # If boxes are already straight OK, else fit a bounding rect
+            if boxes[0].ndim == 3:
+                # Iterate over pages and boxes
+                boxes = [np.concatenate((p_boxes.min(1), p_boxes.max(1)), 1) for p_boxes in boxes]
+
+        _pages = [
+            Page(
+                page,
+                self._build_blocks(
+                    page_boxes,
+                    word_preds,
+                    word_crop_orientations,
+                ),
+                _idx,
+                shape,
+                orientation,
+                language,
+            )
+            for page, _idx, shape, page_boxes, word_preds, word_crop_orientations, orientation, language in zip(
+                pages, range(len(boxes)), page_shapes, boxes, text_preds, crop_orientations, _orientations, _languages
+            )
+        ]
+
+        return Document(_pages)

onnxtr/models/classification/__init__.py

@@ -0,0 +1,2 @@
+from .models import *
+from .zoo import *

onnxtr/models/classification/models/__init__.py

@@ -0,0 +1 @@
+from .mobilenet import *

onnxtr/models/classification/models/mobilenet.py

@@ -0,0 +1,120 @@
+# 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.
+
+# Greatly inspired by https://github.com/pytorch/vision/blob/master/torchvision/models/mobilenetv3.py
+
+from copy import deepcopy
+from typing import Any, Dict, Optional
+
+import numpy as np
+
+from ...engine import Engine
+
+__all__ = [
+    "mobilenet_v3_small_crop_orientation",
+    "mobilenet_v3_small_page_orientation",
+]
+
+default_cfgs: Dict[str, Dict[str, Any]] = {
+    "mobilenet_v3_small_crop_orientation": {
+        "mean": (0.694, 0.695, 0.693),
+        "std": (0.299, 0.296, 0.301),
+        "input_shape": (3, 256, 256),
+        "classes": [0, -90, 180, 90],
+        "url": "https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/mobilenet_v3_small_crop_orientation-5620cf7e.onnx",
+    },
+    "mobilenet_v3_small_page_orientation": {
+        "mean": (0.694, 0.695, 0.693),
+        "std": (0.299, 0.296, 0.301),
+        "input_shape": (3, 512, 512),
+        "classes": [0, -90, 180, 90],
+        "url": "https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/mobilenet_v3_small_page_orientation-d3f76d79.onnx",
+    },
+}
+
+
+class MobileNetV3(Engine):
+    """MobileNetV3 Onnx loader
+
+    Args:
+    ----
+        model_path: path or url to onnx model file
+        cfg: configuration dictionary
+        **kwargs: additional arguments to be passed to `Engine`
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        cfg: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(url=model_path, **kwargs)
+        self.cfg = cfg
+
+    def __call__(
+        self,
+        x: np.ndarray,
+    ) -> np.ndarray:
+        return self.run(x)
+
+
+def _mobilenet_v3(
+    arch: str,
+    model_path: str,
+    **kwargs: Any,
+) -> MobileNetV3:
+    _cfg = deepcopy(default_cfgs[arch])
+    return MobileNetV3(model_path, cfg=_cfg, **kwargs)
+
+
+def mobilenet_v3_small_crop_orientation(
+    model_path: str = default_cfgs["mobilenet_v3_small_crop_orientation"]["url"], **kwargs: Any
+) -> MobileNetV3:
+    """MobileNetV3-Small architecture as described in
+    `"Searching for MobileNetV3",
+    <https://arxiv.org/pdf/1905.02244.pdf>`_.
+
+    >>> import numpy as np
+    >>> from onnxtr.models import mobilenet_v3_small_crop_orientation
+    >>> model = mobilenet_v3_small_crop_orientation()
+    >>> input_tensor = np.random.rand((1, 3, 256, 256))
+    >>> out = model(input_tensor)
+
+    Args:
+    ----
+        model_path: path to onnx model file, defaults to url in default_cfgs
+        **kwargs: keyword arguments of the MobileNetV3 architecture
+
+    Returns:
+    -------
+        MobileNetV3
+    """
+    return _mobilenet_v3("mobilenet_v3_small_crop_orientation", model_path, **kwargs)
+
+
+def mobilenet_v3_small_page_orientation(
+    model_path: str = default_cfgs["mobilenet_v3_small_page_orientation"]["url"], **kwargs: Any
+) -> MobileNetV3:
+    """MobileNetV3-Small architecture as described in
+    `"Searching for MobileNetV3",
+    <https://arxiv.org/pdf/1905.02244.pdf>`_.
+
+    >>> import numpy as np
+    >>> from onnxtr.models import mobilenet_v3_small_page_orientation
+    >>> model = mobilenet_v3_small_page_orientation()
+    >>> input_tensor = np.random.rand((1, 3, 512, 512))
+    >>> out = model(input_tensor)
+
+    Args:
+    ----
+        model_path: path to onnx model file, defaults to url in default_cfgs
+        **kwargs: keyword arguments of the MobileNetV3 architecture
+
+    Returns:
+    -------
+        MobileNetV3
+    """
+    return _mobilenet_v3("mobilenet_v3_small_page_orientation", model_path, **kwargs)

onnxtr/models/classification/predictor/__init__.py

@@ -0,0 +1 @@
+from .base import *

onnxtr/models/classification/predictor/base.py

@@ -0,0 +1,57 @@
+# 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, List, Union
+
+import numpy as np
+from scipy.special import softmax
+
+from onnxtr.models.preprocessor import PreProcessor
+from onnxtr.utils.repr import NestedObject
+
+__all__ = ["OrientationPredictor"]
+
+
+class OrientationPredictor(NestedObject):
+    """Implements an object able to detect the reading direction of a text box or a page.
+    4 possible orientations: 0, 90, 180, 270 (-90) degrees counter clockwise.
+
+    Args:
+    ----
+        pre_processor: transform inputs for easier batched model inference
+        model: core classification architecture (backbone + classification head)
+    """
+
+    _children_names: List[str] = ["pre_processor", "model"]
+
+    def __init__(
+        self,
+        pre_processor: PreProcessor,
+        model: Any,
+    ) -> None:
+        self.pre_processor = pre_processor
+        self.model = model
+
+    def __call__(
+        self,
+        inputs: List[np.ndarray],
+    ) -> List[Union[List[int], List[float]]]:
+        # Dimension check
+        if any(input.ndim != 3 for input in inputs):
+            raise ValueError("incorrect input shape: all inputs are expected to be multi-channel 2D images.")
+
+        processed_batches = self.pre_processor(inputs)
+        predicted_batches = [self.model(batch) for batch in processed_batches]
+
+        # confidence
+        probs = [np.max(softmax(batch, axis=1), axis=1) for batch in predicted_batches]
+        # Postprocess predictions
+        predicted_batches = [np.argmax(out_batch, axis=1) for out_batch in predicted_batches]
+
+        class_idxs = [int(pred) for batch in predicted_batches for pred in batch]
+        classes = [int(self.model.cfg["classes"][idx]) for idx in class_idxs]
+        confs = [round(float(p), 2) for prob in probs for p in prob]
+
+        return [class_idxs, classes, confs]