deepdoctection 0.31__py3-none-any.whl → 0.33__py3-none-any.whl

deepdoctection/extern/doctrocr.py

@@ -18,32 +18,39 @@
 """
 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, Literal, Mapping, Optional, Union
 from zipfile import ZipFile
 
-from ..utils.detection_types import ImageType, Requirement
-from ..utils.env_info import get_device
+from lazy_imports import try_import
+
+from ..utils.env_info import ENV_VARS_TRUE
 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, PageType, TypeOrStr
+from ..utils.types import PathLikeOrStr, PixelValues, Requirement
 from ..utils.viz import viz_handler
-from .base import DetectionResult, ImageTransformer, ObjectDetector, PredictorBase, TextRecognizer
-from .pt.ptutils import set_torch_auto_device
+from .base import DetectionResult, ImageTransformer, ModelCategories, ObjectDetector, 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
@@ -51,70 +58,72 @@ if doctr_available() and ((tf_addons_available() and tf_available()) or pytorch_
     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: Literal["PT", "TF"]) -> None:
-    if lib == "PT" and pytorch_available():
-        state_dict = torch.load(path_weights, map_location=device)
+def _get_doctr_requirements() -> list[Requirement]:
+    if os.environ.get("DD_USE_TF", "0") in ENV_VARS_TRUE:
+        return [get_tensorflow_requirement(), get_doctr_requirement(), get_tf_addons_requirements()]
+    if os.environ.get("DD_USE_TORCH", "0") in ENV_VARS_TRUE:
+        return [get_pytorch_requirement(), get_doctr_requirement()]
+    raise ModuleNotFoundError("Neither Tensorflow nor PyTorch has been installed. Cannot use DoctrTextRecognizer")
+
+
+def _load_model(
+    path_weights: PathLikeOrStr,
+    doctr_predictor: Union[DetectionPredictor, RecognitionPredictor],
+    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(os.fspath(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")
+        is_zip_path = os.fspath(path_weights).endswith(".zip")
         if is_zip_path:
             with ZipFile(path_weights, "r") as file:
                 file.extractall(path=params_path)
                 doctr_predictor.model.load_weights(params_path / "weights")
         else:
-            doctr_predictor.model.load_weights(path_weights)
+            doctr_predictor.model.load_weights(os.fspath(path_weights))
 
 
 def auto_select_lib_for_doctr() -> Literal["PT", "TF"]:
-    """Auto select the DL library from the installed and from environment variables"""
-    if tf_available() and os.environ.get("USE_TF", os.environ.get("USE_TENSORFLOW", False)):
-        os.environ["USE_TF"] = "TRUE"
-        return "TF"
-    if pytorch_available() and os.environ.get("USE_TORCH", os.environ.get("USE_PYTORCH", False)):
+    """Auto select the DL library from environment variables"""
+    if os.environ.get("USE_TORCH", "0") in ENV_VARS_TRUE:
         return "PT"
-    raise DependencyError("Neither Tensorflow nor PyTorch has been installed. Cannot use DoctrTextlineDetector")
+    if os.environ.get("USE_TF", "0") in ENV_VARS_TRUE:
+        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: str) -> List[DetectionResult]:
+def doctr_predict_text_lines(
+    np_img: PixelValues, 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
+            box=box[:4].tolist(), class_id=1, score=box[4], absolute_coords=False, class_name=LayoutType.WORD
         )
         for box in raw_output[0]["words"]
     ]
@@ -122,8 +131,11 @@ def doctr_predict_text_lines(np_img: ImageType, predictor: "DetectionPredictor",
 
 
 def doctr_predict_text(
-    inputs: List[Tuple[str, ImageType]], predictor: "RecognitionPredictor", device: str
-) -> List[DetectionResult]:
+    inputs: list[tuple[str, PixelValues]],
+    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
     returns the recognized text as DetectionResult
@@ -132,15 +144,18 @@ 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)
     ]
@@ -150,15 +165,15 @@ def doctr_predict_text(
 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
+    def __init__(self, categories: Mapping[int, TypeOrStr], lib: Optional[Literal["PT", "TF"]] = None):
+        self.categories = ModelCategories(init_categories=categories)
         self.lib = lib if lib is not None else self.auto_select_lib()
 
-    def possible_categories(self) -> List[ObjectTypes]:
-        return [LayoutType.word]
+    def get_category_names(self) -> tuple[ObjectTypes, ...]:
+        return self.categories.get_categories(as_dict=False)
 
     @staticmethod
-    def get_name(path_weights: str, architecture: str) -> str:
+    def get_name(path_weights: PathLikeOrStr, architecture: str) -> str:
         """Returns the name of the model"""
         return f"doctr_{architecture}" + "_".join(Path(path_weights).parts[-2:])
 
@@ -206,9 +221,9 @@ class DoctrTextlineDetector(DoctrTextlineDetectorMixin):
     def __init__(
         self,
         architecture: str,
-        path_weights: str,
-        categories: Mapping[str, TypeOrStr],
-        device: Optional[Literal["cpu", "cuda"]] = None,
+        path_weights: PathLikeOrStr,
+        categories: Mapping[int, TypeOrStr],
+        device: Optional[Union[Literal["cpu", "cuda"], torch.device, tf.device]] = None,
         lib: Optional[Literal["PT", "TF"]] = None,
     ) -> None:
         """
