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

doctr/models/utils/tensorflow.py

@@ -1,189 +0,0 @@
-# Copyright (C) 2021-2025, 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.
-
-import logging
-from collections.abc import Callable
-from typing import Any
-
-import tensorflow as tf
-import tf2onnx
-import validators
-from tensorflow.keras import Model, layers
-
-from doctr.utils.data import download_from_url
-
-logging.getLogger("tensorflow").setLevel(logging.DEBUG)
-
-
-__all__ = [
-    "load_pretrained_params",
-    "_build_model",
-    "conv_sequence",
-    "IntermediateLayerGetter",
-    "export_model_to_onnx",
-    "_copy_tensor",
-    "_bf16_to_float32",
-]
-
-
-def _copy_tensor(x: tf.Tensor) -> tf.Tensor:
-    return tf.identity(x)
-
-
-def _bf16_to_float32(x: tf.Tensor) -> tf.Tensor:
-    # Convert bfloat16 to float32 for numpy compatibility
-    return tf.cast(x, tf.float32) if x.dtype == tf.bfloat16 else x
-
-
-def _build_model(model: Model):
-    """Build a model by calling it once with dummy input
-
-    Args:
-        model: the model to be built
-    """
-    model(tf.zeros((1, *model.cfg["input_shape"])), training=False)
-
-
-def load_pretrained_params(
-    model: Model,
-    path_or_url: str | None = None,
-    hash_prefix: str | None = None,
-    skip_mismatch: bool = False,
-    **kwargs: Any,
-) -> None:
-    """Load a set of parameters onto a model
-
-    >>> from doctr.models import load_pretrained_params
-    >>> load_pretrained_params(model, "https://yoursource.com/yourcheckpoint-yourhash.weights.h5")
-
-    Args:
-        model: the keras model to be loaded
-        path_or_url: the path or URL to the model parameters (checkpoint)
-        hash_prefix: first characters of SHA256 expected hash
-        skip_mismatch: skip loading layers with mismatched shapes
-        **kwargs: additional arguments to be passed to `doctr.utils.data.download_from_url`
-    """
-    if path_or_url is None:
-        logging.warning("No model URL or Path provided, using default initialization.")
-        return
-
-    archive_path = (
-        download_from_url(path_or_url, hash_prefix=hash_prefix, cache_subdir="models", **kwargs)
-        if validators.url(path_or_url)
-        else path_or_url
-    )
-
-    # Load weights
-    model.load_weights(archive_path, skip_mismatch=skip_mismatch)
-
-
-def conv_sequence(
-    out_channels: int,
-    activation: str | Callable | None = None,
-    bn: bool = False,
-    padding: str = "same",
-    kernel_initializer: str = "he_normal",
-    **kwargs: Any,
-) -> list[layers.Layer]:
-    """Builds a convolutional-based layer sequence
-
-    >>> from tensorflow.keras import Sequential
-    >>> from doctr.models import conv_sequence
-    >>> module = Sequential(conv_sequence(32, 'relu', True, kernel_size=3, input_shape=[224, 224, 3]))
-
-    Args:
-        out_channels: number of output channels
-        activation: activation to be used (default: no activation)
-        bn: should a batch normalization layer be added
-        padding: padding scheme
-        kernel_initializer: kernel initializer
-        **kwargs: additional arguments to be passed to the convolutional layer
-
-    Returns:
-        list of layers
-    """
-    # No bias before Batch norm
-    kwargs["use_bias"] = kwargs.get("use_bias", not bn)
-    # Add activation directly to the conv if there is no BN
-    kwargs["activation"] = activation if not bn else None
-    conv_seq = [layers.Conv2D(out_channels, padding=padding, kernel_initializer=kernel_initializer, **kwargs)]
-
-    if bn:
-        conv_seq.append(layers.BatchNormalization())
-
-    if (isinstance(activation, str) or callable(activation)) and bn:
-        # Activation function can either be a string or a function ('relu' or tf.nn.relu)
-        conv_seq.append(layers.Activation(activation))
-
-    return conv_seq
-
-
-class IntermediateLayerGetter(Model):
-    """Implements an intermediate layer getter
-
-    >>> from tensorflow.keras.applications import ResNet50
-    >>> from doctr.models import IntermediateLayerGetter
-    >>> target_layers = ["conv2_block3_out", "conv3_block4_out", "conv4_block6_out", "conv5_block3_out"]
-    >>> feat_extractor = IntermediateLayerGetter(ResNet50(include_top=False, pooling=False), target_layers)
-
-    Args:
-        model: the model to extract feature maps from
-        layer_names: the list of layers to retrieve the feature map from
-    """
-
-    def __init__(self, model: Model, layer_names: list[str]) -> None:
-        intermediate_fmaps = [model.get_layer(layer_name).get_output_at(0) for layer_name in layer_names]
-        super().__init__(model.input, outputs=intermediate_fmaps)
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}()"
-
-
-def export_model_to_onnx(
-    model: Model, model_name: str, dummy_input: list[tf.TensorSpec], **kwargs: Any
-) -> tuple[str, list[str]]:
-    """Export model to ONNX format.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models.classification import resnet18
-    >>> from doctr.models.utils import export_classification_model_to_onnx
-    >>> model = resnet18(pretrained=True, include_top=True)
-    >>> export_model_to_onnx(model, "my_model",
-    >>> dummy_input=[tf.TensorSpec([None, 32, 32, 3], tf.float32, name="input")])
-
-    Args:
-        model: the keras model to be exported
-        model_name: the name for the exported model
-        dummy_input: the dummy input to the model
-        kwargs: additional arguments to be passed to tf2onnx
-
-    Returns:
-        the path to the exported model and a list with the output layer names
-    """
-    # get the users eager mode
-    eager_mode = tf.executing_eagerly()
-    # set eager mode to true to avoid issues with tf2onnx
-    tf.config.run_functions_eagerly(True)
-    large_model = kwargs.get("large_model", False)
-    model_proto, _ = tf2onnx.convert.from_keras(
-        model,
-        input_signature=dummy_input,
-        output_path=f"{model_name}.zip" if large_model else f"{model_name}.onnx",
-        **kwargs,
-    )
-    # Get the output layer names
-    output = [n.name for n in model_proto.graph.output]
-
-    # reset the eager mode to the users mode
-    tf.config.run_functions_eagerly(eager_mode)
-
-    # models which are too large (weights > 2GB while converting to ONNX) needs to be handled
-    # about an external tensor storage where the graph and weights are seperatly stored in a archive
-    if large_model:
-        logging.info(f"Model exported to {model_name}.zip")
-        return f"{model_name}.zip", output
-
-    logging.info(f"Model exported to {model_name}.zip")
-    return f"{model_name}.onnx", output

