Spaces:
Runtime error
Runtime error
File size: 1,868 Bytes
153628e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
# Copyright (C) 2021-2024, Mindee.
# 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, Tuple, Union
import numpy as np
import tensorflow as tf
from tensorflow import keras
from doctr.models.preprocessor import PreProcessor
from doctr.utils.repr import NestedObject
__all__ = ["DetectionPredictor"]
class DetectionPredictor(NestedObject):
"""Implements an object able to localize text elements in a document
Args:
----
pre_processor: transform inputs for easier batched model inference
model: core detection architecture
"""
_children_names: List[str] = ["pre_processor", "model"]
def __init__(
self,
pre_processor: PreProcessor,
model: keras.Model,
) -> None:
self.pre_processor = pre_processor
self.model = model
def __call__(
self,
pages: List[Union[np.ndarray, tf.Tensor]],
return_maps: bool = False,
**kwargs: Any,
) -> Union[List[Dict[str, np.ndarray]], Tuple[List[Dict[str, np.ndarray]], List[np.ndarray]]]:
# Dimension check
if any(page.ndim != 3 for page in pages):
raise ValueError("incorrect input shape: all pages are expected to be multi-channel 2D images.")
processed_batches = self.pre_processor(pages)
predicted_batches = [
self.model(batch, return_preds=True, return_model_output=True, training=False, **kwargs)
for batch in processed_batches
]
preds = [pred for batch in predicted_batches for pred in batch["preds"]]
if return_maps:
seg_maps = [pred.numpy() for batch in predicted_batches for pred in batch["out_map"]]
return preds, seg_maps
return preds
|