@@ -217,58 +232,54 @@ class DoctrTextlineDetector(DoctrTextlineDetectorMixin):
         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". Will default to "cuda" if the required hardware is available.
+        :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.path_weights = Path(path_weights)
 
         self.name = self.get_name(self.path_weights, self.architecture)
         self.model_id = self.get_model_id()
 
-        if device is None:
-            if self.lib == "TF":
-                device = "cuda" if tf.test.is_gpu_available() else "cpu"
-            elif self.lib == "PT":
-                auto_device = get_device(False)
-                device = "cpu" if auto_device == "mps" else auto_device
-            else:
-                raise DependencyError("Cannot select device automatically. Please set the device manually.")
+        if self.lib == "TF":
+            self.device = get_tf_device(device)
+        if self.lib == "PT":
+            self.device = get_torch_device(device)
 
-        self.device_input = device
-        self.device = _set_device_str(device)
-        self.doctr_predictor = self.get_wrapped_model(self.architecture, self.path_weights, self.device_input, self.lib)
+        self.doctr_predictor = self.get_wrapped_model(self.architecture, self.path_weights, self.device, self.lib)
 
-    def predict(self, np_img: ImageType) -> List[DetectionResult]:
+    def predict(self, np_img: PixelValues) -> list[DetectionResult]:
         """
         Prediction per image.
 
         :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)
-        return detection_results
+        return doctr_predict_text_lines(np_img, self.doctr_predictor, self.device, self.lib)
 
     @classmethod
-    def get_requirements(cls) -> List[Requirement]:
-        if tf_available():
-            return [get_tensorflow_requirement(), get_doctr_requirement(), get_tf_addons_requirements()]
-        if pytorch_available():
-            return [get_pytorch_requirement(), get_doctr_requirement()]
-        raise ModuleNotFoundError("Neither Tensorflow nor PyTorch has been installed. Cannot use DoctrTextlineDetector")
+    def get_requirements(cls) -> list[Requirement]:
+        return _get_doctr_requirements()
 
-    def clone(self) -> PredictorBase:
-        return self.__class__(self.architecture, self.path_weights, self.categories, self.device_input, self.lib)
+    def clone(self) -> DoctrTextlineDetector:
+        return self.__class__(
+            self.architecture, self.path_weights, self.categories.get_categories(), self.device, self.lib
+        )
 
     @staticmethod
-    def load_model(path_weights: str, doctr_predictor: Any, device: str, lib: Literal["PT", "TF"]) -> None:
+    def load_model(
+        path_weights: PathLikeOrStr,
+        doctr_predictor: DetectionPredictor,
+        device: Union[torch.device, tf.device],
+        lib: Literal["PT", "TF"],
+    ) -> None:
         """Loading model weights"""
         _load_model(path_weights, doctr_predictor, device, lib)
 
     @staticmethod
     def get_wrapped_model(
-        architecture: str, path_weights: str, device: Literal["cpu", "cuda"], lib: Literal["PT", "TF"]
+        architecture: str, path_weights: PathLikeOrStr, device: Union[torch.device, tf.device], lib: Literal["PT", "TF"]
     ) -> Any:
         """
         Get the inner (wrapped) model.
