diffbin_textdetector.py

import keras
from keras import layers

from keras_hub.src.api_export import keras_hub_export
from keras_hub.src.models.diffbin.diffbin_backbone import DiffBinBackbone
from keras_hub.src.models.diffbin.diffbin_preprocessor import (
    DiffBinPreprocessor,
)
from keras_hub.src.models.diffbin.losses import DiffBinLoss
from keras_hub.src.models.image_text_detector_preprocessor import (
    ImageTextDetectorPreprocessor,
)


@keras_hub_export("keras_hub.models.DiffBinImageTextDetector")
class DiffBinImageTextDetector(ImageTextDetectorPreprocessor):
    """Differentiable Binarization scene text detection task.

    `DiffBinImageTextDetector` tasks wrap a `keras_hub.models.DiffBinBackbone`
    and a `keras_hub.models.Preprocessor` to create a model that can be used
    for detecting text in natural images.

    The probability map output generated by `predict()` can be translated into
    polygon representation using `model.postprocess_to_polygons()`.

    Args:
        backbone: A `keras_hub.models.DiffBinBackbone`
            instance.
        preprocessor: `None`, a `keras_hub.models.Preprocessor` instance,
            a `keras.Layer` instance, or a callable. If `None` no preprocessing
            will be applied to the inputs.

    Examples:
    ```python
    input_data = np.ones(shape=(8, 224, 224, 3))

    image_encoder = keras_hub.models.ResNetBackbone.from_preset(
        "resnet_vd_50_imagenet"
    )
    backbone = keras_hub.models.DiffBinBackbone(image_encoder)
    detector = keras_hub.models.DiffBinImageTextDetector(
        backbone=backbone
    )

    map_output = detector(input_data)
    ```

    `map_output` now holds a 8x224x224x3 tensor, where the last dimension
    corresponds to the model's probability map, threshold map and binary map
    outputs. Use `postprocess_to_polygons()` to obtain a polygon
    representation:
    ```python
    detector.postprocess_to_polygons(map_output[...,0])
    ```
    """

    backbone_cls = DiffBinBackbone
    preprocessor_cls = DiffBinPreprocessor

    def __init__(
        self,
        backbone,
        preprocessor=None,
        **kwargs,
    ):
        # === Functional Model ===
        inputs = backbone.input
        x = backbone(inputs)
        probability_maps = x["probability_maps"]
        threshold_maps = x["threshold_maps"]
        binary_maps = step_function(probability_maps, threshold_maps)
        outputs = layers.Concatenate(axis=-1)(
            [probability_maps, threshold_maps, binary_maps]
        )

        super().__init__(inputs=inputs, outputs=outputs, **kwargs)

        # === Config ===
        self.backbone = backbone
        self.preprocessor = preprocessor

    def compile(
        self,
        optimizer="auto",
        loss="auto",
        **kwargs,
    ):
        """Configures the `DiffBinImageTextDetector` task for training.

        `DiffBinImageTextDetector` extends the default compilation signature
        of `keras.Model.compile` with defaults for `optimizer` and `loss`. To
        override these defaults, pass any value to these arguments during
        compilation.

        Args:
            optimizer: `"auto"`, an optimizer name, or a `keras.Optimizer`
                instance. Defaults to `"auto"`, which uses the default
                optimizer for `DiffBinImageTextDetector`. See
                `keras.Model.compile` and `keras.optimizers` for more info on
                possible `optimizer` values.
            loss: `"auto"`, a loss name, or a `keras.losses.Loss` instance.
                Defaults to `"auto"`, in which case the default loss
                computation of `DiffBinImageTextDetector` will be applied.
                See `keras.Model.compile` and `keras.losses` for more info on
                possible `loss` values.
            **kwargs: See `keras.Model.compile` for a full list of arguments
                supported by the compile method.
        """
        if optimizer == "auto":
            # parameters from https://arxiv.org/abs/1911.08947
            optimizer = keras.optimizers.SGD(
                learning_rate=0.007, weight_decay=0.0001, momentum=0.9
            )
        if loss == "auto":
            loss = DiffBinLoss()
        super().compile(
            optimizer=optimizer,
            loss=loss,
            **kwargs,
        )


def step_function(x, y, k=50.0):
    return 1.0 / (1.0 + keras.ops.exp(-k * (x - y)))