doctr/transforms/functional/tensorflow.py

@@ -1,254 +0,0 @@
-# Copyright (C) 2021-2025, 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.
-
-import math
-import random
-from collections.abc import Iterable
-from copy import deepcopy
-
-import numpy as np
-import tensorflow as tf
-
-from doctr.utils.geometry import compute_expanded_shape, rotate_abs_geoms
-
-from .base import create_shadow_mask, crop_boxes
-
-__all__ = ["invert_colors", "rotate_sample", "crop_detection", "random_shadow", "rotated_img_tensor"]
-
-
-def invert_colors(img: tf.Tensor, min_val: float = 0.6) -> tf.Tensor:
-    """Invert the colors of an image
-
-    Args:
-        img : tf.Tensor, the image to invert
-        min_val : minimum value of the random shift
-
-    Returns:
-        the inverted image
-    """
-    out = tf.image.rgb_to_grayscale(img)  # Convert to gray
-    # Random RGB shift
-    shift_shape = [img.shape[0], 1, 1, 3] if img.ndim == 4 else [1, 1, 3]
-    rgb_shift = tf.random.uniform(shape=shift_shape, minval=min_val, maxval=1)
-    # Inverse the color
-    if out.dtype == tf.uint8:
-        out = tf.cast(tf.cast(out, dtype=rgb_shift.dtype) * rgb_shift, dtype=tf.uint8)
-    else:
-        out *= tf.cast(rgb_shift, dtype=out.dtype)
-    # Inverse the color
-    out = 255 - out if out.dtype == tf.uint8 else 1 - out
-    return out
-
-
-def rotated_img_tensor(img: tf.Tensor, angle: float, expand: bool = False) -> tf.Tensor:
-    """Rotate image around the center, interpolation=NEAREST, pad with 0 (black)
-
-    Args:
-        img: image to rotate
-        angle: angle in degrees. +: counter-clockwise, -: clockwise
-        expand: whether the image should be padded before the rotation
-
-    Returns:
-        the rotated image (tensor)
-    """
-    # Compute the expanded padding
-    h_crop, w_crop = 0, 0
-    if expand:
-        exp_h, exp_w = compute_expanded_shape(img.shape[:-1], angle)
-        h_diff, w_diff = int(math.ceil(exp_h - img.shape[0])), int(math.ceil(exp_w - img.shape[1]))
-        h_pad, w_pad = max(h_diff, 0), max(w_diff, 0)
-        exp_img = tf.pad(img, tf.constant([[h_pad // 2, h_pad - h_pad // 2], [w_pad // 2, w_pad - w_pad // 2], [0, 0]]))
-        h_crop, w_crop = int(round(max(exp_img.shape[0] - exp_h, 0))), int(round(min(exp_img.shape[1] - exp_w, 0)))
-    else:
-        exp_img = img
-
-    # Compute the rotation matrix
-    height, width = tf.cast(tf.shape(exp_img)[0], tf.float32), tf.cast(tf.shape(exp_img)[1], tf.float32)
-    cos_angle, sin_angle = tf.math.cos(angle * math.pi / 180.0), tf.math.sin(angle * math.pi / 180.0)
-    x_offset = ((width - 1) - (cos_angle * (width - 1) - sin_angle * (height - 1))) / 2.0
-    y_offset = ((height - 1) - (sin_angle * (width - 1) + cos_angle * (height - 1))) / 2.0
-
-    rotation_matrix = tf.convert_to_tensor(
-        [cos_angle, -sin_angle, x_offset, sin_angle, cos_angle, y_offset, 0.0, 0.0],
-        dtype=tf.float32,
-    )
-    # Rotate the image
-    rotated_img = tf.squeeze(
-        tf.raw_ops.ImageProjectiveTransformV3(
-            images=exp_img[None],  # Add a batch dimension for compatibility with ImageProjectiveTransformV3
-            transforms=rotation_matrix[None],  # Add a batch dimension for compatibility with ImageProjectiveTransformV3
-            output_shape=tf.shape(exp_img)[:2],
-            interpolation="NEAREST",
-            fill_mode="CONSTANT",
-            fill_value=tf.constant(0.0, dtype=tf.float32),
-        )
-    )
-    # Crop the rest
-    if h_crop > 0 or w_crop > 0:
-        h_slice = slice(h_crop // 2, -h_crop // 2) if h_crop > 0 else slice(rotated_img.shape[0])
-        w_slice = slice(-w_crop // 2, -w_crop // 2) if w_crop > 0 else slice(rotated_img.shape[1])
-        rotated_img = rotated_img[h_slice, w_slice]
-
-    return rotated_img
-
-
-def rotate_sample(
-    img: tf.Tensor,
-    geoms: np.ndarray,
-    angle: float,
-    expand: bool = False,
-) -> tuple[tf.Tensor, np.ndarray]:
-    """Rotate image around the center, interpolation=NEAREST, pad with 0 (black)
-
-    Args:
-        img: image to rotate
-        geoms: array of geometries of shape (N, 4) or (N, 4, 2)
-        angle: angle in degrees. +: counter-clockwise, -: clockwise
-        expand: whether the image should be padded before the rotation
-
-    Returns:
-        A tuple of rotated img (tensor), rotated boxes (np array)
-    """
-    # Rotated the image
-    rotated_img = rotated_img_tensor(img, angle, expand)
-
-    # Get absolute coords
-    _geoms = deepcopy(geoms)
-    if _geoms.shape[1:] == (4,):
-        if np.max(_geoms) <= 1:
-            _geoms[:, [0, 2]] *= img.shape[1]
-            _geoms[:, [1, 3]] *= img.shape[0]
-    elif _geoms.shape[1:] == (4, 2):
-        if np.max(_geoms) <= 1:
-            _geoms[..., 0] *= img.shape[1]
-            _geoms[..., 1] *= img.shape[0]
-    else:
-        raise AssertionError
-
-    # Rotate the boxes: xmin, ymin, xmax, ymax or polygons --> (4, 2) polygon
-    rotated_geoms: np.ndarray = rotate_abs_geoms(_geoms, angle, img.shape[:-1], expand).astype(np.float32)
-
-    # Always return relative boxes to avoid label confusions when resizing is performed aferwards
-    rotated_geoms[..., 0] = rotated_geoms[..., 0] / rotated_img.shape[1]
-    rotated_geoms[..., 1] = rotated_geoms[..., 1] / rotated_img.shape[0]
-
-    return rotated_img, np.clip(np.around(rotated_geoms, decimals=15), 0, 1)
-
-
-def crop_detection(
-    img: tf.Tensor, boxes: np.ndarray, crop_box: tuple[float, float, float, float]
-) -> tuple[tf.Tensor, np.ndarray]:
-    """Crop and image and associated bboxes
-
-    Args:
-        img: image to crop
-        boxes: array of boxes to clip, absolute (int) or relative (float)
-        crop_box: box (xmin, ymin, xmax, ymax) to crop the image. Relative coords.
-
-    Returns:
-        A tuple of cropped image, cropped boxes, where the image is not resized.
-    """
-    if any(val < 0 or val > 1 for val in crop_box):
-        raise AssertionError("coordinates of arg `crop_box` should be relative")
-    h, w = img.shape[:2]
-    xmin, ymin = int(round(crop_box[0] * (w - 1))), int(round(crop_box[1] * (h - 1)))
-    xmax, ymax = int(round(crop_box[2] * (w - 1))), int(round(crop_box[3] * (h - 1)))
-    cropped_img = tf.image.crop_to_bounding_box(img, ymin, xmin, ymax - ymin, xmax - xmin)
-    # Crop the box
-    boxes = crop_boxes(boxes, crop_box if boxes.max() <= 1 else (xmin, ymin, xmax, ymax))
-
-    return cropped_img, boxes
-
-
-def _gaussian_filter(
-    img: tf.Tensor,
-    kernel_size: int | Iterable[int],
-    sigma: float,
-    mode: str | None = None,
-    pad_value: int = 0,
-):
-    """Apply Gaussian filter to image.
-    Adapted from: https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/filters.py
-
-    Args:
-        img: image to filter of shape (N, H, W, C)
-        kernel_size: kernel size of the filter
-        sigma: standard deviation of the Gaussian filter
-        mode: padding mode, one of "CONSTANT", "REFLECT", "SYMMETRIC"
-        pad_value: value to pad the image with
-
-    Returns:
-        A tensor of shape (N, H, W, C)
-    """
-    ksize = tf.convert_to_tensor(tf.broadcast_to(kernel_size, [2]), dtype=tf.int32)
-    sigma = tf.convert_to_tensor(tf.broadcast_to(sigma, [2]), dtype=img.dtype)
-    assert mode in ("CONSTANT", "REFLECT", "SYMMETRIC"), "mode should be one of 'CONSTANT', 'REFLECT', 'SYMMETRIC'"
-    mode = "CONSTANT" if mode is None else str.upper(mode)
-    constant_values = (
-        tf.zeros([], dtype=img.dtype) if pad_value is None else tf.convert_to_tensor(pad_value, dtype=img.dtype)
-    )
-
-    def kernel1d(ksize: tf.Tensor, sigma: tf.Tensor, dtype: tf.DType):
-        x = tf.range(ksize, dtype=dtype)
-        x = x - tf.cast(tf.math.floordiv(ksize, 2), dtype=dtype)
-        x = x + tf.where(tf.math.equal(tf.math.mod(ksize, 2), 0), tf.cast(0.5, dtype), 0)
-        g = tf.math.exp(-(tf.math.pow(x, 2) / (2 * tf.math.pow(sigma, 2))))
-        g = g / tf.reduce_sum(g)
-        return g
-
-    def kernel2d(ksize: tf.Tensor, sigma: tf.Tensor, dtype: tf.DType):
-        kernel_x = kernel1d(ksize[0], sigma[0], dtype)
-        kernel_y = kernel1d(ksize[1], sigma[1], dtype)
-        return tf.matmul(
-            tf.expand_dims(kernel_x, axis=-1),
-            tf.transpose(tf.expand_dims(kernel_y, axis=-1)),
-        )
-
-    g = kernel2d(ksize, sigma, img.dtype)
-    # Pad the image
-    height, width = ksize[0], ksize[1]
-    paddings = [
-        [0, 0],
-        [(height - 1) // 2, height - 1 - (height - 1) // 2],
-        [(width - 1) // 2, width - 1 - (width - 1) // 2],
-        [0, 0],
-    ]
-    img = tf.pad(img, paddings, mode=mode, constant_values=constant_values)
-
-    channel = tf.shape(img)[-1]
-    shape = tf.concat([ksize, tf.constant([1, 1], ksize.dtype)], axis=0)
-    g = tf.reshape(g, shape)
-    shape = tf.concat([ksize, [channel], tf.constant([1], ksize.dtype)], axis=0)
-    g = tf.broadcast_to(g, shape)
-    return tf.nn.depthwise_conv2d(img, g, [1, 1, 1, 1], padding="VALID", data_format="NHWC")
-
-
-def random_shadow(img: tf.Tensor, opacity_range: tuple[float, float], **kwargs) -> tf.Tensor:
-    """Apply a random shadow to a given image
-
-    Args:
-        img: image to modify
-        opacity_range: the minimum and maximum desired opacity of the shadow
-        **kwargs: additional arguments to pass to `create_shadow_mask`
-
-    Returns:
-        shadowed image
-    """
-    shadow_mask = create_shadow_mask(img.shape[:2], **kwargs)
-
-    opacity = np.random.uniform(*opacity_range)
-    shadow_tensor = 1 - tf.convert_to_tensor(shadow_mask[..., None], dtype=tf.float32)
-
-    # Add some blur to make it believable
-    k = 7 + int(2 * 4 * random.random())
-    sigma = random.uniform(0.5, 5.0)
-    shadow_tensor = _gaussian_filter(
-        shadow_tensor[tf.newaxis, ...],
-        kernel_size=k,
-        sigma=sigma,
-        mode="REFLECT",
-    )
-
-    return tf.squeeze(opacity * shadow_tensor * img + (1 - opacity) * img, axis=0)