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

deepdoctection/extern/d2detect.py

@@ -18,43 +18,40 @@
 """
 D2 GeneralizedRCNN model as predictor for deepdoctection pipeline
 """
+from __future__ import annotations
+
 import io
+import os
 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 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.settings import DefaultType, ObjectTypes, TypeOrStr, get_type
 from ..utils.transform import InferenceResize, ResizeTransform
-from .base import DetectionResult, ObjectDetector, PredictorBase
+from ..utils.types import PathLikeOrStr, PixelValues, Requirement
+from .base import DetectionResult, ModelCategories, ObjectDetector
 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.
@@ -71,11 +68,11 @@ def _d2_post_processing(
 
 
 def d2_predict_image(
-    np_img: ImageType,
-    predictor: "nn.Module",
+    np_img: PixelValues,
+    predictor: nn.Module,
     resizer: InferenceResize,
     nms_thresh_class_agnostic: float,
-) -> List[DetectionResult]:
+) -> list[DetectionResult]:
     """
     Run detection on one image, using the D2 model callable. It will also handle the preprocessing internally which
     is using a custom resizing within some bounds.
@@ -107,8 +104,8 @@ def d2_predict_image(
 
 
 def d2_jit_predict_image(
-    np_img: ImageType, d2_predictor: "nn.Module", resizer: InferenceResize, nms_thresh_class_agnostic: float
-) -> List[DetectionResult]:
+    np_img: PixelValues, 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
     is using a custom resizing within some bounds. Moreover, and different from the setting where Detectron2 is used
@@ -152,7 +149,7 @@ class D2FrcnnDetectorMixin(ObjectDetector, ABC):
 
     def __init__(
         self,
-        categories: Mapping[str, TypeOrStr],
+        categories: Mapping[int, TypeOrStr],
         filter_categories: Optional[Sequence[TypeOrStr]] = None,
     ):
         """
@@ -163,37 +160,31 @@ class D2FrcnnDetectorMixin(ObjectDetector, ABC):
                                   be filtered. Pass a list of category names that must not be returned
         """
 
+        self.categories = ModelCategories(init_categories=categories)
         if filter_categories:
-            filter_categories = [get_type(cat) for cat in filter_categories]
-        self.filter_categories = filter_categories
-        self._categories_d2 = self._map_to_d2_categories(copy(categories))
-        self.categories = {idx: get_type(cat) for idx, cat in categories.items()}
+            self.categories.filter_categories = tuple(get_type(cat) for cat in filter_categories)
 
-    def _map_category_names(self, detection_results: List[DetectionResult]) -> List[DetectionResult]:
+    def _map_category_names(self, detection_results: list[DetectionResult]) -> list[DetectionResult]:
         """
         Populating category names to detection results
 
         :param detection_results: list of detection results. Will also filter categories
         :return: List of detection results with attribute class_name populated
         """
-        filtered_detection_result: List[DetectionResult] = []
+        filtered_detection_result: list[DetectionResult] = []
+        shifted_categories = self.categories.shift_category_ids(shift_by=-1)
         for result in detection_results:
-            result.class_name = self._categories_d2[str(result.class_id)]
-            if isinstance(result.class_id, int):
-                result.class_id += 1
-            if self.filter_categories:
-                if result.class_name not in self.filter_categories:
+            result.class_name = shifted_categories.get(
+                result.class_id if result.class_id is not None else -1, DefaultType.DEFAULT_TYPE
+            )
+            if result.class_name != DefaultType.DEFAULT_TYPE:
+                if result.class_id is not None:
+                    result.class_id += 1
                     filtered_detection_result.append(result)
-            else:
-                filtered_detection_result.append(result)
         return filtered_detection_result
 
-    @classmethod
-    def _map_to_d2_categories(cls, categories: Mapping[str, TypeOrStr]) -> Dict[str, ObjectTypes]:
-        return {str(int(k) - 1): get_type(v) for k, v in categories.items()}
-
-    def possible_categories(self) -> List[ObjectTypes]:
-        return list(self.categories.values())
+    def get_category_names(self) -> tuple[ObjectTypes, ...]:
+        return self.categories.get_categories(as_dict=False)
 
     @staticmethod
     def get_inference_resizer(min_size_test: int, max_size_test: int) -> InferenceResize:
@@ -205,7 +196,7 @@ class D2FrcnnDetectorMixin(ObjectDetector, ABC):
         return InferenceResize(min_size_test, max_size_test)
 
     @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"detectron2_{architecture}" + "_".join(Path(path_weights).parts[-2:])
 
@@ -234,11 +225,11 @@ class D2FrcnnDetector(D2FrcnnDetectorMixin):
 
     def __init__(
         self,
-        path_yaml: str,
-        path_weights: str,
-        categories: Mapping[str, TypeOrStr],
-        config_overwrite: Optional[List[str]] = None,
-        device: Optional[Literal["cpu", "cuda"]] = None,
+        path_yaml: PathLikeOrStr,
+        path_weights: PathLikeOrStr,
+        categories: Mapping[int, TypeOrStr],
+        config_overwrite: Optional[list[str]] = None,
+        device: Optional[Union[Literal["cpu", "cuda"], torch.device]] = None,
         filter_categories: Optional[Sequence[TypeOrStr]] = None,
     ):
         """