@@ -286,10 +297,12 @@ class DoctrTextlineDetector(DoctrTextlineDetectorMixin):
         :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)
-        device_str = _set_device_str(device)
-        DoctrTextlineDetector.load_model(path_weights, doctr_predictor, device_str, lib)
+        DoctrTextlineDetector.load_model(path_weights, doctr_predictor, device, lib)
         return doctr_predictor
 
+    def clear_model(self) -> None:
+        self.doctr_predictor = None
+
 
 class DoctrTextRecognizer(TextRecognizer):
     """
@@ -325,16 +338,15 @@ class DoctrTextRecognizer(TextRecognizer):
 
                  for dp in df:
                      ...
-
     """
 
     def __init__(
         self,
         architecture: str,
-        path_weights: str,
-        device: Optional[Literal["cpu", "cuda"]] = None,
+        path_weights: PathLikeOrStr,
+        device: Optional[Union[Literal["cpu", "cuda"], torch.device, tf.device]] = None,
         lib: Optional[Literal["PT", "TF"]] = None,
-        path_config_json: Optional[str] = None,
+        path_config_json: Optional[PathLikeOrStr] = None,
     ) -> None:
         """
         :param architecture: DocTR supports various text recognition models, e.g. "crnn_vgg16_bn",
@@ -350,30 +362,24 @@ class DoctrTextRecognizer(TextRecognizer):
         self.lib = lib if lib is not None else self.auto_select_lib()
 
         self.architecture = architecture
-        self.path_weights = path_weights
+        self.path_weights = Path(path_weights)
 
         self.name = self.get_name(self.path_weights, self.architecture)
         self.model_id = self.get_model_id()
 
-        if device is None:
-            if self.lib == "TF":
-                device = "cuda" if tf.test.is_gpu_available() else "cpu"
-            if self.lib == "PT":
-                auto_device = get_device(False)
-                device = "cpu" if auto_device == "mps" else auto_device
-            else:
-                raise DependencyError("Cannot select device automatically. Please set the device manually.")
-
-        self.device_input = device
-        self.device = _set_device_str(device)
+        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.architecture, self.path_config_json)
+        self.doctr_predictor = self.build_model(self.architecture, self.lib, 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_input, self.lib, self.path_config_json
+            self.architecture, self.path_weights, self.device, self.lib, self.path_config_json
         )
 
-    def predict(self, images: List[Tuple[str, ImageType]]) -> List[DetectionResult]:
+    def predict(self, images: list[tuple[str, PixelValues]]) -> list[DetectionResult]:
         """
         Prediction on a batch of text lines
 
@@ -381,27 +387,30 @@ 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
-    def get_requirements(cls) -> List[Requirement]:
-        if tf_available():
-            return [get_tensorflow_requirement(), get_doctr_requirement(), get_tf_addons_requirements()]
-        if pytorch_available():
-            return [get_pytorch_requirement(), get_doctr_requirement()]
-        raise ModuleNotFoundError("Neither Tensorflow nor PyTorch has been installed. Cannot use DoctrTextRecognizer")
+    def get_requirements(cls) -> list[Requirement]:
+        return _get_doctr_requirements()
 
