python-doctr 0.7.0__py3-none-any.whl → 0.8.1__py3-none-any.whl

doctr/models/detection/differentiable_binarization/tensorflow.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, Mindee.
 
 # 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.
@@ -15,7 +15,7 @@ from tensorflow.keras import layers
 from tensorflow.keras.applications import ResNet50
 
 from doctr.file_utils import CLASS_NAME
-from doctr.models.utils import IntermediateLayerGetter, conv_sequence, load_pretrained_params
+from doctr.models.utils import IntermediateLayerGetter, _bf16_to_float32, conv_sequence, load_pretrained_params
 from doctr.utils.repr import NestedObject
 
 from ...classification import mobilenet_v3_large
@@ -29,13 +29,13 @@ default_cfgs: Dict[str, Dict[str, Any]] = {
         "mean": (0.798, 0.785, 0.772),
         "std": (0.264, 0.2749, 0.287),
         "input_shape": (1024, 1024, 3),
-        "url": "https://doctr-static.mindee.com/models?id=v0.2.0/db_resnet50-adcafc63.zip&src=0",
+        "url": "https://doctr-static.mindee.com/models?id=v0.7.0/db_resnet50-84171458.zip&src=0",
     },
     "db_mobilenet_v3_large": {
         "mean": (0.798, 0.785, 0.772),
         "std": (0.264, 0.2749, 0.287),
         "input_shape": (1024, 1024, 3),
-        "url": "https://doctr-static.mindee.com/models?id=v0.3.1/db_mobilenet_v3_large-8c16d5bf.zip&src=0",
+        "url": "https://doctr-static.mindee.com/models?id=v0.7.0/db_mobilenet_v3_large-da524564.zip&src=0",
     },
 }
 
@@ -45,6 +45,7 @@ class FeaturePyramidNetwork(layers.Layer, NestedObject):
     <https://arxiv.org/pdf/1612.03144.pdf>`_.
 
     Args:
+    ----
         channels: number of channel to output
     """
 
@@ -66,14 +67,15 @@ class FeaturePyramidNetwork(layers.Layer, NestedObject):
         """Module which performs a 3x3 convolution followed by up-sampling
 
         Args:
+        ----
             channels: number of output channels
             dilation_factor (int): dilation factor to scale the convolution output before concatenation
 
         Returns:
