python-doctr 0.12.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

doctr/__init__.py

@@ -1,3 +1,2 @@
 from . import io, models, datasets, contrib, transforms, utils
-from .file_utils import is_tf_available, is_torch_available
 from .version import __version__  # noqa: F401

doctr/datasets/__init__.py

@@ -1,5 +1,3 @@
-from doctr.file_utils import is_tf_available
-
 from .generator import *
 from .coco_text import *
 from .cord import *
@@ -22,6 +20,3 @@ from .synthtext import *
 from .utils import *
 from .vocabs import *
 from .wildreceipt import *
-
-if is_tf_available():
-    from .loader import *

doctr/datasets/datasets/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *  # type: ignore[assignment]
+from .pytorch import *

doctr/datasets/datasets/pytorch.py

@@ -50,9 +50,9 @@ class AbstractDataset(_AbstractDataset):
     @staticmethod
     def collate_fn(samples: list[tuple[torch.Tensor, Any]]) -> tuple[torch.Tensor, list[Any]]:
         images, targets = zip(*samples)
-        images = torch.stack(images, dim=0)
+        images = torch.stack(images, dim=0)  # type: ignore[assignment]
 
-        return images, list(targets)
+        return images, list(targets)  # type: ignore[return-value]
 
 
 class VisionDataset(AbstractDataset, _VisionDataset):  # noqa: D101

doctr/datasets/generator/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *  # type: ignore[assignment]
+from .pytorch import *

doctr/datasets/vocabs.py

@@ -264,8 +264,6 @@ VOCABS["estonian"] = VOCABS["english"] + "šžõäöüŠŽÕÄÖÜ"
 VOCABS["esperanto"] = re.sub(r"[QqWwXxYy]", "", VOCABS["english"]) + "ĉĝĥĵŝŭĈĜĤĴŜŬ" + "₷"
 
 VOCABS["french"] = VOCABS["english"] + "àâéèêëîïôùûüçÀÂÉÈÊËÎÏÔÙÛÜÇ"
-# NOTE: legacy french is outdated, but kept for compatibility
-VOCABS["legacy_french"] = VOCABS["latin"] + "°" + "àâéèêëîïôùûçÀÂÉÈËÎÏÔÙÛÇ" + _BASE_VOCABS["currency"]
 
 VOCABS["finnish"] = VOCABS["english"] + "äöÄÖ"
 

doctr/file_utils.py

@@ -3,102 +3,13 @@
 # This program is licensed under the Apache License 2.0.
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
 
-# Adapted from https://github.com/huggingface/transformers/blob/master/src/transformers/file_utils.py
-
 import importlib.metadata
-import importlib.util
 import logging
-import os
-
-CLASS_NAME: str = "words"
-
 
-__all__ = ["is_tf_available", "is_torch_available", "requires_package", "CLASS_NAME"]
+__all__ = ["requires_package", "CLASS_NAME"]
 
+CLASS_NAME: str = "words"
 ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"}