@@ -261,18 +252,15 @@ class D2FrcnnDetector(D2FrcnnDetectorMixin):
         """
         super().__init__(categories, filter_categories)
 
-        self.path_weights = path_weights
-        self.path_yaml = path_yaml
+        self.path_weights = Path(path_weights)
+        self.path_yaml = Path(path_yaml)
 
         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 +270,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: PathLikeOrStr, 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_file(os.fspath(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,11 +291,11 @@ 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: PathLikeOrStr) -> None:
         checkpointer = DetectionCheckpointer(wrapped_model)
-        checkpointer.load(path_weights)
+        checkpointer.load(os.fspath(path_weights))
 
-    def predict(self, np_img: ImageType) -> List[DetectionResult]:
+    def predict(self, np_img: PixelValues) -> list[DetectionResult]:
         """
         Prediction per image.
 
@@ -326,23 +311,26 @@ class D2FrcnnDetector(D2FrcnnDetectorMixin):
         return self._map_category_names(detection_results)
 
     @classmethod
-    def get_requirements(cls) -> List[Requirement]:
+    def get_requirements(cls) -> list[Requirement]:
         return [get_pytorch_requirement(), get_detectron2_requirement()]
 
-    def clone(self) -> PredictorBase:
+    def clone(self) -> D2FrcnnDetector:
         return self.__class__(
             self.path_yaml,
             self.path_weights,
-            self.categories,
+            self.categories.get_categories(),
             self.config_overwrite,
             self.device,
-            self.filter_categories,
+            self.categories.filter_categories,
         )
 
     @staticmethod
     def get_wrapped_model(
-        path_yaml: str, path_weights: str, config_overwrite: List[str], device: Literal["cpu", "cuda"]
-    ) -> "GeneralizedRCNN":
+        path_yaml: PathLikeOrStr,
+        path_weights: PathLikeOrStr,
+        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 +353,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)
@@ -374,14 +361,17 @@ class D2FrcnnDetector(D2FrcnnDetectorMixin):
         return model
 
     @staticmethod
-    def _get_d2_config_list(path_weights: str, config_overwrite: List[str]) -> List[str]:
-        d2_conf_list = ["MODEL.WEIGHTS", path_weights]
+    def _get_d2_config_list(path_weights: PathLikeOrStr, config_overwrite: list[str]) -> list[str]:
+        d2_conf_list = ["MODEL.WEIGHTS", os.fspath(path_weights)]
         config_overwrite = config_overwrite if config_overwrite else []
         for conf in config_overwrite:
             key, val = conf.split("=", maxsplit=1)
             d2_conf_list.extend([key, val])
         return d2_conf_list
 
+    def clear_model(self) -> None:
+        self.d2_predictor = None
+
 
 class D2FrcnnTracingDetector(D2FrcnnDetectorMixin):
     """
@@ -409,10 +399,10 @@ class D2FrcnnTracingDetector(D2FrcnnDetectorMixin):
 
     def __init__(
         self,
-        path_yaml: str,
-        path_weights: str,
-        categories: Mapping[str, TypeOrStr],
-        config_overwrite: Optional[List[str]] = None,
+        path_yaml: PathLikeOrStr,
+        path_weights: PathLikeOrStr,
+        categories: Mapping[int, TypeOrStr],
+        config_overwrite: Optional[list[str]] = None,
         filter_categories: Optional[Sequence[TypeOrStr]] = None,
     ):
         """
@@ -432,8 +422,8 @@ class D2FrcnnTracingDetector(D2FrcnnDetectorMixin):
 
         super().__init__(categories, filter_categories)
 
-        self.path_weights = path_weights
-        self.path_yaml = path_yaml
+        self.path_weights = Path(path_weights)
+        self.path_yaml = Path(path_yaml)
 
         self.config_overwrite = copy(config_overwrite)
         self.cfg = self._set_config(self.path_yaml, self.path_weights, self.config_overwrite)
@@ -445,14 +435,16 @@ class D2FrcnnTracingDetector(D2FrcnnDetectorMixin):
         self.d2_predictor = self.get_wrapped_model(self.path_weights)
 
     @staticmethod