+        -------
             a keras.layers.Layer object, wrapping these operations in a sequential module
 
         """
-
         _layers = conv_sequence(channels, "relu", True, kernel_size=3)
 
         if dilation_factor > 1:
@@ -107,8 +109,11 @@ class DBNet(_DBNet, keras.Model, NestedObject):
     <https://arxiv.org/pdf/1911.08947.pdf>`_.
 
     Args:
+    ----
         feature extractor: the backbone serving as feature extractor
         fpn_channels: number of channels each extracted feature maps is mapped to
+        bin_thresh: threshold for binarization
+        box_thresh: minimal objectness score to consider a box
         assume_straight_pages: if True, fit straight bounding boxes only
         exportable: onnx exportable returns only logits
         cfg: the configuration dict of the model
@@ -122,6 +127,7 @@ class DBNet(_DBNet, keras.Model, NestedObject):
         feature_extractor: IntermediateLayerGetter,
         fpn_channels: int = 128,  # to be set to 256 to represent the author's initial idea
         bin_thresh: float = 0.3,
+        box_thresh: float = 0.1,
         assume_straight_pages: bool = True,
         exportable: bool = False,
         cfg: Optional[Dict[str, Any]] = None,
@@ -141,87 +147,96 @@ class DBNet(_DBNet, keras.Model, NestedObject):
         _inputs = [layers.Input(shape=in_shape[1:]) for in_shape in self.feat_extractor.output_shape]
         output_shape = tuple(self.fpn(_inputs).shape)
 
-        self.probability_head = keras.Sequential(
-            [
-                *conv_sequence(64, "relu", True, kernel_size=3, input_shape=output_shape[1:]),
-                layers.Conv2DTranspose(64, 2, strides=2, use_bias=False, kernel_initializer="he_normal"),
-                layers.BatchNormalization(),
-                layers.Activation("relu"),
-                layers.Conv2DTranspose(num_classes, 2, strides=2, kernel_initializer="he_normal"),
-            ]
-        )
-        self.threshold_head = keras.Sequential(
-            [
-                *conv_sequence(64, "relu", True, kernel_size=3, input_shape=output_shape[1:]),
-                layers.Conv2DTranspose(64, 2, strides=2, use_bias=False, kernel_initializer="he_normal"),
-                layers.BatchNormalization(),
-                layers.Activation("relu"),
-                layers.Conv2DTranspose(num_classes, 2, strides=2, kernel_initializer="he_normal"),
-            ]
+        self.probability_head = keras.Sequential([
+            *conv_sequence(64, "relu", True, kernel_size=3, input_shape=output_shape[1:]),
+            layers.Conv2DTranspose(64, 2, strides=2, use_bias=False, kernel_initializer="he_normal"),
+            layers.BatchNormalization(),
+            layers.Activation("relu"),
+            layers.Conv2DTranspose(num_classes, 2, strides=2, kernel_initializer="he_normal"),
+        ])
+        self.threshold_head = keras.Sequential([
+            *conv_sequence(64, "relu", True, kernel_size=3, input_shape=output_shape[1:]),
+            layers.Conv2DTranspose(64, 2, strides=2, use_bias=False, kernel_initializer="he_normal"),
+            layers.BatchNormalization(),
+            layers.Activation("relu"),
+            layers.Conv2DTranspose(num_classes, 2, strides=2, kernel_initializer="he_normal"),
+        ])
+
+        self.postprocessor = DBPostProcessor(
+            assume_straight_pages=assume_straight_pages, bin_thresh=bin_thresh, box_thresh=box_thresh
         )
 
-        self.postprocessor = DBPostProcessor(assume_straight_pages=assume_straight_pages, bin_thresh=bin_thresh)
-
     def compute_loss(
         self,
         out_map: tf.Tensor,
         thresh_map: tf.Tensor,
         target: List[Dict[str, np.ndarray]],
+        gamma: float = 2.0,
+        alpha: float = 0.5,
+        eps: float = 1e-8,
     ) -> tf.Tensor:
         """Compute a batch of gts, masks, thresh_gts, thresh_masks from a list of boxes
         and a list of masks for each image. From there it computes the loss with the model output
 
         Args:
+        ----
             out_map: output feature map of the model of shape (N, H, W, C)
             thresh_map: threshold map of shape (N, H, W, C)
             target: list of dictionary where each dict has a `boxes` and a `flags` entry
+            gamma: modulating factor in the focal loss formula
+            alpha: balancing factor in the focal loss formula
+            eps: epsilon factor in dice loss
 
         Returns:
