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

deepdoctection/extern/d2detect.py

@@ -18,43 +18,39 @@
 """
 D2 GeneralizedRCNN model as predictor for deepdoctection pipeline
 """
+from __future__ import annotations
+
 import io
 from abc import ABC
 from copy import copy
 from pathlib import Path
-from typing import Any, Dict, List, Literal, Mapping, Optional, Sequence
+from typing import Any, Dict, List, Literal, Mapping, Optional, Sequence, Union
 
 import numpy as np
+from lazy_imports import try_import
 
 from ..utils.detection_types import ImageType, Requirement
-from ..utils.file_utils import (
-    detectron2_available,
-    get_detectron2_requirement,
-    get_pytorch_requirement,
-    pytorch_available,
-)
+from ..utils.file_utils import get_detectron2_requirement, get_pytorch_requirement
 from ..utils.metacfg import AttrDict, set_config_by_yaml
 from ..utils.settings import ObjectTypes, TypeOrStr, get_type
 from ..utils.transform import InferenceResize, ResizeTransform
 from .base import DetectionResult, ObjectDetector, PredictorBase
 from .pt.nms import batched_nms
-from .pt.ptutils import set_torch_auto_device
+from .pt.ptutils import get_torch_device
 
-if pytorch_available():
+with try_import() as pt_import_guard:
     import torch
     import torch.cuda
     from torch import nn  # pylint: disable=W0611
 
-if detectron2_available():
+with try_import() as d2_import_guard:
     from detectron2.checkpoint import DetectionCheckpointer
     from detectron2.config import CfgNode, get_cfg  # pylint: disable=W0611
     from detectron2.modeling import GeneralizedRCNN, build_model  # pylint: disable=W0611
     from detectron2.structures import Instances  # pylint: disable=W0611
 
 
-def _d2_post_processing(
-    predictions: Dict[str, "Instances"], nms_thresh_class_agnostic: float
-) -> Dict[str, "Instances"]:
+def _d2_post_processing(predictions: Dict[str, Instances], nms_thresh_class_agnostic: float) -> Dict[str, Instances]:
     """
     D2 postprocessing steps, so that detection outputs are aligned with outputs of other packages (e.g. Tensorpack).
     Apply a class agnostic NMS.
@@ -72,7 +68,7 @@ def _d2_post_processing(
 
 def d2_predict_image(
     np_img: ImageType,
-    predictor: "nn.Module",
+    predictor: nn.Module,
     resizer: InferenceResize,
     nms_thresh_class_agnostic: float,
 ) -> List[DetectionResult]:
@@ -107,7 +103,7 @@ def d2_predict_image(
 
 
 def d2_jit_predict_image(
-    np_img: ImageType, d2_predictor: "nn.Module", resizer: InferenceResize, nms_thresh_class_agnostic: float
+    np_img: ImageType, d2_predictor: nn.Module, resizer: InferenceResize, nms_thresh_class_agnostic: float
 ) -> List[DetectionResult]:
     """
     Run detection on an image using torchscript. It will also handle the preprocessing internally which
@@ -238,7 +234,7 @@ class D2FrcnnDetector(D2FrcnnDetectorMixin):
         path_weights: str,
         categories: Mapping[str, TypeOrStr],
         config_overwrite: Optional[List[str]] = None,
