focal.model_pipeline module

Model pipeline module for the FOCAL application.

This module provides classes for building, training, and managing CNN model for fiber cleave quality classification. It uses binary classification to place the image as either a good or bad cleave.

class focal.model_pipeline.CustomModel(train_ds: DatasetV2, test_ds: DatasetV2, num_classes: int, classification_type: str | None = 'binary')[source]

Bases: object

Class for defining custom models using pre-trained MobileNetV2.

This class provides functionality for building, compiling, and training CNN models for fiber cleave classification.

compile_custom_model(image_shape: Tuple[int, int, int], learning_rate: float = 0.001, metrics: List[str] | None = None, num_classes: int | None = 5) → Model[source]

Compile custom model after calling build_custom_model function.

Parameters:

  • image_shape – Dimensions of input images

  • param_shape – Dimensions of numerical parameters

  • learning_rate – Learning rate for optimization

  • metrics – List of metrics to monitor

  • num_class – number of output classes

Returns:

Compiled model ready for training

Return type:

tf.keras.Model

compile_image_only_model(image_shape: Tuple[int, int, int], learning_rate: float = 0.001, metrics: List[str] | None = None, backbone: str | None = 'mobilenet', num_classes: int = 5, dropout1: float | None = 0.1, dense1: int | None = 32, dropout2: float | None = 0.2, l2_factor: float | None = None, unfreeze_from: int | None = None) → Model[source]

Compile an image-only model.

Parameters:

  • image_shape – Dimensions of input images

  • learning_rate – Learning rate for optimization

  • metrics – List of metrics to monitor

Returns:

Compiled image-only model

Return type:

tf.keras.Model

compile_model(image_shape: Tuple[int, int, int], param_shape: Tuple[int, ...], dropout1: float, dense1: int, dropout2: float, dense2: int, dropout3: float, brightness: float, height: float, width: float, contrast: float, rotation: float, learning_rate: float = 0.001, metrics: List[str] | None = None, unfreeze_from: int | None = None, backbone: str | None = 'mobilenet') → Model[source]

Compile the custom model head on top of pre-trained backbone.

Parameters:

  • image_shape (Tuple[int, int, int]) – Shape of each image.

  • param_shape (Tuple[int, ...]) – Shape of feature vector.

  • dropout1 (float) – Amount of dropout for image layer.

  • dense1 (int) – Amount of neurons for first FC layer.

  • dropout2 (float) – Amount of dropout for feature input.

  • dense2 (int) – Amount of neruons for second FC layer.

  • dropout3 (float) – Amount of dropout for concatenated images+features.

  • brightness (float) – Brightness amount for image.

  • height (float) – Height of augmented image.

  • width (float) – Width of augmented image.

  • contrast (float) – Contrast amount of augmented image.

  • rotation (float) – Rotation amount of augmented image.

  • learning_rate (float, optional) – Size of steps to take during optimization. Defaults to 0.001.

  • metrics (Optional[List[str]], optional) – Saved metrics. Defaults to None.

  • unfreeze_from (Optional[int], optional) – Layer to start training pre-trained backbone. Defaults to None.

  • backbone (Optional[str], optional) – Name of pre-trained backbone. Defaults to “mobilenet”.

Returns:

Compiled custom head on pre-trained model backbone.

Return type:

Model

create_checkpoints(checkpoint_filepath: str = './checkpoints/model.keras', monitor: str = 'val_accuracy', mode: str = 'max', save_best_only: bool = True) → ModelCheckpoint[source]

Create model checkpoints to avoid losing data while training.

Parameters:

  • checkpoint_filepath – Path to save model checkpoints

  • monitor – Metric to monitor during training

  • mode – Method to determine stopping point of metric (max, min, auto)

  • save_best_only – Whether to save only the best model

Returns:

Checkpoint callback

Return type:

tf.keras.callbacks.ModelCheckpoint

create_early_stopping(patience: int = 3, mode: str = 'max', monitor: str = 'val_accuracy') → EarlyStopping[source]

Create early stopping callback to monitor training success and prevent overfitting.

Parameters:

  • patience – Number of epochs to wait before stopping when monitor plateaus

  • mode – Method to track monitor (max, min, auto)

  • monitor – Metric to monitor during training

Returns:

Early stopping callback

Return type:

tf.keras.callbacks.EarlyStopping

create_tensorboard_callback(log_dir: str = './logs', histogram_freq: int = 1) → TensorBoard[source]

Create TensorBoard callback for monitoring training.

Parameters:

  • log_dir – Directory for TensorBoard logs

  • histogram_freq – Frequency for computing weight histograms

Returns:

TensorBoard callback

Return type:

tf.keras.callbacks.TensorBoard

static get_averages_from_kfold(kfold_histories: List[History]) → None[source]

Calculate and display average metrics from k-fold cross validation.

Parameters:

kfold_histories – List of training histories from k-fold training

get_data_augmentation(rotation: float, brightness: float, height: float, width: float, contrast: float)[source]

Get data augmentation parameters for training model.

Parameters:

  • rotation (float) – Rotation amount of image.

  • brightness (float) – Brightness amount of image.

  • height (float) – Height of image.

  • width (float) – Widht of image.

  • contrast (float) – Contrast amount of image.

Returns:

Data augmentation model.

Return type:

tf.keras.models.Model

plot_metric(title: str, metric_1: List[float], metric_2: List[float], metric_1_label: str, metric_2_label: str, x_label: str, y_label: str, model_path: str) → None[source]

Plot training metrics for visualization.

Parameters:

  • title – Title for the plot

  • metric_1 – First metric values to plot

  • metric_2 – Second metric values to plot

  • metric_1_label – Label for first metric

  • metric_2_label – Label for second metric

  • x_label – Label for x-axis

  • y_label – Label for y-axis

reduce_on_plateau(patience: int = 3, mode: str = 'auto', factor: float = 2.0, monitor: str = 'val_accuracy') → ReduceLROnPlateau[source]

Create reduced learning rate callback if monitored value stops improving.

Parameters:

  • patience – Number of epochs before reducing learning rate

  • mode – Method to monitor (min, max, auto)

  • factor – Factor by which to reduce learning rate

  • monitor – Value to monitor before reducing learning rate

Returns:

Reduce learning rate callback

Return type:

tf.keras.callbacks.ReduceLROnPlateau

static train_kfold(datasets: List[Tuple], image_shape: Tuple[int, int, int], param_shape: Tuple[int, ...], learning_rate: float, num_classes: int, dropout1: float, dense1: int, dropout2: float, dense2: int, dropout3: float, brightness: float, height: float, width: float, contrast: float, rotation: float, metrics: List[str] | None = None, epochs: int = 5, initial_epoch: int = 0, history_file: str | None = None, save_model_file: str | None = None, callbacks: List | None = None) → Tuple[List[Model], List[History]][source]

Train model using k-fold cross validation.

Parameters:

  • datasets – List of (train_ds, test_ds) tuples for each fold

  • image_shape – Dimensions of input images

  • param_shape – Dimensions of numerical parameters

  • dropout1 – Amount of dropout for image layer.

  • dense1 – Amount of neurons for first FC layer.

  • dropout2 – Amount of dropout for feature input.

  • dense2 – Amount of neruons for second FC layer.

  • dropout3 – Amount of dropout for concatenated images+features.

  • brightness – Brightness amount for image.

  • height – Height of augmented image.

  • width – Width of augmented image.

  • contrast – Contrast amount of augmented image.

  • rotation – Rotation amount of augmented image.

  • learning_rate – Learning rate for optimization.

  • metrics – List of metrics to monitor.

  • epochs – Number of training epochs.

  • initial_epoch – Starting epoch number.

  • history_file – Base filename for saving training history.

  • model_file – Base filename for saving models.

Returns:

Tuple of (list of trained models, list of training histories)

train_model(class_weights, model: Model, epochs: int = 5, initial_epoch: int = 0, callbacks: List | None = None, history_file: str | None = None, save_model_file: str | None = None) → History[source]

Train model with possible callbacks to prevent overfitting.

Parameters:

  • model – Model to be trained

  • checkpoints – Checkpoints to save model

  • epochs – Number of training epochs

  • initial_epoch – Starting epoch number

  • early_stopping – Early stopping callback

  • reduce_lr – Reduce learning rate callback

  • tensorboard – TensorBoard callback

  • history_file – File to save training history

  • model_file – File to save trained model

Returns:

Training history

Return type:

tf.keras.callbacks.History