+        -------
             A loss tensor
         """
+        if gamma < 0:
+            raise ValueError("Value of gamma should be greater than or equal to zero.")
 
         prob_map = tf.math.sigmoid(out_map)
         thresh_map = tf.math.sigmoid(thresh_map)
 
-        seg_target, seg_mask, thresh_target, thresh_mask = self.build_target(target, out_map.shape, True)
+        seg_target, seg_mask, thresh_target, thresh_mask = self.build_target(target, out_map.shape[1:], True)
         seg_target = tf.convert_to_tensor(seg_target, dtype=out_map.dtype)
         seg_mask = tf.convert_to_tensor(seg_mask, dtype=tf.bool)
+        seg_mask = tf.cast(seg_mask, tf.float32)
         thresh_target = tf.convert_to_tensor(thresh_target, dtype=out_map.dtype)
         thresh_mask = tf.convert_to_tensor(thresh_mask, dtype=tf.bool)
 
-        # Compute balanced BCE loss for proba_map
-        bce_scale = 5.0
-        bce_loss = tf.keras.losses.binary_crossentropy(
-            seg_target[..., None],
-            out_map[..., None],
-            from_logits=True,
-        )[seg_mask]
-
-        neg_target = 1 - seg_target[seg_mask]
-        positive_count = tf.math.reduce_sum(seg_target[seg_mask])
-        negative_count = tf.math.reduce_min([tf.math.reduce_sum(neg_target), 3.0 * positive_count])
-        negative_loss = bce_loss * neg_target
-        negative_loss, _ = tf.nn.top_k(negative_loss, tf.cast(negative_count, tf.int32))
-        sum_losses = tf.math.reduce_sum(bce_loss * seg_target[seg_mask]) + tf.math.reduce_sum(negative_loss)
-        balanced_bce_loss = sum_losses / (positive_count + negative_count + 1e-6)
-
-        # Compute dice loss for approxbin_map
-        bin_map = 1 / (1 + tf.exp(-50.0 * (prob_map[seg_mask] - thresh_map[seg_mask])))
-
-        bce_min = tf.math.reduce_min(bce_loss)
-        weights = (bce_loss - bce_min) / (tf.math.reduce_max(bce_loss) - bce_min) + 1.0
-        inter = tf.math.reduce_sum(bin_map * seg_target[seg_mask] * weights)
-        union = tf.math.reduce_sum(bin_map) + tf.math.reduce_sum(seg_target[seg_mask]) + 1e-8
-        dice_loss = 1 - 2.0 * inter / union
+        # Focal loss
+        focal_scale = 10.0
+        bce_loss = tf.keras.losses.binary_crossentropy(seg_target[..., None], out_map[..., None], from_logits=True)
+
+        # Convert logits to prob, compute gamma factor
+        p_t = (seg_target * prob_map) + ((1 - seg_target) * (1 - prob_map))
+        alpha_t = seg_target * alpha + (1 - seg_target) * (1 - alpha)
+        # Unreduced loss
+        focal_loss = alpha_t * (1 - p_t) ** gamma * bce_loss
+        # Class reduced
+        focal_loss = tf.reduce_sum(seg_mask * focal_loss, (0, 1, 2, 3)) / tf.reduce_sum(seg_mask, (0, 1, 2, 3))
+
+        # Compute dice loss for each class or for approx binary_map
+        if len(self.class_names) > 1:
+            dice_map = tf.nn.softmax(out_map, axis=-1)
+        else:
+            # compute binary map instead
+            dice_map = 1.0 / (1.0 + tf.exp(-50 * (prob_map - thresh_map)))
+        # Class-reduced dice loss
+        inter = tf.reduce_sum(seg_mask * dice_map * seg_target, axis=[0, 1, 2])
+        cardinality = tf.reduce_sum(seg_mask * (dice_map + seg_target), axis=[0, 1, 2])
+        dice_loss = tf.reduce_mean(1 - 2 * inter / (cardinality + eps))
 
         # Compute l1 loss for thresh_map
-        l1_scale = 10.0
         if tf.reduce_any(thresh_mask):
-            l1_loss = tf.math.reduce_mean(tf.math.abs(thresh_map[thresh_mask] - thresh_target[thresh_mask]))
+            thresh_mask = tf.cast(thresh_mask, tf.float32)
+            l1_loss = tf.reduce_sum(tf.abs(thresh_map - thresh_target) * thresh_mask) / (
+                tf.reduce_sum(thresh_mask) + eps
+            )
         else:
             l1_loss = tf.constant(0.0)
 
-        return l1_scale * l1_loss + bce_scale * balanced_bce_loss + dice_loss
+        return l1_loss + focal_scale * focal_loss + dice_loss
 
     def call(
         self,
@@ -241,7 +256,7 @@ class DBNet(_DBNet, keras.Model, NestedObject):
             return out
 
         if return_model_output or target is None or return_preds:
-            prob_map = tf.math.sigmoid(logits)
+            prob_map = _bf16_to_float32(tf.math.sigmoid(logits))
 
         if return_model_output:
             out["out_map"] = prob_map
@@ -342,12 +357,14 @@ def db_resnet50(pretrained: bool = False, **kwargs: Any) -> DBNet:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text detection dataset
+        **kwargs: keyword arguments of the DBNet architecture
 
     Returns:
+    -------
         text detection architecture
     """