-        device: Optional[Literal["cpu", "cuda"]] = None,
+        device: Optional[Union[Literal["cpu", "cuda"], torch.device]] = None,
         filter_categories: Optional[Sequence[TypeOrStr]] = None,
     ):
         """
@@ -266,13 +262,10 @@ class D2FrcnnDetector(D2FrcnnDetectorMixin):
 
         config_overwrite = config_overwrite if config_overwrite else []
         self.config_overwrite = config_overwrite
-        if device is not None:
-            self.device = device
-        else:
-            self.device = set_torch_auto_device()
+        self.device = get_torch_device(device)
 
         d2_conf_list = self._get_d2_config_list(path_weights, config_overwrite)
-        self.cfg = self._set_config(path_yaml, d2_conf_list, device)
+        self.cfg = self._set_config(path_yaml, d2_conf_list, self.device)
 
         self.name = self.get_name(path_weights, self.cfg.MODEL.META_ARCHITECTURE)
         self.model_id = self.get_model_id()
@@ -282,21 +275,18 @@ class D2FrcnnDetector(D2FrcnnDetectorMixin):
         self.resizer = self.get_inference_resizer(self.cfg.INPUT.MIN_SIZE_TEST, self.cfg.INPUT.MAX_SIZE_TEST)
 
     @staticmethod
-    def _set_config(
-        path_yaml: str, d2_conf_list: List[str], device: Optional[Literal["cpu", "cuda"]] = None
-    ) -> "CfgNode":
+    def _set_config(path_yaml: str, d2_conf_list: List[str], device: torch.device) -> CfgNode:
         cfg = get_cfg()
         # additional attribute with default value, so that the true value can be loaded from the configs
         cfg.NMS_THRESH_CLASS_AGNOSTIC = 0.1
         cfg.merge_from_file(path_yaml)
         cfg.merge_from_list(d2_conf_list)
-        if not torch.cuda.is_available() or device == "cpu":
-            cfg.MODEL.DEVICE = "cpu"
+        cfg.MODEL.DEVICE = str(device)
         cfg.freeze()
         return cfg
 
     @staticmethod
-    def _set_model(config: "CfgNode") -> "GeneralizedRCNN":
+    def _set_model(config: CfgNode) -> GeneralizedRCNN:
         """
         Build the D2 model. It uses the available builtin tools of D2
 
@@ -306,7 +296,7 @@ class D2FrcnnDetector(D2FrcnnDetectorMixin):
         return build_model(config.clone()).eval()
 
     @staticmethod
-    def _instantiate_d2_predictor(wrapped_model: "GeneralizedRCNN", path_weights: str) -> None:
+    def _instantiate_d2_predictor(wrapped_model: GeneralizedRCNN, path_weights: str) -> None:
         checkpointer = DetectionCheckpointer(wrapped_model)
         checkpointer.load(path_weights)
 
@@ -341,8 +331,11 @@ class D2FrcnnDetector(D2FrcnnDetectorMixin):
 
     @staticmethod
     def get_wrapped_model(
-        path_yaml: str, path_weights: str, config_overwrite: List[str], device: Literal["cpu", "cuda"]
-    ) -> "GeneralizedRCNN":
+        path_yaml: str,
+        path_weights: str,
+        config_overwrite: List[str],
+        device: Optional[Union[Literal["cpu", "cuda"], torch.device]] = None,
+    ) -> GeneralizedRCNN:
         """
         Get the wrapped model. Useful if one do not want to build the wrapper but only needs the instantiated model.
 
@@ -365,8 +358,7 @@ class D2FrcnnDetector(D2FrcnnDetectorMixin):
         :return: Detectron2 GeneralizedRCNN model
         """
 
-        if device is None:
-            device = set_torch_auto_device()
+        device = get_torch_device(device)
         d2_conf_list = D2FrcnnDetector._get_d2_config_list(path_weights, config_overwrite)
         cfg = D2FrcnnDetector._set_config(path_yaml, d2_conf_list, device)
         model = D2FrcnnDetector._set_model(cfg)

deepdoctection/extern/deskew.py

@@ -21,13 +21,15 @@ jdeskew estimator and rotator to deskew images: <https://github.com/phamquiluan/
 
 from typing import List
 
+from lazy_imports import try_import
+
 from ..utils.detection_types import ImageType, Requirement
-from ..utils.file_utils import get_jdeskew_requirement, jdeskew_available
+from ..utils.file_utils import get_jdeskew_requirement
 from ..utils.settings import PageType
 from ..utils.viz import viz_handler
 from .base import DetectionResult, ImageTransformer
 
-if jdeskew_available():
+with try_import() as import_guard:
     from jdeskew.estimator import get_angle
 
 

deepdoctection/extern/doctrocr.py

@@ -18,32 +18,40 @@
 """
 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.env_info import get_device
 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.viz import viz_handler
 from .base import DetectionResult, ImageTransformer, ObjectDetector, PredictorBase, TextRecognizer
-from .pt.ptutils import set_torch_auto_device
+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,32 +59,19 @@ 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():
+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")
@@ -89,29 +84,33 @@ def _load_model(path_weights: str, doctr_predictor: Any, device: str, lib: Liter
 
 
 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"):
         return "PT"
-    raise DependencyError("Neither Tensorflow nor PyTorch has been installed. Cannot use DoctrTextlineDetector")
+    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: str) -> List[DetectionResult]:
+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
@@ -122,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
@@ -132,15 +134,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)
     ]
@@ -208,7 +213,7 @@ class DoctrTextlineDetector(DoctrTextlineDetectorMixin):
         architecture: str,
         path_weights: str,
         categories: Mapping[str, TypeOrStr],
-        device: Optional[Literal["cpu", "cuda"]] = None,
+        device: Optional[Union[Literal["cpu", "cuda"], torch.device, tf.device]] = None,
         lib: Optional[Literal["PT", "TF"]] = None,
     ) -> None:
         """
