rebase and update

docs/source/using_doctr/using_model_export.rst

@@ -31,7 +31,7 @@ Advantages:
         .. code:: python3
 
             import tensorflow as tf
-            from keras import mixed_precision
+            from tensorflow.keras import mixed_precision
             mixed_precision.set_global_policy('mixed_float16')
             predictor = ocr_predictor(reco_arch="crnn_mobilenet_v3_small", det_arch="linknet_resnet34", pretrained=True)
 

doctr/file_utils.py

@@ -35,7 +35,7 @@
     logging.info("Disabling PyTorch because USE_TF is set")
     _torch_available = False
 
-# Compatibility fix to make sure tf.keras stays at Keras 2
+# 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"
 

doctr/io/image/tensorflow.py

@@ -7,8 +7,8 @@
 
 import numpy as np
 import tensorflow as tf
-from keras.utils import img_to_array
 from PIL import Image
+from tensorflow.keras.utils import img_to_array
 
 from doctr.utils.common_types import AbstractPath
 

doctr/models/classification/magc_resnet/tensorflow.py

@@ -9,8 +9,8 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import tensorflow as tf
-from keras import activations, layers
-from keras.models import Sequential
+from tensorflow.keras import activations, layers
+from tensorflow.keras.models import Sequential
 
 from doctr.datasets import VOCABS
 

doctr/models/classification/mobilenet/tensorflow.py

@@ -9,8 +9,8 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import tensorflow as tf
-from keras import layers
-from keras.models import Sequential
+from tensorflow.keras import layers
+from tensorflow.keras.models import Sequential
 
 from ....datasets import VOCABS
 from ...utils import conv_sequence, load_pretrained_params

doctr/models/classification/predictor/tensorflow.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model
+from tensorflow.keras import Model
 
 from doctr.models.preprocessor import PreProcessor
 from doctr.utils.repr import NestedObject

doctr/models/classification/resnet/tensorflow.py

@@ -7,9 +7,9 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import tensorflow as tf
-from keras import layers
-from keras.applications import ResNet50
-from keras.models import Sequential
+from tensorflow.keras import layers
+from tensorflow.keras.applications import ResNet50
+from tensorflow.keras.models import Sequential
 
 from doctr.datasets import VOCABS
 

doctr/models/classification/textnet/tensorflow.py

@@ -7,7 +7,7 @@
 from copy import deepcopy
 from typing import Any, Dict, List, Optional, Tuple
 
-from keras import Sequential, layers
+from tensorflow.keras import Sequential, layers
 
 from doctr.datasets import VOCABS
 

doctr/models/classification/vgg/tensorflow.py

@@ -6,8 +6,8 @@
 from copy import deepcopy
 from typing import Any, Dict, List, Optional, Tuple
 
-from keras import layers
-from keras.models import Sequential
+from tensorflow.keras import layers
+from tensorflow.keras.models import Sequential
 
 from doctr.datasets import VOCABS
 

doctr/models/classification/vit/tensorflow.py

@@ -7,7 +7,7 @@
 from typing import Any, Dict, Optional, Tuple
 
 import tensorflow as tf
-from keras import Sequential, layers
+from tensorflow.keras import Sequential, layers
 
 from doctr.datasets import VOCABS
 from doctr.models.modules.transformer import EncoderBlock

doctr/models/detection/differentiable_binarization/tensorflow.py

@@ -10,8 +10,8 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, Sequential, layers, losses
-from keras.applications import ResNet50
+from tensorflow.keras import Model, Sequential, layers, losses
+from tensorflow.keras.applications import ResNet50
 
 from doctr.file_utils import CLASS_NAME
 from doctr.models.utils import IntermediateLayerGetter, _bf16_to_float32, conv_sequence, load_pretrained_params

doctr/models/detection/fast/tensorflow.py

@@ -10,7 +10,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, Sequential, layers
+from tensorflow.keras import Model, Sequential, layers
 
 from doctr.file_utils import CLASS_NAME
 from doctr.models.utils import IntermediateLayerGetter, _bf16_to_float32, load_pretrained_params

doctr/models/detection/linknet/tensorflow.py

@@ -10,7 +10,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, Sequential, layers, losses
+from tensorflow.keras import Model, Sequential, layers, losses
 
 from doctr.file_utils import CLASS_NAME
 from doctr.models.classification import resnet18, resnet34, resnet50

doctr/models/detection/predictor/tensorflow.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model
+from tensorflow.keras import Model
 
 from doctr.models.detection._utils import _remove_padding
 from doctr.models.preprocessor import PreProcessor