-    def _set_config(path_yaml: str, path_weights: str, config_overwrite: Optional[List[str]]) -> AttrDict:
+    def _set_config(
+        path_yaml: PathLikeOrStr, path_weights: PathLikeOrStr, config_overwrite: Optional[list[str]]
+    ) -> AttrDict:
         cfg = set_config_by_yaml(path_yaml)
         config_overwrite = config_overwrite if config_overwrite else []
-        config_overwrite.extend([f"MODEL.WEIGHTS={path_weights}"])
+        config_overwrite.extend([f"MODEL.WEIGHTS={os.fspath(path_weights)}"])
         cfg.update_args(config_overwrite)
         return cfg
 
-    def predict(self, np_img: ImageType) -> List[DetectionResult]:
+    def predict(self, np_img: PixelValues) -> list[DetectionResult]:
         """
         Prediction per image.
 
@@ -468,46 +460,23 @@ class D2FrcnnTracingDetector(D2FrcnnDetectorMixin):
         return self._map_category_names(detection_results)
 
     @classmethod
-    def get_requirements(cls) -> List[Requirement]:
+    def get_requirements(cls) -> list[Requirement]:
         return [get_pytorch_requirement()]
 
-    @classmethod
-    def _map_to_d2_categories(cls, categories: Mapping[str, TypeOrStr]) -> Dict[str, ObjectTypes]:
-        return {str(int(k) - 1): get_type(v) for k, v in categories.items()}
-
-    def clone(self) -> PredictorBase:
+    def clone(self) -> D2FrcnnTracingDetector:
         return self.__class__(
             self.path_yaml,
             self.path_weights,
-            self.categories,
+            self.categories.get_categories(),
             self.config_overwrite,
-            self.filter_categories,
+            self.categories.filter_categories,
         )
 
-    def _map_category_names(self, detection_results: List[DetectionResult]) -> List[DetectionResult]:
-        """
-        Populating category names to detection results
-
-        :param detection_results: list of detection results. Will also filter categories
-        :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_d2[str(result.class_id)]
-            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
-
-    def possible_categories(self) -> List[ObjectTypes]:
-        return list(self.categories.values())
+    def get_category_names(self) -> tuple[ObjectTypes, ...]:
+        return self.categories.get_categories(as_dict=False)
 
     @staticmethod
-    def get_wrapped_model(path_weights: str) -> Any:
+    def get_wrapped_model(path_weights: PathLikeOrStr) -> torch.jit.ScriptModule:
         """
         Get the wrapped model. Useful if one do not want to build the wrapper but only needs the instantiated model.
 
@@ -518,3 +487,6 @@ class D2FrcnnTracingDetector(D2FrcnnDetectorMixin):
             buffer = io.BytesIO(file.read())
         # Load all tensors to the original device
         return torch.jit.load(buffer)
+
+    def clear_model(self) -> None:
+        self.d2_predictor = None  # type: ignore

deepdoctection/extern/deskew.py

@@ -18,16 +18,17 @@
 """
 jdeskew estimator and rotator to deskew images: <https://github.com/phamquiluan/jdeskew>
 """
+from __future__ import annotations
 
-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.settings import PageType
+from ..utils.file_utils import get_jdeskew_requirement
+from ..utils.settings import ObjectTypes, PageType
+from ..utils.types import PixelValues, Requirement
 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
 
 
@@ -42,7 +43,7 @@ class Jdeskewer(ImageTransformer):
         self.model_id = self.get_model_id()
         self.min_angle_rotation = min_angle_rotation
 
-    def transform(self, np_img: ImageType, specification: DetectionResult) -> ImageType:
+    def transform(self, np_img: PixelValues, specification: DetectionResult) -> PixelValues:
         """
         Rotation of the image according to the angle determined by the jdeskew estimator.
 
@@ -59,7 +60,7 @@ class Jdeskewer(ImageTransformer):
             return viz_handler.rotate_image(np_img, specification.angle)  # type: ignore
         return np_img
 
-    def predict(self, np_img: ImageType) -> DetectionResult:
+    def predict(self, np_img: PixelValues) -> DetectionResult:
         """
         Predict the angle of the image to deskew it.
 
@@ -69,12 +70,14 @@ class Jdeskewer(ImageTransformer):
         return DetectionResult(angle=round(float(get_angle(np_img)), 4))
 
     @classmethod
-    def get_requirements(cls) -> List[Requirement]:
+    def get_requirements(cls) -> list[Requirement]:
         """
         Get a list of requirements for running the detector
         """
         return [get_jdeskew_requirement()]
 
-    @staticmethod
-    def possible_category() -> PageType:
-        return PageType.angle
+    def clone(self) -> Jdeskewer:
+        return self.__class__(self.min_angle_rotation)
+
+    def get_category_names(self) -> tuple[ObjectTypes, ...]:
+        return (PageType.ANGLE,)