@@ -217,7 +222,7 @@ 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)
@@ -227,18 +232,12 @@ class DoctrTextlineDetector(DoctrTextlineDetectorMixin):
         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]:
         """
@@ -247,28 +246,30 @@ class DoctrTextlineDetector(DoctrTextlineDetectorMixin):
         :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)
+        return self.__class__(self.architecture, self.path_weights, self.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: str, doctr_predictor: Any, 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: str, device: Union[torch.device, tf.device], lib: Literal["PT", "TF"]
     ) -> Any:
         """
         Get the inner (wrapped) model.
@@ -286,8 +287,7 @@ 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
 
 
@@ -325,14 +325,13 @@ class DoctrTextRecognizer(TextRecognizer):
 
                  for dp in df:
                      ...
-
     """
 
     def __init__(
         self,
         architecture: str,
         path_weights: str,
-        device: Optional[Literal["cpu", "cuda"]] = None,
+        device: Optional[Union[Literal["cpu", "cuda"], torch.device, tf.device]] = None,
         lib: Optional[Literal["PT", "TF"]] = None,
         path_config_json: Optional[str] = None,
     ) -> None:
@@ -355,22 +354,16 @@ class DoctrTextRecognizer(TextRecognizer):
         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.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]:
@@ -381,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
@@ -393,10 +386,12 @@ 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)
 
     @staticmethod
-    def load_model(path_weights: str, doctr_predictor: Any, device: str, lib: Literal["PT", "TF"]) -> None:
+    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(path_weights, doctr_predictor, device, lib)
 
@@ -438,7 +433,7 @@ class DoctrTextRecognizer(TextRecognizer):
     def get_wrapped_model(
         architecture: str,
         path_weights: str,
-        device: Literal["cpu", "cuda"],
+        device: Union[torch.device, tf.device],
         lib: Literal["PT", "TF"],
         path_config_json: Optional[str] = None,
     ) -> Any:
@@ -456,8 +451,7 @@ class DoctrTextRecognizer(TextRecognizer):
         :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)
+        DoctrTextRecognizer.load_model(path_weights, doctr_predictor, device, lib)
         return doctr_predictor
 
     @staticmethod

deepdoctection/extern/fastlang.py

@@ -23,11 +23,13 @@ from copy import copy
 from pathlib import Path
 from typing import Any, List, Mapping, Tuple, Union
 
-from ..utils.file_utils import Requirement, fasttext_available, get_fasttext_requirement
+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
 
 

deepdoctection/extern/hfdetr.py

@@ -18,27 +18,25 @@
 """
 HF Detr model for object detection.
 """
+from __future__ import annotations
 
 from abc import ABC
 from pathlib import Path
-from typing import List, Literal, Mapping, Optional, Sequence
+from typing import List, Literal, Mapping, Optional, Sequence, Union
+
+from lazy_imports import try_import
 
 from ..utils.detection_types import ImageType, Requirement
-from ..utils.file_utils import (
-    get_pytorch_requirement,
-    get_transformers_requirement,
-    pytorch_available,
-    transformers_available,
-)
+from ..utils.file_utils import get_pytorch_requirement, get_transformers_requirement
 from ..utils.settings import TypeOrStr, get_type
 from .base import DetectionResult, ObjectDetector
-from .pt.ptutils import set_torch_auto_device
+from .pt.ptutils import get_torch_device
 
-if pytorch_available():
+with try_import() as pt_import_guard:
     import torch  # pylint: disable=W0611
     from torchvision.ops import boxes as box_ops  # type: ignore
 
-if transformers_available():
+with try_import() as tr_import_guard:
     from transformers import (  # pylint: disable=W0611
         AutoFeatureExtractor,
         DetrFeatureExtractor,
@@ -48,16 +46,16 @@ if transformers_available():
 
 
 def _detr_post_processing(
-    boxes: "torch.Tensor", scores: "torch.Tensor", labels: "torch.Tensor", nms_thresh: float
-) -> "torch.Tensor":
+    boxes: torch.Tensor, scores: torch.Tensor, labels: torch.Tensor, nms_thresh: float
+) -> torch.Tensor:
     return box_ops.batched_nms(boxes.float(), scores, labels, nms_thresh)
 
 
 def detr_predict_image(
     np_img: ImageType,
-    predictor: "TableTransformerForObjectDetection",
-    feature_extractor: "DetrFeatureExtractor",
-    device: Literal["cpu", "cuda"],
+    predictor: TableTransformerForObjectDetection,
+    feature_extractor: DetrFeatureExtractor,
+    device: torch.device,
     threshold: float,
     nms_threshold: float,
 ) -> List[DetectionResult]:
@@ -168,7 +166,7 @@ class HFDetrDerivedDetector(HFDetrDerivedDetectorMixin):
         path_weights: str,
         path_feature_extractor_config_json: str,
         categories: Mapping[str, TypeOrStr],
-        device: Optional[Literal["cpu", "cuda"]] = None,
+        device: Optional[Union[Literal["cpu", "cuda"], torch.device]] = None,
         filter_categories: Optional[Sequence[TypeOrStr]] = None,
     ):
         """