-ENV_VARS_TRUE_AND_AUTO_VALUES = ENV_VARS_TRUE_VALUES.union({"AUTO"})
-
-USE_TF = os.environ.get("USE_TF", "AUTO").upper()
-USE_TORCH = os.environ.get("USE_TORCH", "AUTO").upper()
-
-
-if USE_TORCH in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TF not in ENV_VARS_TRUE_VALUES:
-    _torch_available = importlib.util.find_spec("torch") is not None
-    if _torch_available:
-        try:
-            _torch_version = importlib.metadata.version("torch")
-            logging.info(f"PyTorch version {_torch_version} available.")
-        except importlib.metadata.PackageNotFoundError:  # pragma: no cover
-            _torch_available = False
-else:  # pragma: no cover
-    logging.info("Disabling PyTorch because USE_TF is set")
-    _torch_available = False
-
-# Compatibility fix to make sure tensorflow.keras stays at Keras 2
-if "TF_USE_LEGACY_KERAS" not in os.environ:
-    os.environ["TF_USE_LEGACY_KERAS"] = "1"
-
-elif os.environ["TF_USE_LEGACY_KERAS"] != "1":
-    raise ValueError(
-        "docTR is only compatible with Keras 2, but you have explicitly set `TF_USE_LEGACY_KERAS` to `0`. "
-    )
-
-
-def ensure_keras_v2() -> None:  # pragma: no cover
-    if not os.environ.get("TF_USE_LEGACY_KERAS") == "1":
-        os.environ["TF_USE_LEGACY_KERAS"] = "1"
-
-
-if USE_TF in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TORCH not in ENV_VARS_TRUE_VALUES:
-    _tf_available = importlib.util.find_spec("tensorflow") is not None
-    if _tf_available:
-        candidates = (
-            "tensorflow",
-            "tensorflow-cpu",
-            "tensorflow-gpu",
-            "tf-nightly",
-            "tf-nightly-cpu",
-            "tf-nightly-gpu",
-            "intel-tensorflow",
-            "tensorflow-rocm",
-            "tensorflow-macos",
-        )
-        _tf_version = None
-        # For the metadata, we have to look for both tensorflow and tensorflow-cpu
-        for pkg in candidates:
-            try:
-                _tf_version = importlib.metadata.version(pkg)
-                break
-            except importlib.metadata.PackageNotFoundError:
-                pass
-        _tf_available = _tf_version is not None
-    if _tf_available:
-        if int(_tf_version.split(".")[0]) < 2:  # type: ignore[union-attr]  # pragma: no cover
-            logging.info(f"TensorFlow found but with version {_tf_version}. DocTR requires version 2 minimum.")
-            _tf_available = False
-        else:
-            logging.info(f"TensorFlow version {_tf_version} available.")
-            ensure_keras_v2()
-
-        import warnings
-
-        warnings.simplefilter("always", DeprecationWarning)
-        warnings.warn(
-            "Support for TensorFlow in DocTR is deprecated and will be removed in the next major release (v1.0.0). "
-            "Please switch to the PyTorch backend.",
-            DeprecationWarning,
-        )
-
-else:  # pragma: no cover
-    logging.info("Disabling Tensorflow because USE_TORCH is set")
-    _tf_available = False
-
-
-if not _torch_available and not _tf_available:  # pragma: no cover
-    raise ModuleNotFoundError(
-        "DocTR requires either TensorFlow or PyTorch to be installed. Please ensure one of them"
-        " is installed and that either USE_TF or USE_TORCH is enabled."
-    )
 
 
 def requires_package(name: str, extra_message: str | None = None) -> None:  # pragma: no cover
@@ -117,13 +28,3 @@ def requires_package(name: str, extra_message: str | None = None) -> None:  # pr
             f"\n\n{extra_message if extra_message is not None else ''} "
             f"\nPlease install it with the following command: pip install {name}\n"
         )
-
-
-def is_torch_available():
-    """Whether PyTorch is installed."""
-    return _torch_available
-
-
-def is_tf_available():
-    """Whether TensorFlow is installed."""
-    return _tf_available

doctr/io/image/__init__.py

@@ -1,8 +1,2 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
 from .base import *
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *
+from .pytorch import *

doctr/io/image/pytorch.py

@@ -95,4 +95,4 @@ def tensor_from_numpy(npy_img: np.ndarray, dtype: torch.dtype = torch.float32) -
 
 def get_img_shape(img: torch.Tensor) -> tuple[int, int]:
     """Get the shape of an image"""
-    return img.shape[-2:]
+    return img.shape[-2:]  # type: ignore[return-value]

doctr/models/_utils.py