-
     return _db_resnet(
         "db_resnet50",
         pretrained,
@@ -368,12 +385,14 @@ def db_mobilenet_v3_large(pretrained: bool = False, **kwargs: Any) -> DBNet:
     >>> out = model(input_tensor)
 
     Args:
+    ----
         pretrained (bool): If True, returns a model pre-trained on our text detection dataset
+        **kwargs: keyword arguments of the DBNet architecture
 
     Returns:
+    -------
         text detection architecture
     """
-
     return _db_mobilenet(
         "db_mobilenet_v3_large",
         pretrained,

doctr/models/detection/fast/__init__.py

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

doctr/models/detection/fast/base.py

@@ -0,0 +1,256 @@
+# Copyright (C) 2021-2024, Mindee.
+
+# 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.
+
+# Credits: post-processing adapted from https://github.com/xuannianz/DifferentiableBinarization
+
+from typing import Dict, List, Tuple, Union
+
+import cv2
+import numpy as np
+import pyclipper
+from shapely.geometry import Polygon
+
+from doctr.models.core import BaseModel
+
+from ..core import DetectionPostProcessor
+
+__all__ = ["_FAST", "FASTPostProcessor"]
+
+
+class FASTPostProcessor(DetectionPostProcessor):
+    """Implements a post processor for FAST model.
+
+    Args:
+    ----
+        bin_thresh: threshold used to binzarized p_map at inference time
+        box_thresh: minimal objectness score to consider a box
+        assume_straight_pages: whether the inputs were expected to have horizontal text elements
+    """
+
+    def __init__(
+        self,
+        bin_thresh: float = 0.3,
+        box_thresh: float = 0.1,
+        assume_straight_pages: bool = True,
+    ) -> None:
+        super().__init__(box_thresh, bin_thresh, assume_straight_pages)
+        self.unclip_ratio = 1.0
+
+    def polygon_to_box(
+        self,
+        points: np.ndarray,
+    ) -> np.ndarray:
+        """Expand a polygon (points) by a factor unclip_ratio, and returns a polygon
+
+        Args:
+        ----
+            points: The first parameter.
+
+        Returns:
+        -------
+            a box in absolute coordinates (xmin, ymin, xmax, ymax) or (4, 2) array (quadrangle)
+        """
+        if not self.assume_straight_pages:
+            # Compute the rectangle polygon enclosing the raw polygon
+            rect = cv2.minAreaRect(points)
+            points = cv2.boxPoints(rect)
+            # Add 1 pixel to correct cv2 approx
+            area = (rect[1][0] + 1) * (1 + rect[1][1])
+            length = 2 * (rect[1][0] + rect[1][1]) + 2
+        else:
+            poly = Polygon(points)
+            area = poly.area
+            length = poly.length
+        distance = area * self.unclip_ratio / length  # compute distance to expand polygon
+        offset = pyclipper.PyclipperOffset()
+        offset.AddPath(points, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+        _points = offset.Execute(distance)
+        # Take biggest stack of points
+        idx = 0
+        if len(_points) > 1:
+            max_size = 0
+            for _idx, p in enumerate(_points):
+                if len(p) > max_size:
+                    idx = _idx
+                    max_size = len(p)
+            # We ensure that _points can be correctly casted to a ndarray
+            _points = [_points[idx]]
+        expanded_points: np.ndarray = np.asarray(_points)  # expand polygon
+        if len(expanded_points) < 1:
+            return None  # type: ignore[return-value]
+        return (
+            cv2.boundingRect(expanded_points)  # type: ignore[return-value]
+            if self.assume_straight_pages
+            else np.roll(cv2.boxPoints(cv2.minAreaRect(expanded_points)), -1, axis=0)
+        )
+
+    def bitmap_to_boxes(
+        self,
+        pred: np.ndarray,
+        bitmap: np.ndarray,
+    ) -> np.ndarray:
+        """Compute boxes from a bitmap/pred_map: find connected components then filter boxes
+
+        Args:
+        ----
+            pred: Pred map from differentiable linknet output
+            bitmap: Bitmap map computed from pred (binarized)
+            angle_tol: Comparison tolerance of the angle with the median angle across the page
+            ratio_tol: Under this limit aspect ratio, we cannot resolve the direction of the crop
+
+        Returns:
+        -------
+            np tensor boxes for the bitmap, each box is a 6-element list
+                containing x, y, w, h, alpha, score for the box
+        """
+        height, width = bitmap.shape[:2]
+        boxes: List[Union[np.ndarray, List[float]]] = []
+        # get contours from connected components on the bitmap
+        contours, _ = cv2.findContours(bitmap.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        for contour in contours:
+            # Check whether smallest enclosing bounding box is not too small
+            if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < 2):
+                continue
+            # Compute objectness
+            if self.assume_straight_pages:
+                x, y, w, h = cv2.boundingRect(contour)
+                points: np.ndarray = np.array([[x, y], [x, y + h], [x + w, y + h], [x + w, y]])
+                score = self.box_score(pred, points, assume_straight_pages=True)
+            else:
+                score = self.box_score(pred, contour, assume_straight_pages=False)
+
+            if score < self.box_thresh:  # remove polygons with a weak objectness
+                continue
+
+            if self.assume_straight_pages:
+                _box = self.polygon_to_box(points)
+            else:
+                _box = self.polygon_to_box(np.squeeze(contour))
+
+            if self.assume_straight_pages:
+                # compute relative polygon to get rid of img shape
+                x, y, w, h = _box
+                xmin, ymin, xmax, ymax = x / width, y / height, (x + w) / width, (y + h) / height
+                boxes.append([xmin, ymin, xmax, ymax, score])
+            else:
+                # compute relative box to get rid of img shape
+                _box[:, 0] /= width
+                _box[:, 1] /= height
+                boxes.append(_box)
+
+        if not self.assume_straight_pages:
+            return np.clip(np.asarray(boxes), 0, 1) if len(boxes) > 0 else np.zeros((0, 4, 2), dtype=pred.dtype)
+        else:
+            return np.clip(np.asarray(boxes), 0, 1) if len(boxes) > 0 else np.zeros((0, 5), dtype=pred.dtype)
+
+
+class _FAST(BaseModel):
+    """FAST as described in `"FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation"
+    <https://arxiv.org/pdf/2111.02394.pdf>`_.
+    """
+
+    min_size_box: int = 3
+    assume_straight_pages: bool = True
+    shrink_ratio = 0.1
+
+    def build_target(
+        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
+        """
+        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'.")
+        if any(np.any((t[:, :4] > 1) | (t[:, :4] < 0)) for tgt in target for t in tgt.values()):
+            raise ValueError("the 'boxes' entry of the target is expected to take values between 0 & 1.")
+
+        h: int
+        w: int
+        if channels_last:
+            h, w, num_classes = output_shape
+        else:
+            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)
+        seg_mask: np.ndarray = np.ones(target_shape, dtype=bool)
+        shrunken_kernel: np.ndarray = np.zeros(target_shape, dtype=np.uint8)
+
+        for idx, tgt in enumerate(target):
+            for class_idx, _tgt in enumerate(tgt.values()):
+                # Draw each polygon on gt
+                if _tgt.shape[0] == 0:
+                    # Empty image, full masked
+                    seg_mask[idx, class_idx] = False
+
+                # Absolute bounding boxes
+                abs_boxes = _tgt.copy()
+
+                if abs_boxes.ndim == 3:
+                    abs_boxes[:, :, 0] *= w
+                    abs_boxes[:, :, 1] *= h
+                    polys = abs_boxes
+                    boxes_size = np.linalg.norm(abs_boxes[:, 2, :] - abs_boxes[:, 0, :], axis=-1)
+                    abs_boxes = np.concatenate((abs_boxes.min(1), abs_boxes.max(1)), -1).round().astype(np.int32)
+                else:
+                    abs_boxes[:, [0, 2]] *= w
+                    abs_boxes[:, [1, 3]] *= h
+                    abs_boxes = abs_boxes.round().astype(np.int32)
+                    polys = np.stack(
+                        [
+                            abs_boxes[:, [0, 1]],
+                            abs_boxes[:, [0, 3]],
+                            abs_boxes[:, [2, 3]],
+                            abs_boxes[:, [2, 1]],
+                        ],
+                        axis=1,
+                    )
+                    boxes_size = np.minimum(abs_boxes[:, 2] - abs_boxes[:, 0], abs_boxes[:, 3] - abs_boxes[:, 1])
+
+                for poly, box, box_size in zip(polys, abs_boxes, boxes_size):
+                    # Mask boxes that are too small
+                    if box_size < self.min_size_box:
+                        seg_mask[idx, class_idx, box[1] : box[3] + 1, box[0] : box[2] + 1] = False
+                        continue
+
+                    # Negative shrink for gt, as described in paper
+                    polygon = Polygon(poly)
+                    distance = polygon.area * (1 - np.power(self.shrink_ratio, 2)) / polygon.length
+                    subject = [tuple(coor) for coor in poly]
+                    padding = pyclipper.PyclipperOffset()
+                    padding.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+                    shrunken = padding.Execute(-distance)
+
+                    # Draw polygon on gt if it is valid
+                    if len(shrunken) == 0:
+                        seg_mask[idx, class_idx, box[1] : box[3] + 1, box[0] : box[2] + 1] = False
+                        continue
+                    shrunken = np.array(shrunken[0]).reshape(-1, 2)
+                    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]
+                    # 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))
+
+        return seg_target, seg_mask, shrunken_kernel