doctr/models/modules/layers/tensorflow.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import layers
+from tensorflow.keras import layers
 
 from doctr.utils.repr import NestedObject
 

doctr/models/modules/transformer/tensorflow.py

@@ -7,7 +7,7 @@
 from typing import Any, Callable, Optional, Tuple
 
 import tensorflow as tf
-from keras import layers
+from tensorflow.keras import layers
 
 from doctr.utils.repr import NestedObject
 

doctr/models/modules/vision_transformer/tensorflow.py

@@ -7,7 +7,7 @@
 from typing import Any, Tuple
 
 import tensorflow as tf
-from keras import layers
+from tensorflow.keras import layers
 
 from doctr.utils.repr import NestedObject
 

doctr/models/recognition/crnn/tensorflow.py

@@ -7,8 +7,8 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import tensorflow as tf
-from keras import layers
-from keras.models import Model, Sequential
+from tensorflow.keras import layers
+from tensorflow.keras.models import Model, Sequential
 
 from doctr.datasets import VOCABS
 

doctr/models/recognition/master/tensorflow.py

@@ -7,7 +7,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import tensorflow as tf
-from keras import Model, layers
+from tensorflow.keras import Model, layers
 
 from doctr.datasets import VOCABS
 from doctr.models.classification import magc_resnet31

doctr/models/recognition/parseq/tensorflow.py

@@ -10,7 +10,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, layers
+from tensorflow.keras import Model, layers
 
 from doctr.datasets import VOCABS
 from doctr.models.modules.transformer import MultiHeadAttention, PositionwiseFeedForward

doctr/models/recognition/sar/tensorflow.py

@@ -7,7 +7,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import tensorflow as tf
-from keras import Model, Sequential, layers
+from tensorflow.keras import Model, Sequential, layers
 
 from doctr.datasets import VOCABS
 from doctr.utils.repr import NestedObject

doctr/models/recognition/vitstr/tensorflow.py

@@ -7,7 +7,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import tensorflow as tf
-from keras import Model, layers
+from tensorflow.keras import Model, layers
 
 from doctr.datasets import VOCABS
 

doctr/models/utils/tensorflow.py

@@ -8,7 +8,7 @@
 
 import tensorflow as tf
 import tf2onnx
-from keras import Model, layers
+from tensorflow.keras import Model, layers
 
 from doctr.utils.data import download_from_url
 
@@ -77,7 +77,7 @@ def conv_sequence(
 ) -> List[layers.Layer]:
     """Builds a convolutional-based layer sequence
 
-    >>> from keras import Sequential
+    >>> 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]))
 
@@ -113,7 +113,7 @@ def conv_sequence(
 class IntermediateLayerGetter(Model):
     """Implements an intermediate layer getter
 
-    >>> from keras.applications import ResNet50
+    >>> 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)

pyproject.toml