@@ -63,7 +63,7 @@ def estimate_orientation(
         thresh = img.astype(np.uint8)
 
     page_orientation, orientation_confidence = general_page_orientation or (None, 0.0)
-    if page_orientation and orientation_confidence >= min_confidence:
+    if page_orientation is not None and orientation_confidence >= min_confidence:
         # We rotate the image to the general orientation which improves the detection
         # No expand needed bitmap is already padded
         thresh = rotate_image(thresh, -page_orientation)
@@ -100,7 +100,7 @@ def estimate_orientation(
         estimated_angle = -round(median) if abs(median) != 0 else 0
 
     # combine with the general orientation and the estimated angle
-    if page_orientation and orientation_confidence >= min_confidence:
+    if page_orientation is not None and orientation_confidence >= min_confidence:
         # special case where the estimated angle is mostly wrong:
         # case 1: - and + swapped
         # case 2: estimated angle is completely wrong
@@ -184,7 +184,7 @@ def invert_data_structure(
         dictionary of list when x is a list of dictionaries or a list of dictionaries when x is dictionary of lists
     """
     if isinstance(x, dict):
-        assert len({len(v) for v in x.values()}) == 1, "All the lists in the dictionnary should have the same length."
+        assert len({len(v) for v in x.values()}) == 1, "All the lists in the dictionary should have the same length."
         return [dict(zip(x, t)) for t in zip(*x.values())]
     elif isinstance(x, list):
         return {k: [dic[k] for dic in x] for k in x[0]}

doctr/models/classification/magc_resnet/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *
+from .pytorch import *

doctr/models/classification/magc_resnet/pytorch.py

@@ -14,7 +14,7 @@ from torch import nn
 
 from doctr.datasets import VOCABS
 
-from ..resnet.pytorch import ResNet
+from ..resnet import ResNet
 
 __all__ = ["magc_resnet31"]
 
@@ -72,7 +72,7 @@ class MAGC(nn.Module):
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         batch, _, height, width = inputs.size()
         # (N * headers, C / headers, H , W)
-        x = inputs.view(batch * self.headers, self.single_header_inplanes, height, width)
+        x = inputs.contiguous().view(batch * self.headers, self.single_header_inplanes, height, width)
         shortcut = x
         # (N * headers, C / headers, H * W)
         shortcut = shortcut.view(batch * self.headers, self.single_header_inplanes, height * width)

doctr/models/classification/mobilenet/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *
+from .pytorch import *

doctr/models/classification/predictor/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *  # type: ignore[assignment]
+from .pytorch import *

doctr/models/classification/predictor/pytorch.py

@@ -35,7 +35,7 @@ class OrientationPredictor(nn.Module):
     @torch.inference_mode()
     def forward(
         self,
-        inputs: list[np.ndarray | torch.Tensor],
+        inputs: list[np.ndarray],
     ) -> list[list[int] | list[float]]:
         # Dimension check
         if any(input.ndim != 3 for input in inputs):

doctr/models/classification/resnet/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *  # type: ignore[assignment]
+from .pytorch import *

doctr/models/classification/textnet/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *
+from .pytorch import *

doctr/models/classification/textnet/pytorch.py

@@ -11,7 +11,7 @@ from torch import nn
 
 from doctr.datasets import VOCABS
 
-from ...modules.layers.pytorch import FASTConvLayer
+from ...modules.layers import FASTConvLayer
 from ...utils import conv_sequence_pt, load_pretrained_params
 
 __all__ = ["textnet_tiny", "textnet_small", "textnet_base"]

doctr/models/classification/vgg/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *
+from .pytorch import *

doctr/models/classification/vip/__init__.py

@@ -1,4 +1 @@
-from doctr.file_utils import is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
+from .pytorch import *

doctr/models/classification/vip/layers/__init__.py

@@ -1,4 +1 @@
-from doctr.file_utils import is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
+from .pytorch import *

doctr/models/classification/vip/layers/pytorch.py

@@ -433,7 +433,7 @@ class LePEAttention(nn.Module):
         Returns:
             A float tensor of shape (b, h, w, c).
         """
-        b_merged = int(img_splits_hw.shape[0] / (h * w / h_sp / w_sp))
+        b_merged = img_splits_hw.shape[0] // ((h * w) // (h_sp * w_sp))
         img = img_splits_hw.view(b_merged, h // h_sp, w // w_sp, h_sp, w_sp, -1)
         # contiguous() required to ensure the tensor has a contiguous memory layout
         # after permute, allowing the subsequent view operation to work correctly.

doctr/models/classification/vit/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *
+from .pytorch import *

doctr/models/classification/vit/pytorch.py

@@ -11,9 +11,9 @@ from torch import nn
 
 from doctr.datasets import VOCABS
 from doctr.models.modules.transformer import EncoderBlock
-from doctr.models.modules.vision_transformer.pytorch import PatchEmbedding
+from doctr.models.modules.vision_transformer import PatchEmbedding
 
-from ...utils.pytorch import load_pretrained_params
+from ...utils import load_pretrained_params
 
 __all__ = ["vit_s", "vit_b"]
 

doctr/models/classification/zoo.py

@@ -5,7 +5,7 @@
 
 from typing import Any
 
-from doctr.file_utils import is_tf_available, is_torch_available
+from doctr.models.utils import _CompiledModule
 
 from .. import classification
 from ..preprocessor import PreProcessor
@@ -30,11 +30,10 @@ ARCHS: list[str] = [
     "vgg16_bn_r",
     "vit_s",
     "vit_b",
+    "vip_tiny",
+    "vip_base",
 ]
 
-if is_torch_available():
-    ARCHS.extend(["vip_tiny", "vip_base"])
-
 ORIENTATION_ARCHS: list[str] = ["mobilenet_v3_small_crop_orientation", "mobilenet_v3_small_page_orientation"]
 
 
@@ -52,12 +51,8 @@ def _orientation_predictor(
         # Load directly classifier from backbone
         _model = classification.__dict__[arch](pretrained=pretrained)
     else:
-        allowed_archs = [classification.MobileNetV3]
-        if is_torch_available():
-            # Adding the type for torch compiled models to the allowed architectures
-            from doctr.models.utils import _CompiledModule
-
-            allowed_archs.append(_CompiledModule)
+        # Adding the type for torch compiled models to the allowed architectures
+        allowed_archs = [classification.MobileNetV3, _CompiledModule]
 
         if not isinstance(arch, tuple(allowed_archs)):
             raise ValueError(f"unknown architecture: {type(arch)}")
@@ -66,7 +61,7 @@ def _orientation_predictor(
     kwargs["mean"] = kwargs.get("mean", _model.cfg["mean"])
     kwargs["std"] = kwargs.get("std", _model.cfg["std"])
     kwargs["batch_size"] = kwargs.get("batch_size", 128 if model_type == "crop" else 4)
-    input_shape = _model.cfg["input_shape"][:-1] if is_tf_available() else _model.cfg["input_shape"][1:]
+    input_shape = _model.cfg["input_shape"][1:]
     predictor = OrientationPredictor(
         PreProcessor(input_shape, preserve_aspect_ratio=True, symmetric_pad=True, **kwargs), _model
     )

doctr/models/detection/_utils/__init__.py

@@ -1,7 +1,2 @@
-from doctr.file_utils import is_tf_available, is_torch_available
 from .base import *
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *
+from .pytorch import *

doctr/models/detection/core.py

@@ -53,7 +53,7 @@ class DetectionPostProcessor(NestedObject):
 
         else:
             mask: np.ndarray = np.zeros((h, w), np.int32)
-            cv2.fillPoly(mask, [points.astype(np.int32)], 1.0)  # type: ignore[call-overload]
+            cv2.fillPoly(mask, [points.astype(np.int32)], 1.0)
             product = pred * mask
             return np.sum(product) / np.count_nonzero(product)
 

doctr/models/detection/differentiable_binarization/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *  # type: ignore[assignment]
+from .pytorch import *

doctr/models/detection/differentiable_binarization/base.py

@@ -224,7 +224,7 @@ class _DBNet:
         padded_polygon: np.ndarray = np.array(padding.Execute(distance)[0])
 
         # Fill the mask with 1 on the new padded polygon
-        cv2.fillPoly(mask, [padded_polygon.astype(np.int32)], 1.0)  # type: ignore[call-overload]
+        cv2.fillPoly(mask, [padded_polygon.astype(np.int32)], 1.0)
 
         # Get min/max to recover polygon after distance computation
         xmin = padded_polygon[:, 0].min()
@@ -269,7 +269,6 @@ class _DBNet:
         self,
         target: list[dict[str, np.ndarray]],
         output_shape: tuple[int, int, int],
-        channels_last: bool = True,
     ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
         if any(t.dtype != np.float32 for tgt in target for t in tgt.values()):
             raise AssertionError("the expected dtype of target 'boxes' entry is 'np.float32'.")
@@ -280,10 +279,8 @@ class _DBNet:
 
         h: int
         w: int
-        if channels_last:
-            h, w, num_classes = output_shape
-        else:
-            num_classes, h, w = output_shape
+
+        num_classes, h, w = output_shape
         target_shape = (len(target), num_classes, h, w)
 
         seg_target: np.ndarray = np.zeros(target_shape, dtype=np.uint8)
@@ -343,17 +340,12 @@ class _DBNet:
                     if shrunken.shape[0] <= 2 or not Polygon(shrunken).is_valid:
                         seg_mask[idx, class_idx, box[1] : box[3] + 1, box[0] : box[2] + 1] = False
                         continue
-                    cv2.fillPoly(seg_target[idx, class_idx], [shrunken.astype(np.int32)], 1.0)  # type: ignore[call-overload]
+                    cv2.fillPoly(seg_target[idx, class_idx], [shrunken.astype(np.int32)], 1.0)
 
                     # Draw on both thresh map and thresh mask
                     poly, thresh_target[idx, class_idx], thresh_mask[idx, class_idx] = self.draw_thresh_map(
                         poly, thresh_target[idx, class_idx], thresh_mask[idx, class_idx]
                     )
-        if channels_last:
-            seg_target = seg_target.transpose((0, 2, 3, 1))
-            seg_mask = seg_mask.transpose((0, 2, 3, 1))
-            thresh_target = thresh_target.transpose((0, 2, 3, 1))
-            thresh_mask = thresh_mask.transpose((0, 2, 3, 1))
 
         thresh_target = thresh_target.astype(input_dtype) * (self.thresh_max - self.thresh_min) + self.thresh_min
 

doctr/models/detection/differentiable_binarization/pytorch.py

@@ -215,7 +215,7 @@ class DBNet(_DBNet, nn.Module):
 
         if target is None or return_preds:
             # Disable for torch.compile compatibility
-            @torch.compiler.disable  # type: ignore[attr-defined]
+            @torch.compiler.disable
             def _postprocess(prob_map: torch.Tensor) -> list[dict[str, Any]]:
                 return [
                     dict(zip(self.class_names, preds))
@@ -261,7 +261,7 @@ class DBNet(_DBNet, nn.Module):
         prob_map = torch.sigmoid(out_map)
         thresh_map = torch.sigmoid(thresh_map)
 
-        targets = self.build_target(target, out_map.shape[1:], False)  # type: ignore[arg-type]
+        targets = self.build_target(target, out_map.shape[1:])  # type: ignore[arg-type]
 
         seg_target, seg_mask = torch.from_numpy(targets[0]), torch.from_numpy(targets[1])
         seg_target, seg_mask = seg_target.to(out_map.device), seg_mask.to(out_map.device)
@@ -285,7 +285,7 @@ class DBNet(_DBNet, nn.Module):
                 dice_map = torch.softmax(out_map, dim=1)
             else:
                 # compute binary map instead
-                dice_map = 1 / (1 + torch.exp(-50.0 * (prob_map - thresh_map)))  # type: ignore[assignment]
+                dice_map = 1 / (1 + torch.exp(-50.0 * (prob_map - thresh_map)))
             # Class reduced
             inter = (seg_mask * dice_map * seg_target).sum((0, 2, 3))
             cardinality = (seg_mask * (dice_map + seg_target)).sum((0, 2, 3))

doctr/models/detection/fast/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *  # type: ignore[assignment]
+from .pytorch import *

doctr/models/detection/fast/base.py

@@ -153,14 +153,12 @@ class _FAST(BaseModel):
         self,
         target: list[dict[str, np.ndarray]],
         output_shape: tuple[int, int, int],
-        channels_last: bool = True,
     ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
         """Build the target, and it's mask to be used from loss computation.
 
         Args:
             target: target coming from dataset
             output_shape: shape of the output of the model without batch_size
-            channels_last: whether channels are last or not
 
         Returns:
             the new formatted target, mask and shrunken text kernel
@@ -172,10 +170,8 @@ class _FAST(BaseModel):
 
         h: int
         w: int
-        if channels_last:
-            h, w, num_classes = output_shape
-        else:
-            num_classes, h, w = output_shape
+
+        num_classes, h, w = output_shape
         target_shape = (len(target), num_classes, h, w)
 
         seg_target: np.ndarray = np.zeros(target_shape, dtype=np.uint8)
@@ -235,14 +231,8 @@ class _FAST(BaseModel):
                     if shrunken.shape[0] <= 2 or not Polygon(shrunken).is_valid:
                         seg_mask[idx, class_idx, box[1] : box[3] + 1, box[0] : box[2] + 1] = False
                         continue
-                    cv2.fillPoly(shrunken_kernel[idx, class_idx], [shrunken.astype(np.int32)], 1.0)  # type: ignore[call-overload]
+                    cv2.fillPoly(shrunken_kernel[idx, class_idx], [shrunken.astype(np.int32)], 1.0)
                     # draw the original polygon on the segmentation target
-                    cv2.fillPoly(seg_target[idx, class_idx], [poly.astype(np.int32)], 1.0)  # type: ignore[call-overload]
-
-        # Don't forget to switch back to channel last if Tensorflow is used
-        if channels_last:
-            seg_target = seg_target.transpose((0, 2, 3, 1))
-            seg_mask = seg_mask.transpose((0, 2, 3, 1))
-            shrunken_kernel = shrunken_kernel.transpose((0, 2, 3, 1))
+                    cv2.fillPoly(seg_target[idx, class_idx], [poly.astype(np.int32)], 1.0)
 
         return seg_target, seg_mask, shrunken_kernel

doctr/models/detection/fast/pytorch.py

@@ -206,7 +206,7 @@ class FAST(_FAST, nn.Module):
 
         if target is None or return_preds:
             # Disable for torch.compile compatibility
-            @torch.compiler.disable  # type: ignore[attr-defined]
+            @torch.compiler.disable
             def _postprocess(prob_map: torch.Tensor) -> list[dict[str, Any]]:
                 return [
                     dict(zip(self.class_names, preds))
@@ -238,7 +238,7 @@ class FAST(_FAST, nn.Module):
         Returns:
             A loss tensor
         """
-        targets = self.build_target(target, out_map.shape[1:], False)  # type: ignore[arg-type]
+        targets = self.build_target(target, out_map.shape[1:])  # type: ignore[arg-type]
 
         seg_target, seg_mask = torch.from_numpy(targets[0]), torch.from_numpy(targets[1])
         shrunken_kernel = torch.from_numpy(targets[2]).to(out_map.device)
@@ -303,7 +303,7 @@ def reparameterize(model: FAST | nn.Module) -> FAST:
 
     for module in model.modules():
         if hasattr(module, "reparameterize_layer"):
-            module.reparameterize_layer()
+            module.reparameterize_layer()  # type: ignore[operator]
 
     for name, child in model.named_children():
         if isinstance(child, nn.BatchNorm2d):
@@ -315,7 +315,7 @@ def reparameterize(model: FAST | nn.Module) -> FAST:
 
             factor = child.weight / torch.sqrt(child.running_var + child.eps)  # type: ignore
             last_conv.weight = nn.Parameter(conv_w * factor.reshape([last_conv.out_channels, 1, 1, 1]))
-            last_conv.bias = nn.Parameter((conv_b - child.running_mean) * factor + child.bias)
+            last_conv.bias = nn.Parameter((conv_b - child.running_mean) * factor + child.bias)  # type: ignore[operator]
             model._modules[last_conv_name] = last_conv  # type: ignore[index]
             model._modules[name] = nn.Identity()
             last_conv = None

doctr/models/detection/linknet/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *  # type: ignore[assignment]
+from .pytorch import *