@@ -195,10 +193,7 @@ class HFDetrDerivedDetector(HFDetrDerivedDetectorMixin):
         self.hf_detr_predictor = self.get_model(self.path_weights, self.config)
         self.feature_extractor = self.get_pre_processor(self.path_feature_extractor_config)
 
-        if device is not None:
-            self.device = device
-        else:
-            self.device = set_torch_auto_device()
+        self.device = get_torch_device(device)
         self.hf_detr_predictor.to(self.device)
 
     def predict(self, np_img: ImageType) -> List[DetectionResult]:
@@ -213,7 +208,7 @@ class HFDetrDerivedDetector(HFDetrDerivedDetectorMixin):
         return self._map_category_names(results)
 
     @staticmethod
-    def get_model(path_weights: str, config: "PretrainedConfig") -> "TableTransformerForObjectDetection":
+    def get_model(path_weights: str, config: PretrainedConfig) -> TableTransformerForObjectDetection:
         """
         Builds the Detr model
 
@@ -226,7 +221,7 @@ class HFDetrDerivedDetector(HFDetrDerivedDetectorMixin):
         )
 
     @staticmethod
-    def get_pre_processor(path_feature_extractor_config: str) -> "DetrFeatureExtractor":
+    def get_pre_processor(path_feature_extractor_config: str) -> DetrFeatureExtractor:
         """
         Builds the feature extractor
 
@@ -235,7 +230,7 @@ class HFDetrDerivedDetector(HFDetrDerivedDetectorMixin):
         return AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path=path_feature_extractor_config)
 
     @staticmethod
-    def get_config(path_config: str) -> "PretrainedConfig":
+    def get_config(path_config: str) -> PretrainedConfig:
         """
         Builds the config
 
@@ -252,15 +247,15 @@ class HFDetrDerivedDetector(HFDetrDerivedDetectorMixin):
     def get_requirements(cls) -> List[Requirement]:
         return [get_pytorch_requirement(), get_transformers_requirement()]
 
-    def clone(self) -> "HFDetrDerivedDetector":
+    def clone(self) -> HFDetrDerivedDetector:
         return self.__class__(
             self.path_config, self.path_weights, self.path_feature_extractor_config, self.categories, self.device
         )
 
     @staticmethod
     def get_wrapped_model(
-        path_config_json: str, path_weights: str, device: Optional[Literal["cpu", "cuda"]] = None
-    ) -> "TableTransformerForObjectDetection":
+        path_config_json: str, path_weights: str, device: Optional[Union[Literal["cpu", "cuda"], torch.device]] = None
+    ) -> TableTransformerForObjectDetection:
         """
         Get the wrapped model
 
@@ -271,6 +266,5 @@ class HFDetrDerivedDetector(HFDetrDerivedDetectorMixin):
         """
         config = HFDetrDerivedDetector.get_config(path_config_json)
         hf_detr_predictor = HFDetrDerivedDetector.get_model(path_weights, config)
-        if device is None:
-            device = set_torch_auto_device()
+        device = get_torch_device()
         return hf_detr_predictor.to(device)