-    def clone(self) -> PredictorBase:
-        return self.__class__(self.architecture, self.path_weights, self.device_input, self.lib)
+    def clone(self) -> DoctrTextRecognizer:
+        return self.__class__(self.architecture, self.path_weights, self.device, self.lib)
 
     @staticmethod
-    def load_model(path_weights: str, doctr_predictor: Any, device: str, lib: Literal["PT", "TF"]) -> None:
+    def load_model(
+        path_weights: PathLikeOrStr,
+        doctr_predictor: RecognitionPredictor,
+        device: Union[torch.device, tf.device],
+        lib: Literal["PT", "TF"],
+    ) -> None:
         """Loading model weights"""
         _load_model(path_weights, doctr_predictor, device, lib)
 
     @staticmethod
-    def build_model(architecture: str, path_config_json: Optional[str] = None) -> "RecognitionPredictor":
+    def build_model(
+        architecture: str, lib: Literal["TF", "PT"], path_config_json: Optional[PathLikeOrStr] = None
+    ) -> RecognitionPredictor:
         """Building the model"""
 
         # inspired and adapted from https://github.com/mindee/doctr/blob/main/doctr/models/recognition/zoo.py
@@ -424,6 +433,7 @@ class DoctrTextRecognizer(TextRecognizer):
 
             model = recognition.__dict__[architecture](pretrained=True, pretrained_backbone=True, **custom_configs)
         else:
+            # This is not documented, but you can also directly pass the model class to architecture
             if not isinstance(
                 architecture,
                 (recognition.CRNN, recognition.SAR, recognition.MASTER, recognition.ViTSTR, recognition.PARSeq),
@@ -431,16 +441,16 @@ class DoctrTextRecognizer(TextRecognizer):
                 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:]
+        input_shape = model.cfg["input_shape"][:2] if lib == "TF" 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: Literal["cpu", "cuda"],
+        path_weights: PathLikeOrStr,
+        device: Union[torch.device, tf.device],
         lib: Literal["PT", "TF"],
-        path_config_json: Optional[str] = None,
+        path_config_json: Optional[PathLikeOrStr] = None,
     ) -> Any:
         """
         Get the inner (wrapped) model.
@@ -455,13 +465,12 @@ class DoctrTextRecognizer(TextRecognizer):
         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)
-        device_str = _set_device_str(device)
-        DoctrTextRecognizer.load_model(path_weights, doctr_predictor, device_str, lib)
+        doctr_predictor = DoctrTextRecognizer.build_model(architecture, lib, 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:
+    def get_name(path_weights: PathLikeOrStr, architecture: str) -> str:
         """Returns the name of the model"""
         return f"doctr_{architecture}" + "_".join(Path(path_weights).parts[-2:])
 
@@ -470,6 +479,9 @@ class DoctrTextRecognizer(TextRecognizer):
         """Auto select the DL library from the installed and from environment variables"""
         return auto_select_lib_for_doctr()
 
+    def clear_model(self) -> None:
+        self.doctr_predictor = None
+
 
 class DocTrRotationTransformer(ImageTransformer):
     """
@@ -503,7 +515,7 @@ class DocTrRotationTransformer(ImageTransformer):
         self.ratio_threshold_for_lines = ratio_threshold_for_lines
         self.name = "doctr_rotation_transformer"
 
-    def transform(self, np_img: ImageType, specification: DetectionResult) -> ImageType:
+    def transform(self, np_img: PixelValues, specification: DetectionResult) -> PixelValues:
         """
         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.
@@ -514,19 +526,18 @@ class DocTrRotationTransformer(ImageTransformer):
         """
         return viz_handler.rotate_image(np_img, specification.angle)  # type: ignore
 
-    def predict(self, np_img: ImageType) -> DetectionResult:
+    def predict(self, np_img: PixelValues) -> 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]:
+    def get_requirements(cls) -> list[Requirement]:
         return [get_doctr_requirement()]
 
-    def clone(self) -> PredictorBase:
+    def clone(self) -> DocTrRotationTransformer:
         return self.__class__(self.number_contours, self.ratio_threshold_for_lines)
 
-    @staticmethod
-    def possible_category() -> PageType:
-        return PageType.angle
+    def get_category_names(self) -> tuple[ObjectTypes, ...]:
+        return (PageType.ANGLE,)

deepdoctection/extern/fastlang.py

@@ -18,16 +18,22 @@
 """
 Deepdoctection wrappers for fasttext language detection models
 """
+from __future__ import annotations
+
+import os
 from abc import ABC
-from copy import copy
 from pathlib import Path
-from typing import Any, List, Mapping, Tuple, Union
+from types import MappingProxyType
+from typing import Any, Mapping, Union
+
+from lazy_imports import try_import
 
-from ..utils.file_utils import Requirement, fasttext_available, get_fasttext_requirement
+from ..utils.file_utils import Requirement, get_fasttext_requirement
 from ..utils.settings import TypeOrStr, get_type
-from .base import DetectionResult, LanguageDetector, PredictorBase
+from ..utils.types import PathLikeOrStr
+from .base import DetectionResult, LanguageDetector, ModelCategories
 
-if fasttext_available():
+with try_import() as import_guard:
     from fasttext import load_model  # type: ignore
 
 
@@ -36,22 +42,23 @@ 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:
+    def __init__(self, categories: Mapping[int, TypeOrStr], categories_orig: 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()})
+        self.categories = ModelCategories(init_categories=categories)
+        self.categories_orig = MappingProxyType({cat_orig: get_type(cat) for cat_orig, cat in categories_orig.items()})
 
-    def output_to_detection_result(self, output: Union[Tuple[Any, Any]]) -> DetectionResult:
+    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])
+        return DetectionResult(text=self.categories_orig[output[0][0]], score=output[1][0])
 
     @staticmethod
-    def get_name(path_weights: str) -> str:
+    def get_name(path_weights: PathLikeOrStr) -> str:
         """Returns the name of the model"""
         return "fasttext_" + "_".join(Path(path_weights).parts[-2:])
 
@@ -78,15 +85,17 @@ class FasttextLangDetector(FasttextLangDetectorMixin):
 
     """
 
-    def __init__(self, path_weights: str, categories: Mapping[str, TypeOrStr]):
+    def __init__(
+        self, path_weights: PathLikeOrStr, categories: Mapping[int, TypeOrStr], categories_orig: Mapping[str, TypeOrStr]
+    ):
         """
         :param path_weights: path to model weights
         :param categories: A dict with the model output label and value. We use as convention the ISO 639-2 language
                            code.
         """
-        super().__init__(categories)
+        super().__init__(categories, categories_orig)
 
-        self.path_weights = path_weights
+        self.path_weights = Path(path_weights)
 
         self.name = self.get_name(self.path_weights)
         self.model_id = self.get_model_id()
@@ -98,16 +107,16 @@ class FasttextLangDetector(FasttextLangDetectorMixin):
         return self.output_to_detection_result(output)
 
     @classmethod
-    def get_requirements(cls) -> List[Requirement]:
+    def get_requirements(cls) -> list[Requirement]:
         return [get_fasttext_requirement()]
 
-    def clone(self) -> PredictorBase:
-        return self.__class__(self.path_weights, self.categories)
+    def clone(self) -> FasttextLangDetector:
+        return self.__class__(self.path_weights, self.categories.get_categories(), self.categories_orig)
 
     @staticmethod
-    def get_wrapped_model(path_weights: str) -> Any:
+    def get_wrapped_model(path_weights: PathLikeOrStr) -> Any:
         """
         Get the wrapped model
         :param path_weights: path to model weights
         """
-        return load_model(path_weights)
+        return load_model(os.fspath(path_weights))