@@ -56,8 +56,6 @@ tf = [
     "tensorflow>=2.15.0,<3.0.0",
     "tf-keras>=2.15.0,<3.0.0",  # Keep keras 2 compatibility
     "tf2onnx>=1.16.0,<2.0.0",  # cf. https://github.com/onnx/tensorflow-onnx/releases/tag/v1.16.0
-    # TODO: This is a temporary fix until we can upgrade to a newer version of tensorflow
-    "numpy>=1.16.0,<2.0.0",
 ]
 torch = [
     "torch>=1.12.0,<3.0.0",

references/classification/train_tensorflow_character.py

@@ -17,7 +17,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, mixed_precision, optimizers
+from tensorflow.keras import Model, mixed_precision, optimizers
 from tqdm.auto import tqdm
 
 from doctr.models import login_to_hub, push_to_hf_hub

references/classification/train_tensorflow_orientation.py

@@ -17,7 +17,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, mixed_precision, optimizers
+from tensorflow.keras import Model, mixed_precision, optimizers
 from tqdm.auto import tqdm
 
 from doctr.models import login_to_hub, push_to_hf_hub

references/detection/evaluate_tensorflow.py

@@ -18,7 +18,7 @@
 from pathlib import Path
 
 import tensorflow as tf
-from keras import mixed_precision
+from tensorflow.keras import mixed_precision
 from tqdm import tqdm
 
 gpu_devices = tf.config.experimental.list_physical_devices("GPU")

references/detection/train_tensorflow.py

@@ -18,7 +18,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, mixed_precision, optimizers
+from tensorflow.keras import Model, mixed_precision, optimizers
 from tqdm.auto import tqdm
 
 from doctr.models import login_to_hub, push_to_hf_hub

references/recognition/evaluate_tensorflow.py

@@ -15,7 +15,7 @@
 import time
 
 import tensorflow as tf
-from keras import mixed_precision
+from tensorflow.keras import mixed_precision
 from tqdm import tqdm
 
 gpu_devices = tf.config.experimental.list_physical_devices("GPU")

references/recognition/train_tensorflow.py

@@ -19,7 +19,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, mixed_precision, optimizers
+from tensorflow.keras import Model, mixed_precision, optimizers
 from tqdm.auto import tqdm
 
 from doctr.models import login_to_hub, push_to_hf_hub

tests/tensorflow/test_models_classification_tf.py

@@ -2,7 +2,6 @@
 import tempfile
 
 import cv2
-import keras
 import numpy as np
 import onnxruntime
 import psutil
@@ -38,7 +37,7 @@
 def test_classification_architectures(arch_name, input_shape, output_size):
     # Model
     batch_size = 2
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     model = classification.__dict__[arch_name](pretrained=True, include_top=True, input_shape=input_shape)
     # Forward
     out = model(tf.random.uniform(shape=[batch_size, *input_shape], maxval=1, dtype=tf.float32))
@@ -47,7 +46,7 @@ def test_classification_architectures(arch_name, input_shape, output_size):
     assert out.dtype == tf.float32
     assert out.numpy().shape == (batch_size, *output_size)
     # Check that you can load pretrained up to the classification layer with differing number of classes to fine-tune
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     assert classification.__dict__[arch_name](
         pretrained=True, include_top=True, input_shape=input_shape, num_classes=10
     )
@@ -63,7 +62,7 @@ def test_classification_architectures(arch_name, input_shape, output_size):
 def test_classification_models(arch_name, input_shape):
     batch_size = 8
     reco_model = classification.__dict__[arch_name](pretrained=True, input_shape=input_shape)
-    assert isinstance(reco_model, keras.Model)
+    assert isinstance(reco_model, tf.keras.Model)
     input_tensor = tf.random.uniform(shape=[batch_size, *input_shape], minval=0, maxval=1)
 
     out = reco_model(input_tensor)
@@ -232,7 +231,7 @@ def test_page_orientation_model(mock_payslip):
 def test_models_onnx_export(arch_name, input_shape, output_size):
     # Model
     batch_size = 2
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     if "orientation" in arch_name:
         model = classification.__dict__[arch_name](pretrained=True, input_shape=input_shape)
     else:

tests/tensorflow/test_models_detection_tf.py

@@ -2,7 +2,6 @@
 import os
 import tempfile
 
-import keras
 import numpy as np
 import onnxruntime
 import psutil
@@ -38,13 +37,13 @@
 )
 def test_detection_models(arch_name, input_shape, output_size, out_prob, train_mode):
     batch_size = 2
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     if arch_name == "fast_tiny_rep":
         model = reparameterize(detection.fast_tiny(pretrained=True, input_shape=input_shape))
         train_mode = False  # Reparameterized model is not trainable
     else:
         model = detection.__dict__[arch_name](pretrained=True, input_shape=input_shape)
-    assert isinstance(model, keras.Model)
+    assert isinstance(model, tf.keras.Model)
     input_tensor = tf.random.uniform(shape=[batch_size, *input_shape], minval=0, maxval=1)
     target = [
         {CLASS_NAME: np.array([[0.5, 0.5, 1, 1], [0.5, 0.5, 0.8, 0.8]], dtype=np.float32)},
@@ -153,7 +152,7 @@ def test_rotated_detectionpredictor(mock_pdf):
 )
 def test_detection_zoo(arch_name):
     # Model
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     predictor = detection.zoo.detection_predictor(arch_name, pretrained=False)
     # object check
     assert isinstance(predictor, DetectionPredictor)
@@ -178,7 +177,7 @@ def test_fast_reparameterization():
     base_model_params = np.sum([np.prod(v.shape) for v in base_model.trainable_variables])
     assert math.isclose(base_model_params, 13535296)  # base model params
     base_out = base_model(dummy_input, training=False)["logits"]
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     rep_model = reparameterize(base_model)
     rep_model_params = np.sum([np.prod(v.shape) for v in base_model.trainable_variables])
     assert math.isclose(rep_model_params, 8520256)  # reparameterized model params
@@ -242,7 +241,7 @@ def test_dilate():
 def test_models_onnx_export(arch_name, input_shape, output_size):
     # Model
     batch_size = 2
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     if arch_name == "fast_tiny_rep":
         model = reparameterize(detection.fast_tiny(pretrained=True, exportable=True, input_shape=input_shape))
     else: