deepdoctection 0.30__py3-none-any.whl → 0.32__py3-none-any.whl

deepdoctection/extern/doctrocr.py

@@ -18,60 +18,60 @@
 """
 Deepdoctection wrappers for DocTr OCR text line detection and text recognition models
 """
+from __future__ import annotations
 
+import os
+from abc import ABC
 from pathlib import Path
-from typing import Any, List, Literal, Mapping, Optional, Tuple
+from typing import Any, List, Literal, Mapping, Optional, Tuple, Union
 from zipfile import ZipFile
 
+from lazy_imports import try_import
+
 from ..utils.detection_types import ImageType, Requirement
+from ..utils.error import DependencyError
 from ..utils.file_utils import (
-    doctr_available,
     get_doctr_requirement,
     get_pytorch_requirement,
     get_tensorflow_requirement,
     get_tf_addons_requirements,
     pytorch_available,
-    tf_addons_available,
     tf_available,
 )
 from ..utils.fs import load_json
-from ..utils.settings import LayoutType, ObjectTypes, TypeOrStr
-from .base import DetectionResult, ObjectDetector, PredictorBase, TextRecognizer
-from .pt.ptutils import set_torch_auto_device
+from ..utils.settings import LayoutType, ObjectTypes, PageType, TypeOrStr
+from ..utils.viz import viz_handler
+from .base import DetectionResult, ImageTransformer, ObjectDetector, PredictorBase, TextRecognizer
+from .pt.ptutils import get_torch_device
+from .tp.tfutils import get_tf_device
+
+with try_import() as pt_import_guard:
+    import torch
 
-if doctr_available() and ((tf_addons_available() and tf_available()) or pytorch_available()):
+with try_import() as tf_import_guard:
+    import tensorflow as tf  # type: ignore  # pylint: disable=E0401
+
+with try_import() as doctr_import_guard:
+    from doctr.models._utils import estimate_orientation
     from doctr.models.detection.predictor import DetectionPredictor  # pylint: disable=W0611
     from doctr.models.detection.zoo import detection_predictor
     from doctr.models.preprocessor import PreProcessor
     from doctr.models.recognition.predictor import RecognitionPredictor  # pylint: disable=W0611
     from doctr.models.recognition.zoo import ARCHS, recognition
 
-if pytorch_available():
-    import torch
-
-if tf_available():
-    import tensorflow as tf  # type: ignore  # pylint: disable=E0401
-
-
-def _set_device_str(device: Optional[str] = None) -> str:
-    if device is not None:
-        if tf_available():
-            device = "/" + device.replace("cuda", "gpu") + ":0"
-    elif pytorch_available():
-        device = set_torch_auto_device()
-    else:
-        device = "/gpu:0"  # we impose to install tensorflow-gpu because of Tensorpack models
-    return device
-
 
-def _load_model(path_weights: str, doctr_predictor: Any, device: str, lib: str) -> None:
-    if lib == "PT" and pytorch_available():
+def _load_model(
+    path_weights: str, doctr_predictor: Any, device: Union[torch.device, tf.device], lib: Literal["PT", "TF"]
+) -> None:
+    """Loading a model either in TF or PT. We only shift the model to the device when using PyTorch. The shift of
+    the model to the device in Tensorflow is done in the predict function."""
+    if lib == "PT":
         state_dict = torch.load(path_weights, map_location=device)
         for key in list(state_dict.keys()):
             state_dict["model." + key] = state_dict.pop(key)
         doctr_predictor.load_state_dict(state_dict)
         doctr_predictor.to(device)
-    elif lib == "TF" and tf_available():
+    elif lib == "TF":
         # Unzip the archive
         params_path = Path(path_weights).parent
         is_zip_path = path_weights.endswith(".zip")
@@ -83,20 +83,34 @@ def _load_model(path_weights: str, doctr_predictor: Any, device: str, lib: str)
             doctr_predictor.model.load_weights(path_weights)
 
 
-def doctr_predict_text_lines(np_img: ImageType, predictor: "DetectionPredictor", device: str) -> List[DetectionResult]:
+def auto_select_lib_for_doctr() -> Literal["PT", "TF"]:
+    """Auto select the DL library from environment variables"""
+    if os.environ.get("USE_TORCH"):
+        return "PT"
+    if os.environ.get("USE_TF"):
+        return "TF"
+    raise DependencyError("At least one of the env variables USE_TORCH or USE_TF must be set.")
+
+
+def doctr_predict_text_lines(
+    np_img: ImageType, predictor: DetectionPredictor, device: Union[torch.device, tf.device], lib: Literal["TF", "PT"]
+) -> List[DetectionResult]:
     """
     Generating text line DetectionResult based on Doctr DetectionPredictor.
 
     :param np_img: Image in np.array.
     :param predictor: `doctr.models.detection.predictor.DetectionPredictor`
     :param device: Will only be used in tensorflow settings. Either /gpu:0 or /cpu:0
+    :param lib: "TF" or "PT"
     :return: A list of text line detection results (without text).
     """
-    if tf_available() and device is not None:
-        with tf.device(device):
+    if lib == "TF":
+        with device:
             raw_output = predictor([np_img])
-    else:
+    elif lib == "PT":
         raw_output = predictor([np_img])
