Remove prediction_model

src/model/management.py

@@ -383,39 +383,3 @@ def verify_charlist_length(charlist: List[str],
             f"Charlist length ({len(charlist)}) does not match "
             f"model output length ({expected_length}). If the charlist "
             "is correct, try setting use_mask to True.")
-
-
-def get_prediction_model(model: tf.keras.Model) -> tf.keras.Model:
-    """
-    Extracts a prediction model from a given Keras model.
-
-    Parameters
-    ----------
-    model : tf.keras.Model
-        The complete Keras model from which the prediction model is to be
-        extracted.
-
-    Returns
-    -------
-    tf.keras.Model
-        The prediction model extracted from the given model, typically up to
-        the last dense layer.
-
-    Raises
-    ------
-    ValueError
-        If no dense layer is found in the given model.
-    """
-
-    last_dense_layer = None
-    for layer in reversed(model.layers):
-        if layer.name.startswith('dense'):
-            last_dense_layer = layer
-            break
-    if last_dense_layer is None:
-        raise ValueError("No dense layer found in the model")
-
-    prediction_model = tf.keras.models.Model(
-        model.get_layer(name="image").input, last_dense_layer.output
-    )
-    return prediction_model

src/modes/inference.py

@@ -9,7 +9,6 @@
 
 # > Local dependencies
 from data.manager import DataManager
-from model.management import get_prediction_model
 from setup.config import Config
 from utils.decoding import decode_batch_predictions
 from utils.text import Tokenizer
@@ -50,12 +49,11 @@ def perform_inference(config: Config,
     """
 
     tokenizer = Tokenizer(charlist, config["use_mask"])
-    prediction_model = get_prediction_model(model)
 
     with open(config["results_file"], "w", encoding="utf-8") as results_file:
         for batch_no, batch in enumerate(inference_dataset):
             # Get the predictions
-            predictions = prediction_model.predict_on_batch(batch[0])
+            predictions = model.predict_on_batch(batch[0])
             y_pred = decode_batch_predictions(predictions,
                                               tokenizer,
                                               config["greedy"],

src/modes/test.py

@@ -9,7 +9,6 @@
 
 # > Local dependencies
 from data.manager import DataManager
-from model.management import get_prediction_model
 from setup.config import Config
 from utils.calculate import calc_95_confidence_interval, calculate_cers, \
     increment_counters, calculate_edit_distances
@@ -146,7 +145,6 @@ def perform_test(config: Config,
     logging.info("Performing test...")
 
     tokenizer = Tokenizer(charlist, config["use_mask"])
-    prediction_model = get_prediction_model(model)
 
     # Setup WordBeamSearch if needed
     wbs = setup_word_beam_search(config, charlist) \
@@ -161,7 +159,7 @@ def perform_test(config: Config,
         X, y_true, _ = batch
 
         logging.info("Batch %s/%s", batch_no + 1, len(test_dataset))
-        batch_counter = process_batch((X, y_true), prediction_model, tokenizer,
+        batch_counter = process_batch((X, y_true), model, tokenizer,
                                       config, wbs, data_manager, charlist)
 
         # Update the total counter

src/modes/validation.py

@@ -11,7 +11,6 @@
 
 # > Local dependencies
 from data.manager import DataManager
-from model.management import get_prediction_model
 from setup.config import Config
 from utils.calculate import calc_95_confidence_interval, \
     calculate_edit_distances, update_statistics, increment_counters