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

deepdoctection/extern/hfdetr.py

@@ -18,25 +18,25 @@
 """
 HF Detr model for object detection.
 """
+from __future__ import annotations
 
-from typing import List, Literal, Mapping, Optional, Sequence
+from abc import ABC
+from pathlib import Path
+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,
@@ -46,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]:
@@ -94,7 +94,48 @@ def detr_predict_image(
     ]
 
 
-class HFDetrDerivedDetector(ObjectDetector):
+class HFDetrDerivedDetectorMixin(ObjectDetector, ABC):
+    """Base class for Detr object detector. This class only implements the basic wrapper functions"""
+
+    def __init__(self, categories: Mapping[str, TypeOrStr], filter_categories: Optional[Sequence[TypeOrStr]] = None):
+        """
+
+        :param categories: A dict with key (indices) and values (category names).
+        :param filter_categories: The model might return objects that are not supposed to be predicted and that should
+                                  be filtered. Pass a list of category names that must not be returned
+        """
+        self.categories = {idx: get_type(cat) for idx, cat in categories.items()}
+        if filter_categories:
+            filter_categories = [get_type(cat) for cat in filter_categories]
+        self.filter_categories = filter_categories
+
+    def _map_category_names(self, detection_results: List[DetectionResult]) -> List[DetectionResult]:
+        """
+        Populating category names to detection results. Will also filter categories
+
+        :param detection_results: list of detection results
+        :return: List of detection results with attribute class_name populated
+        """
+        filtered_detection_result: List[DetectionResult] = []
+        for result in detection_results:
+            result.class_name = self.categories[str(result.class_id + 1)]  # type: ignore
+            if isinstance(result.class_id, int):
+                result.class_id += 1
+            if self.filter_categories:
+                if result.class_name not in self.filter_categories:
+                    filtered_detection_result.append(result)
+            else:
+                filtered_detection_result.append(result)
+
+        return filtered_detection_result
+
+    @staticmethod
+    def get_name(path_weights: str) -> str:
+        """Returns the name of the model"""
+        return "Transformers_Tatr_" + "_".join(Path(path_weights).parts[-2:])
+
+
+class HFDetrDerivedDetector(HFDetrDerivedDetectorMixin):
     """
     Model wrapper for TableTransformerForObjectDetection that again is based on
 
@@ -125,7 +166,7 @@ class HFDetrDerivedDetector(ObjectDetector):
         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,
     ):
         """
@@ -138,26 +179,22 @@ class HFDetrDerivedDetector(ObjectDetector):
         :param filter_categories: The model might return objects that are not supposed to be predicted and that should
                                   be filtered. Pass a list of category names that must not be returned
         """
-        self.name = "Detr"
-        self.categories = {idx: get_type(cat) for idx, cat in categories.items()}
+        super().__init__(categories, filter_categories)
+
         self.path_config = path_config_json
         self.path_weights = path_weights
         self.path_feature_extractor_config = path_feature_extractor_config_json
-        self.config = PretrainedConfig.from_pretrained(pretrained_model_name_or_path=self.path_config)
-        self.config.use_timm_backbone = True
-        self.config.threshold = 0.1
-        self.config.nms_threshold = 0.05
-        self.hf_detr_predictor = self.set_model(path_weights)
-        self.feature_extractor = self.set_pre_processor()
-
-        if device is not None:
-            self.device = device
-        else:
-            self.device = set_torch_auto_device()
+
+        self.name = self.get_name(self.path_weights)
+        self.model_id = self.get_model_id()
+
+        self.config = self.get_config(path_config_json)
+
+        self.hf_detr_predictor = self.get_model(self.path_weights, self.config)
+        self.feature_extractor = self.get_pre_processor(self.path_feature_extractor_config)
+
+        self.device = get_torch_device(device)
         self.hf_detr_predictor.to(self.device)
-        if filter_categories:
-            filter_categories = [get_type(cat) for cat in filter_categories]
-        self.filter_categories = filter_categories
 
     def predict(self, np_img: ImageType) -> List[DetectionResult]:
         results = detr_predict_image(
@@ -170,50 +207,64 @@ class HFDetrDerivedDetector(ObjectDetector):
         )
         return self._map_category_names(results)
 
-    def set_model(self, path_weights: str) -> "TableTransformerForObjectDetection":
+    @staticmethod
+    def get_model(path_weights: str, config: PretrainedConfig) -> TableTransformerForObjectDetection:
         """
         Builds the Detr model
 
-        :param path_weights: weights
+        :param path_weights: The path to the model checkpoint.
+        :param config: `PretrainedConfig`
         :return: TableTransformerForObjectDetection instance
         """
         return TableTransformerForObjectDetection.from_pretrained(
-            pretrained_model_name_or_path=path_weights, config=self.config
+            pretrained_model_name_or_path=path_weights, config=config
         )
 
-    def set_pre_processor(self) -> "DetrFeatureExtractor":
+    @staticmethod
+    def get_pre_processor(path_feature_extractor_config: str) -> DetrFeatureExtractor:
         """
         Builds the feature extractor
 
         :return: DetrFeatureExtractor
         """
-        return AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path=self.path_feature_extractor_config)
+        return AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path=path_feature_extractor_config)
 
-    def _map_category_names(self, detection_results: List[DetectionResult]) -> List[DetectionResult]:
+    @staticmethod
+    def get_config(path_config: str) -> PretrainedConfig:
         """
-        Populating category names to detection results. Will also filter categories
+        Builds the config
 
-        :param detection_results: list of detection results
-        :return: List of detection results with attribute class_name populated
+        :param path_config: The path to the json config.
+        :return: PretrainedConfig instance
         """
-        filtered_detection_result: List[DetectionResult] = []
-        for result in detection_results:
-            result.class_name = self.categories[str(result.class_id + 1)]  # type: ignore
-            if isinstance(result.class_id, int):
-                result.class_id += 1
-            if self.filter_categories:
-                if result.class_name not in self.filter_categories:
-                    filtered_detection_result.append(result)
-            else:
-                filtered_detection_result.append(result)
-
-        return filtered_detection_result
+        config = PretrainedConfig.from_pretrained(pretrained_model_name_or_path=path_config)
+        config.use_timm_backbone = True
+        config.threshold = 0.1
+        config.nms_threshold = 0.05
+        return config
 
     @classmethod
     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[Union[Literal["cpu", "cuda"], torch.device]] = None
+    ) -> TableTransformerForObjectDetection:
+        """
+        Get the wrapped model
+
+        :param path_config_json: The path to the json config.
+        :param path_weights: The path to the model checkpoint.
+        :param device: "cpu" or "cuda". If not specified will auto select depending on what is available
+        :return: TableTransformerForObjectDetection instance
+        """
+        config = HFDetrDerivedDetector.get_config(path_config_json)
+        hf_detr_predictor = HFDetrDerivedDetector.get_model(path_weights, config)
+        device = get_torch_device()
+        return hf_detr_predictor.to(device)