+    else:
+        raise DependencyError("Tensorflow or PyTorch must be installed.")
     detection_results = [
         DetectionResult(
             box=box[:4].tolist(), class_id=1, score=box[4], absolute_coords=False, class_name=LayoutType.word
@@ -107,7 +121,10 @@ def doctr_predict_text_lines(np_img: ImageType, predictor: "DetectionPredictor",
 
 
 def doctr_predict_text(
-    inputs: List[Tuple[str, ImageType]], predictor: "RecognitionPredictor", device: str
+    inputs: List[Tuple[str, ImageType]],
+    predictor: RecognitionPredictor,
+    device: Union[torch.device, tf.device],
+    lib: Literal["TF", "PT"],
 ) -> List[DetectionResult]:
     """
     Calls Doctr text recognition model on a batch of numpy arrays (text lines predicted from a text line detector) and
@@ -117,22 +134,46 @@ def doctr_predict_text(
                    text line
     :param predictor: `doctr.models.detection.predictor.RecognitionPredictor`
     :param device: Will only be used in tensorflow settings. Either /gpu:0 or /cpu:0
+    :param lib: "TF" or "PT"
     :return: A list of DetectionResult containing recognized text.
     """
 
     uuids, images = list(zip(*inputs))
-    if tf_available() and device is not None:
-        with tf.device(device):
+    if lib == "TF":
+        with device:
             raw_output = predictor(list(images))
-    else:
+    elif lib == "PT":
         raw_output = predictor(list(images))
+    else:
+        raise DependencyError("Tensorflow or PyTorch must be installed.")
     detection_results = [
         DetectionResult(score=output[1], text=output[0], uuid=uuid) for uuid, output in zip(uuids, raw_output)
     ]
     return detection_results
 
 
-class DoctrTextlineDetector(ObjectDetector):
+class DoctrTextlineDetectorMixin(ObjectDetector, ABC):
+    """Base class for Doctr textline detector. This class only implements the basic wrapper functions"""
+
+    def __init__(self, categories: Mapping[str, TypeOrStr], lib: Optional[Literal["PT", "TF"]] = None):
+        self.categories = categories  # type: ignore
+        self.lib = lib if lib is not None else self.auto_select_lib()
+
+    def possible_categories(self) -> List[ObjectTypes]:
+        return [LayoutType.word]
+
+    @staticmethod
+    def get_name(path_weights: str, architecture: str) -> str:
+        """Returns the name of the model"""
+        return f"doctr_{architecture}" + "_".join(Path(path_weights).parts[-2:])
+
+    @staticmethod
+    def auto_select_lib() -> Literal["PT", "TF"]:
+        """Auto select the DL library from the installed and from environment variables"""
+        return auto_select_lib_for_doctr()
+
+
+class DoctrTextlineDetector(DoctrTextlineDetectorMixin):
     """
     A deepdoctection wrapper of DocTr text line detector. We model text line detection as ObjectDetector
     and assume to use this detector in a ImageLayoutService.
@@ -165,8 +206,6 @@ class DoctrTextlineDetector(ObjectDetector):
 
                  for dp in df:
                      ...
-
-
     """
 
     def __init__(
@@ -174,21 +213,31 @@ class DoctrTextlineDetector(ObjectDetector):
         architecture: str,
         path_weights: str,
         categories: Mapping[str, TypeOrStr],
-        device: Optional[Literal["cpu", "cuda"]] = None,
-        lib: str = "TF",
+        device: Optional[Union[Literal["cpu", "cuda"], torch.device, tf.device]] = None,
+        lib: Optional[Literal["PT", "TF"]] = None,
     ) -> None:
-        self.lib = lib
-        self.name = "doctr_text_detector"
+        """
+        :param architecture: DocTR supports various text line detection models, e.g. "db_resnet50",
+        "db_mobilenet_v3_large". The full list can be found here:
+        https://github.com/mindee/doctr/blob/main/doctr/models/detection/zoo.py#L20
+        :param path_weights: Path to the weights of the model
+        :param categories: A dict with the model output label and value
+        :param device: "cpu" or "cuda" or any tf.device or torch.device. The device must be compatible with the dll
+        :param lib: "TF" or "PT" or None. If None, env variables USE_TENSORFLOW, USE_PYTORCH will be used.
+        """
+        super().__init__(categories, lib)
         self.architecture = architecture
         self.path_weights = path_weights
-        self.doctr_predictor = detection_predictor(
-            arch=self.architecture, pretrained=False, pretrained_backbone=False
-        )  # we will be loading the model
-        # later because there is no easy way in doctr to load a model by giving only a path to its weights
-        self.categories = categories  # type: ignore
-        self.device_input = device
-        self.device = _set_device_str(device)
-        self.load_model()
+
+        self.name = self.get_name(self.path_weights, self.architecture)
+        self.model_id = self.get_model_id()
+
+        if self.lib == "TF":
+            self.device = get_tf_device(device)
+        if self.lib == "PT":
+            self.device = get_torch_device(device)
+
+        self.doctr_predictor = self.get_wrapped_model(self.architecture, self.path_weights, self.device, self.lib)
 
     def predict(self, np_img: ImageType) -> List[DetectionResult]:
         """
@@ -197,26 +246,49 @@ class DoctrTextlineDetector(ObjectDetector):
         :param np_img: image as numpy array
         :return: A list of DetectionResult
         """
-        detection_results = doctr_predict_text_lines(np_img, self.doctr_predictor, self.device)
+        detection_results = doctr_predict_text_lines(np_img, self.doctr_predictor, self.device, self.lib)
         return detection_results
 
     @classmethod
     def get_requirements(cls) -> List[Requirement]:
-        if tf_available():
+        if os.environ.get("DD_USE_TF"):
             return [get_tensorflow_requirement(), get_doctr_requirement(), get_tf_addons_requirements()]
-        if pytorch_available():
+        if os.environ.get("DD_USE_TORCH"):
             return [get_pytorch_requirement(), get_doctr_requirement()]
         raise ModuleNotFoundError("Neither Tensorflow nor PyTorch has been installed. Cannot use DoctrTextlineDetector")
 
     def clone(self) -> PredictorBase:
-        return self.__class__(self.architecture, self.path_weights, self.categories, self.device_input, self.lib)
-
-    def possible_categories(self) -> List[ObjectTypes]:
-        return [LayoutType.word]
+        return self.__class__(self.architecture, self.path_weights, self.categories, self.device, self.lib)
 
-    def load_model(self) -> None:
+    @staticmethod
+    def load_model(
+        path_weights: str, doctr_predictor: Any, device: Union[torch.device, tf.device], lib: Literal["PT", "TF"]
+    ) -> None:
         """Loading model weights"""
-        _load_model(self.path_weights, self.doctr_predictor, self.device, self.lib)
+        _load_model(path_weights, doctr_predictor, device, lib)
+
+    @staticmethod
+    def get_wrapped_model(
+        architecture: str, path_weights: str, device: Union[torch.device, tf.device], lib: Literal["PT", "TF"]
+    ) -> Any:
+        """
+        Get the inner (wrapped) model.
+
+        :param architecture: DocTR supports various text line detection models, e.g. "db_resnet50",
+        "db_mobilenet_v3_large". The full list can be found here:
+        https://github.com/mindee/doctr/blob/main/doctr/models/detection/zoo.py#L20
+        :param path_weights: Path to the weights of the model
+        :param device: "cpu" or "cuda". Will default to "cuda" if the required hardware is available.
+        :param lib: "TF" or "PT" or None. If None, env variables USE_TENSORFLOW, USE_PYTORCH will be used. Make sure,
+                    these variables are set. If not, use
+
+                        deepdoctection.utils.env_info.auto_select_lib_and_device
+
+        :return: Inner model which is a "nn.Module" in PyTorch or a "tf.keras.Model" in Tensorflow
+        """
+        doctr_predictor = detection_predictor(arch=architecture, pretrained=False, pretrained_backbone=False)
+        DoctrTextlineDetector.load_model(path_weights, doctr_predictor, device, lib)
+        return doctr_predictor
 
 
 class DoctrTextRecognizer(TextRecognizer):
@@ -253,15 +325,14 @@ class DoctrTextRecognizer(TextRecognizer):
 
                  for dp in df:
                      ...
-
     """
 
     def __init__(
         self,
         architecture: str,
         path_weights: str,
-        device: Optional[Literal["cpu", "cuda"]] = None,
-        lib: str = "TF",
+        device: Optional[Union[Literal["cpu", "cuda"], torch.device, tf.device]] = None,
+        lib: Optional[Literal["PT", "TF"]] = None,
         path_config_json: Optional[str] = None,
     ) -> None:
         """
@@ -270,19 +341,30 @@ class DoctrTextRecognizer(TextRecognizer):
         https://github.com/mindee/doctr/blob/main/doctr/models/recognition/zoo.py#L16.
         :param path_weights: Path to the weights of the model
         :param device: "cpu" or "cuda". Will default to "cuda" if the required hardware is available.
-        :param lib: "TF" or "PT". Will default to "TF".
+        :param lib: "TF" or "PT" or None. If None, env variables USE_TENSORFLOW, USE_PYTORCH will be used.
         :param path_config_json: Path to a json file containing the configuration of the model. Useful, if you have
         a model trained on custom vocab.
         """
-        self.lib = lib
-        self.name = "doctr_text_recognizer"
+
+        self.lib = lib if lib is not None else self.auto_select_lib()
+
         self.architecture = architecture
         self.path_weights = path_weights
-        self.device_input = device
-        self.device = _set_device_str(device)
+
+        self.name = self.get_name(self.path_weights, self.architecture)
+        self.model_id = self.get_model_id()
+
+        if self.lib == "TF":
+            self.device = get_tf_device(device)
+        if self.lib == "PT":
+            self.device = get_torch_device(device)
+
         self.path_config_json = path_config_json
-        self.doctr_predictor = self.build_model()
-        self.load_model()
+        self.doctr_predictor = self.build_model(self.architecture, self.path_config_json)
+        self.load_model(self.path_weights, self.doctr_predictor, self.device, self.lib)
+        self.doctr_predictor = self.get_wrapped_model(
+            self.architecture, self.path_weights, self.device, self.lib, self.path_config_json
+        )
 
     def predict(self, images: List[Tuple[str, ImageType]]) -> List[DetectionResult]:
         """
@@ -292,7 +374,7 @@ class DoctrTextRecognizer(TextRecognizer):
         :return: A list of DetectionResult
         """
         if images:
-            return doctr_predict_text(images, self.doctr_predictor, self.device)
+            return doctr_predict_text(images, self.doctr_predictor, self.device, self.lib)
         return []
 
     @classmethod
@@ -304,21 +386,25 @@ class DoctrTextRecognizer(TextRecognizer):
         raise ModuleNotFoundError("Neither Tensorflow nor PyTorch has been installed. Cannot use DoctrTextRecognizer")
 
     def clone(self) -> PredictorBase:
-        return self.__class__(self.architecture, self.path_weights, self.device_input, self.lib)
+        return self.__class__(self.architecture, self.path_weights, self.device, self.lib)
 
-    def load_model(self) -> None:
+    @staticmethod
+    def load_model(
+        path_weights: str, doctr_predictor: Any, device: Union[torch.device, tf.device], lib: Literal["PT", "TF"]
+    ) -> None:
         """Loading model weights"""
-        _load_model(self.path_weights, self.doctr_predictor, self.device, self.lib)
+        _load_model(path_weights, doctr_predictor, device, lib)
 
-    def build_model(self) -> "RecognitionPredictor":
+    @staticmethod
+    def build_model(architecture: str, path_config_json: Optional[str] = None) -> "RecognitionPredictor":
         """Building the model"""
 
         # inspired and adapted from https://github.com/mindee/doctr/blob/main/doctr/models/recognition/zoo.py
         custom_configs = {}
         batch_size = 32
         recognition_configs = {}
-        if self.path_config_json:
-            custom_configs = load_json(self.path_config_json)
+        if path_config_json:
+            custom_configs = load_json(path_config_json)
             custom_configs.pop("arch", None)
             custom_configs.pop("url", None)
             custom_configs.pop("task", None)
@@ -327,18 +413,114 @@ class DoctrTextRecognizer(TextRecognizer):
             batch_size = custom_configs.pop("batch_size")
         recognition_configs["batch_size"] = batch_size
 
-        if isinstance(self.architecture, str):
-            if self.architecture not in ARCHS:
-                raise ValueError(f"unknown architecture '{self.architecture}'")
+        if isinstance(architecture, str):
+            if architecture not in ARCHS:
+                raise ValueError(f"unknown architecture '{architecture}'")
 
-            model = recognition.__dict__[self.architecture](pretrained=True, pretrained_backbone=True, **custom_configs)
+            model = recognition.__dict__[architecture](pretrained=True, pretrained_backbone=True, **custom_configs)
         else:
             if not isinstance(
-                self.architecture,
+                architecture,
                 (recognition.CRNN, recognition.SAR, recognition.MASTER, recognition.ViTSTR, recognition.PARSeq),
             ):
-                raise ValueError(f"unknown architecture: {type(self.architecture)}")
-            model = self.architecture
+                raise ValueError(f"unknown architecture: {type(architecture)}")
+            model = architecture
 
         input_shape = model.cfg["input_shape"][:2] if tf_available() else model.cfg["input_shape"][-2:]
         return RecognitionPredictor(PreProcessor(input_shape, preserve_aspect_ratio=True, **recognition_configs), model)
+
+    @staticmethod
+    def get_wrapped_model(
+        architecture: str,
+        path_weights: str,
+        device: Union[torch.device, tf.device],
+        lib: Literal["PT", "TF"],
+        path_config_json: Optional[str] = None,
+    ) -> Any:
+        """
+        Get the inner (wrapped) model.
+
+        :param architecture: DocTR supports various text recognition models, e.g. "crnn_vgg16_bn",
+        "crnn_mobilenet_v3_small". The full list can be found here:
+        https://github.com/mindee/doctr/blob/main/doctr/models/recognition/zoo.py#L16.
+        :param path_weights: Path to the weights of the model
+        :param device: "cpu" or "cuda". Will default to "cuda" if the required hardware is available.
+        :param lib: "TF" or "PT" or None. If None, env variables USE_TENSORFLOW, USE_PYTORCH will be used.
+        :param path_config_json: Path to a json file containing the configuration of the model. Useful, if you have
+        a model trained on custom vocab.
+        :return: Inner model which is a "nn.Module" in PyTorch or a "tf.keras.Model" in Tensorflow
+        """
+        doctr_predictor = DoctrTextRecognizer.build_model(architecture, path_config_json)
+        DoctrTextRecognizer.load_model(path_weights, doctr_predictor, device, lib)
+        return doctr_predictor
+
+    @staticmethod
+    def get_name(path_weights: str, architecture: str) -> str:
+        """Returns the name of the model"""
+        return f"doctr_{architecture}" + "_".join(Path(path_weights).parts[-2:])
+
+    @staticmethod
+    def auto_select_lib() -> Literal["PT", "TF"]:
+        """Auto select the DL library from the installed and from environment variables"""
+        return auto_select_lib_for_doctr()
+
+
+class DocTrRotationTransformer(ImageTransformer):
+    """
+    The `DocTrRotationTransformer` class is a specialized image transformer that is designed to handle image rotation
+    in the context of Optical Character Recognition (OCR) tasks. It inherits from the `ImageTransformer` base class and
+    implements methods for predicting and applying rotation transformations to images.
+
+    The `predict` method determines the angle of the rotated image using the `estimate_orientation` function from the
+    `doctr.models._utils` module. The `n_ct` and `ratio_threshold_for_lines` parameters for this function can be
+    configured when instantiating the class.
+
+    The `transform` method applies the predicted rotation to the image, effectively rotating the image backwards.
+    This method uses either the Pillow library or OpenCV for the rotation operation, depending on the configuration.
+
+    This class can be particularly useful in OCR tasks where the orientation of the text in the image matters.
+    The class also provides methods for cloning itself and for getting the requirements of the OCR system.
+
+    **Example:**
+                    transformer = DocTrRotationTransformer()
+                    detection_result = transformer.predict(np_img)
+                    rotated_image = transformer.transform(np_img, detection_result)
+    """
+
+    def __init__(self, number_contours: int = 50, ratio_threshold_for_lines: float = 5):
+        """
+
+        :param number_contours: the number of contours used for the orientation estimation
+        :param ratio_threshold_for_lines: this is the ratio w/h used to discriminates lines
+        """
+        self.number_contours = number_contours
+        self.ratio_threshold_for_lines = ratio_threshold_for_lines
+        self.name = "doctr_rotation_transformer"
+
+    def transform(self, np_img: ImageType, specification: DetectionResult) -> ImageType:
+        """
+        Applies the predicted rotation to the image, effectively rotating the image backwards.
+        This method uses either the Pillow library or OpenCV for the rotation operation, depending on the configuration.
+
+        :param np_img: The input image as a numpy array.
+        :param specification: A `DetectionResult` object containing the predicted rotation angle.
+        :return: The rotated image as a numpy array.
+        """
+        return viz_handler.rotate_image(np_img, specification.angle)  # type: ignore
+
+    def predict(self, np_img: ImageType) -> DetectionResult:
+        angle = estimate_orientation(np_img, self.number_contours, self.ratio_threshold_for_lines)
+        if angle < 0:
+            angle += 360
+        return DetectionResult(angle=round(angle, 2))
+
+    @classmethod
+    def get_requirements(cls) -> List[Requirement]:
+        return [get_doctr_requirement()]
+
+    def clone(self) -> PredictorBase:
+        return self.__class__(self.number_contours, self.ratio_threshold_for_lines)
+
+    @staticmethod
+    def possible_category() -> PageType:
+        return PageType.angle

deepdoctection/extern/fastlang.py

@@ -18,18 +18,47 @@
 """
 Deepdoctection wrappers for fasttext language detection models
 """
+from abc import ABC
 from copy import copy
-from typing import List, Mapping
+from pathlib import Path
+from typing import Any, List, Mapping, Tuple, Union
 
-from ..utils.file_utils import Requirement, fasttext_available, get_fasttext_requirement
-from ..utils.settings import TypeOrStr
+from lazy_imports import try_import
+
+from ..utils.file_utils import Requirement, get_fasttext_requirement
+from ..utils.settings import TypeOrStr, get_type
 from .base import DetectionResult, LanguageDetector, PredictorBase
 
-if fasttext_available():
+with try_import() as import_guard:
     from fasttext import load_model  # type: ignore
 
 
-class FasttextLangDetector(LanguageDetector):
+class FasttextLangDetectorMixin(LanguageDetector, ABC):
+    """
+    Base class for Fasttext language detection implementation. This class only implements the basic wrapper functions.
+    """
+
+    def __init__(self, categories: Mapping[str, TypeOrStr]) -> None:
+        """
+        :param categories: A dict with the model output label and value. We use as convention the ISO 639-2 language
+        """
+        self.categories = copy({idx: get_type(cat) for idx, cat in categories.items()})
+
+    def output_to_detection_result(self, output: Union[Tuple[Any, Any]]) -> DetectionResult:
+        """
+        Generating `DetectionResult` from model output
+        :param output: FastText model output
+        :return: `DetectionResult` filled with `text` and `score`
+        """
+        return DetectionResult(text=self.categories[output[0][0]], score=output[1][0])
+
+    @staticmethod
+    def get_name(path_weights: str) -> str:
+        """Returns the name of the model"""
+        return "fasttext_" + "_".join(Path(path_weights).parts[-2:])
+
+
+class FasttextLangDetector(FasttextLangDetectorMixin):
     """
     Fasttext language detector wrapper. Two models provided in the fasttext library can be used to identify languages.
     The background to the models can be found in the works:
@@ -57,15 +86,18 @@ class FasttextLangDetector(LanguageDetector):
         :param categories: A dict with the model output label and value. We use as convention the ISO 639-2 language
                            code.
         """
+        super().__init__(categories)
 
-        self.name = "fasttest_lang_detector"
         self.path_weights = path_weights
-        self.model = load_model(self.path_weights)
-        self.categories = copy(categories)  # type: ignore
+
+        self.name = self.get_name(self.path_weights)
+        self.model_id = self.get_model_id()
+
+        self.model = self.get_wrapped_model(self.path_weights)
 
     def predict(self, text_string: str) -> DetectionResult:
         output = self.model.predict(text_string)
-        return DetectionResult(text=self.categories[output[0][0]], score=output[1][0])
+        return self.output_to_detection_result(output)
 
     @classmethod
     def get_requirements(cls) -> List[Requirement]:
@@ -73,3 +105,11 @@ class FasttextLangDetector(LanguageDetector):
 
     def clone(self) -> PredictorBase:
         return self.__class__(self.path_weights, self.categories)
+
+    @staticmethod
+    def get_wrapped_model(path_weights: str) -> Any:
+        """
+        Get the wrapped model
+        :param path_weights: path to model weights
+        """
+        return load_model(path_weights)