iBrokeTheCode/Multimodal_Product_Classification
0
1import os2from itertools import cycle3 4import matplotlib5import tensorflow as tf6 7# 💬 NOTE: Handle plots issues when running tests or displaying in notebooks8try:9 get_ipython # Only exists in Jupyter10 matplotlib.use("module://matplotlib_inline.backend_inline")11except Exception:12 matplotlib.use("Agg") # Fix error with tests13 14import matplotlib.pyplot as plt15import numpy as np16import pandas as pd17import seaborn as sns18from sklearn.metrics import (19 accuracy_score,20 classification_report,21 confusion_matrix,22 f1_score,23 precision_score,24 recall_score,25 roc_auc_score,26 roc_curve,27)28from sklearn.preprocessing import LabelEncoder29from sklearn.utils.class_weight import compute_class_weight30from tensorflow.keras import Input, Model31from tensorflow.keras.callbacks import EarlyStopping32from tensorflow.keras.layers import BatchNormalization, Concatenate, Dense, Dropout33from tensorflow.keras.losses import CategoricalCrossentropy34from tensorflow.keras.optimizers import SGD, Adam35from tensorflow.keras.utils import Sequence36 37 38class MultimodalDataset(Sequence):39 """40 Custom Keras Dataset class for multimodal data handling, designed for models that41 take both text and image data as inputs. It facilitates batching and shuffling42 of data for efficient training in Keras models.43 44 This class supports loading and batching multimodal data (text and images), as well as handling45 label encoding. It is compatible with Keras and can be used to train models that require both46 text and image inputs. It also supports optional shuffling at the end of each epoch for better47 training performance.48 49 Args:50 df (pd.DataFrame): The DataFrame containing the dataset with text, image, and label columns.51 text_cols (list): List of column names corresponding to text data. Can be a single column or multiple columns.52 image_cols (list): List of column names corresponding to image data (usually file paths or image pixel data).53 label_col (str): Column name corresponding to the target labels.54 encoder (LabelEncoder, optional): A pre-fitted LabelEncoder instance for encoding the labels.55 If None, a new LabelEncoder is fitted based on the provided data.56 batch_size (int, optional): Number of samples per batch. Default is 32.57 shuffle (bool, optional): Whether to shuffle the dataset at the end of each epoch. Default is True.58 59 Attributes:60 text_data (np.ndarray): Array of text data from the DataFrame. None if `text_cols` is not provided.61 image_data (np.ndarray): Array of image data from the DataFrame. None if `image_cols` is not provided.62 labels (np.ndarray): One-hot encoded labels corresponding to the dataset's classes.63 encoder (LabelEncoder): Fitted LabelEncoder used to encode target labels.64 batch_size (int): Number of samples per batch.65 shuffle (bool): Flag indicating whether to shuffle the data after each epoch.66 indices (np.ndarray): Array of indices representing the dataset. Used for shuffling batches.67 68 Methods:69 -------70 __len__():71 Returns the number of batches per epoch based on the dataset size and batch size.72 73 __getitem__(idx):74 Retrieves a single batch of data, including both text and image inputs and the corresponding labels.75 The method returns a tuple in the format ({'text': text_batch, 'image': image_batch}, label_batch),76 where 'text' and 'image' are only included if their respective columns were provided.77 78 on_epoch_end():79 Updates the index order after each epoch, shuffling if needed.80 """81 82 def __init__(83 self,84 df,85 text_cols,86 image_cols,87 label_col,88 encoder=None,89 batch_size=32,90 shuffle=True,91 ):92 """93 Initializes the MultimodalDataset object.94 95 Args:96 df (pd.DataFrame): The dataset as a DataFrame, containing text, image, and label data.97 text_cols (list): List of column names representing text features.98 image_cols (list): List of column names representing image features (e.g., file paths or pixel data).99 label_col (str): Column name corresponding to the target labels.100 encoder (LabelEncoder, optional): LabelEncoder for encoding the target labels. If None, a new LabelEncoder will be created.101 batch_size (int, optional): Batch size for loading data. Default is 32.102 shuffle (bool, optional): Whether to shuffle the data at the end of each epoch. Default is True.103 104 Raises:105 ValueError: If both text_cols and image_cols are None or empty.106 """107 if text_cols:108 # Get the text data from the DataFrame as a NumPy array109 self.text_data = df[text_cols].astype(np.float32).values110 else:111 # Else, set text data to None112 self.text_data = None113 114 if image_cols:115 # Get the image data from the DataFrame as a NumPy array116 self.image_data = df[image_cols].astype(np.float32).values117 else:118 # Else, set image data to None119 self.image_data = None120 121 if not text_cols and not image_cols:122 raise ValueError(123 "At least one of text_cols or image_cols must be provided."124 )125 126 # Get the labels from the DataFrame and encode them127 self.labels = df[label_col].values128 129 # Use provided encoder or fit a new one130 if encoder is None:131 self.encoder = LabelEncoder()132 self.labels = self.encoder.fit_transform(self.labels)133 else:134 self.encoder = encoder135 self.labels = self.encoder.transform(self.labels)136 137 # One-hot encode labels for multi-class classification138 num_classes = len(self.encoder.classes_)139 self.labels = np.eye(num_classes)[self.labels]140 141 self.batch_size = batch_size142 self.shuffle = shuffle143 self.on_epoch_end()144 145 def __len__(self):146 """147 Returns the number of batches per epoch based on the dataset size and batch size.148 149 Returns:150 -------151 int:152 The number of batches per epoch.153 """154 return int(np.floor(len(self.labels) / self.batch_size))155 156 def __getitem__(self, idx):157 """158 Retrieves a single batch of data (text and/or image) and the corresponding labels.159 160 Args:161 idx (int): Index of the batch to retrieve.162 163 Returns:164 -------165 tuple:166 A tuple containing the batch of text and/or image inputs and the corresponding labels.167 The input data is returned as a dictionary with keys 'text' and 'image', depending on the provided columns.168 If no text or image columns were provided, only the other is returned.169 """170 indices = self.indices[idx * self.batch_size : (idx + 1) * self.batch_size]171 172 if self.text_data is not None:173 text_batch = self.text_data[indices]174 if self.image_data is not None:175 image_batch = self.image_data[indices]176 label_batch = self.labels[indices]177 178 if self.text_data is None:179 return {"image": image_batch}, label_batch180 if self.image_data is None:181 return {"text": text_batch}, label_batch182 else:183 return {"text": text_batch, "image": image_batch}, label_batch184 185 def on_epoch_end(self):186 """187 Updates the index order after each epoch, shuffling the data if needed.188 189 This method is called at the end of each epoch and will shuffle the data if the `shuffle` flag is set to True.190 """191 self.indices = np.arange(len(self.labels))192 if self.shuffle:193 np.random.shuffle(self.indices)194 195 196# Early Fusion Model197def create_early_fusion_model(198 text_input_size, image_input_size, output_size, hidden=[128], p=0.2199):200 """201 Creates a multimodal early fusion model combining text and image inputs. The model concatenates the text and202 image features, passes them through fully connected layers with optional dropout and batch normalization,203 and produces a multi-class classification output.204 205 Args:206 text_input_size (int): Size of the input vector for the text data.207 image_input_size (int): Size of the input vector for the image data.208 output_size (int): Number of classes for the output layer (i.e., size of the softmax output).209 hidden (int or list, optional): Specifies the number of hidden units in the dense layers.210 If an integer, a single dense layer with the specified units is created.211 If a list, multiple dense layers are created with the respective units. Default is [128].212 p (float, optional): Dropout rate to apply after each dense layer. Default is 0.2.213 214 Returns:215 Model (keras.Model): A compiled Keras model with text and image inputs and a softmax output for classification.216 217 Model Architecture:218 - The model accepts two inputs: one for text features and one for image features.219 - The features are concatenated into a single vector.220 - Dense layers with ReLU activation are applied, followed by dropout and batch normalization (if multiple hidden layers are specified).221 - The output layer uses a softmax activation for multi-class classification.222 223 Example:224 model = create_early_fusion_model(text_input_size=300, image_input_size=2048, output_size=10, hidden=[128, 64], p=0.3)225 model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])226 """227 228 if text_input_size is None and image_input_size is None:229 raise ValueError(230 "At least one of text_input_size and image_input_size must be provided."231 )232 233 # Define inputs234 if text_input_size is not None:235 # Define text input layer for only text data236 text_input = Input(shape=(text_input_size,), name="text")237 if image_input_size is not None:238 # Define image input layer for only image data239 image_input = Input(shape=(image_input_size,), name="image")240 241 # Merge or select inputs242 if text_input_size is not None and image_input_size is not None:243 # Concatenate text and image inputs if both are provided244 x = Concatenate(name="fusion_layer")([text_input, image_input])245 elif text_input_size is not None:246 x = text_input247 elif image_input_size is not None:248 x = image_input249 250 # Hidden layers251 if isinstance(hidden, int):252 # Add a single dense layer, activation, dropout and normalization253 x = Dense(hidden, activation="relu")(x)254 x = Dropout(p)(x)255 x = BatchNormalization()(x)256 elif isinstance(hidden, list):257 for h in hidden:258 # Add multiple dense layers based on the hidden list, activation, dropout and normalization259 x = Dense(h, activation="relu")(x)260 x = Dropout(p)(x)261 x = BatchNormalization()(x)262 263 # Output layer264 # Add the output layer with softmax activation265 output = Dense(output_size, activation="softmax", name="output")(x)266 267 # Create the model268 if text_input_size is not None and image_input_size is not None:269 # Define the model with both text and image inputs270 model = Model(inputs=[text_input, image_input], outputs=output)271 elif text_input_size is not None:272 # Define the model with only text input273 model = Model(inputs=text_input, outputs=output)274 elif image_input_size is not None:275 # Define the model with only image input276 model = Model(inputs=image_input, outputs=output)277 else:278 raise ValueError(279 "At least one of text_input_size and image_input_size must be provided."280 )281 282 return model283 284 285def test_model(y_test, y_pred, y_prob=None, encoder=None):286 """287 Evaluates a trained model's performance using various metrics such as accuracy, precision, recall, F1-score,288 and visualizations including a confusion matrix and ROC curves.289 290 Args:291 y_test (np.ndarray): Ground truth one-hot encoded labels for the test data.292 y_pred (np.ndarray): Predicted class labels by the model for the test data (after argmax transformation).293 y_prob (np.ndarray, optional): Predicted probabilities for each class from the model. Required for ROC curves. Default is None.294 encoder (LabelEncoder, optional): A fitted LabelEncoder instance used to inverse transform one-hot encoded and predicted labels to their original categorical form.295 296 Returns:297 accuracy (float): Accuracy score of the model on the test data.298 precision (float): Weighted precision score of the model on the test data.299 recall (float): Weighted recall score of the model on the test data.300 f1 (float): Weighted F1 score of the model on the test data.301 302 This function performs the following steps:303 - Inverse transforms the one-hot encoded `y_test` and predicted `y_pred` values to their original labels using the provided LabelEncoder.304 - Computes the confusion matrix and plots it as a heatmap using Seaborn.305 - If `y_prob` is provided, computes and plots the ROC curves for each class.306 - Prints the classification report, which includes precision, recall, F1-score, and support for each class.307 - Returns the overall accuracy, weighted precision, recall, and F1-score of the model.308 309 Visualizations:310 - Confusion Matrix: A heatmap of the confusion matrix comparing the true labels with the predicted labels.311 - ROC Curves: Plots ROC curves for each class if predicted probabilities are provided (`y_prob`).312 313 Example:314 accuracy, precision, recall, f1 = test_model(y_test, y_pred, y_prob, encoder)315 """316 # Handle label decoding317 y_test_binarized = y_test318 y_test = encoder.inverse_transform(np.argmax(y_test, axis=1))319 y_pred = encoder.inverse_transform(y_pred)320 321 cm = confusion_matrix(y_test, y_pred)322 fig, ax = plt.subplots(figsize=(15, 15))323 sns.heatmap(cm, annot=True, cmap="Blues", fmt="g", ax=ax)324 plt.xlabel("Predicted")325 plt.ylabel("True")326 plt.title("Confusion Matrix")327 plt.show()328 329 if y_prob is not None:330 fig, ax = plt.subplots(figsize=(15, 15))331 332 colors = cycle(["aqua", "darkorange", "cornflowerblue"])333 334 for i, color in zip(range(y_prob.shape[1]), colors):335 fpr, tpr, _ = roc_curve(y_test_binarized[:, i], y_prob[:, i])336 ax.plot(fpr, tpr, color=color, lw=2, label=f"Class {i}")337 338 ax.plot([0, 1], [0, 1], "k--")339 plt.title("ROC Curve")340 plt.ylabel("True Positive Rate")341 plt.xlabel("False Positive Rate")342 plt.legend()343 plt.show()344 345 cr = classification_report(y_test, y_pred)346 print(cr)347 348 accuracy = accuracy_score(y_test, y_pred)349 precision = precision_score(y_test, y_pred, average="weighted")350 recall = recall_score(y_test, y_pred, average="weighted")351 f1 = f1_score(y_test, y_pred, average="weighted")352 353 return accuracy, precision, recall, f1354 355 356def train_mlp(357 train_loader,358 test_loader,359 text_input_size,360 image_input_size,361 output_size,362 num_epochs=50,363 report=False,364 lr=0.001,365 set_weights=True,366 adam=False,367 p=0.0,368 seed=1,369 patience=40,370 save_results=True,371 train_model=True,372 test_mlp_model=True,373):374 """375 Trains a multimodal early fusion model using both text and image data.376 377 The function handles the training process of the model by combining text and image features,378 computes class weights if needed, applies an optimizer (SGD or Adam), and implements early stopping379 to prevent overfitting. The model is evaluated on the test set, and key performance metrics are computed.380 381 Args:382 train_loader (MultimodalDataset): Keras-compatible data loader for the training set with both text and image data.383 test_loader (MultimodalDataset): Keras-compatible data loader for the test set with both text and image data.384 text_input_size (int): The size of the input vector for the text data.385 image_input_size (int): The size of the input vector for the image data.386 output_size (int): Number of output classes for the softmax layer.387 num_epochs (int, optional): Number of training epochs. Default is 50.388 report (bool, optional): Whether to generate a detailed classification report and display metrics. Default is False.389 lr (float, optional): Learning rate for the optimizer. Default is 0.001.390 set_weights (bool, optional): Whether to compute and apply class weights to handle imbalanced datasets. Default is True.391 adam (bool, optional): Whether to use the Adam optimizer instead of SGD. Default is False.392 p (float, optional): Dropout rate for regularization in the model. Default is 0.0.393 seed (int, optional): Seed for random number generators to ensure reproducibility. Default is 1.394 patience (int, optional): Number of epochs with no improvement on validation loss before early stopping. Default is 40.395 396 Returns:397 None398 399 Side Effects:400 - Trains the early fusion model and saves the best weights based on validation loss.401 - Generates plots showing the training and validation accuracy over epochs.402 - If `report` is True, calls `test_model` to print detailed evaluation metrics and plots.403 404 Training Process:405 - The function creates a fusion model combining text and image inputs.406 - Class weights are computed to balance the dataset if `set_weights` is True.407 - The model is trained using categorical cross-entropy loss and the chosen optimizer (Adam or SGD).408 - Early stopping is applied based on validation loss to prevent overfitting.409 - After training, the model is evaluated on the test set, and accuracy, F1-score, and AUC are calculated.410 411 Example:412 train_mlp(train_loader, test_loader, text_input_size=300, image_input_size=2048, output_size=10, num_epochs=30, lr=0.001, adam=True, report=True)413 414 Notes:415 - `train_loader` and `test_loader` should be instances of `MultimodalDataset` or compatible Keras data loaders.416 - If the dataset is imbalanced, setting `set_weights=True` is recommended to ensure better model performance on minority classes.417 """418 419 if seed is not None:420 np.random.seed(seed)421 tf.random.set_seed(seed)422 423 # Create an early fusion model using the provided input sizes and output size424 model = create_early_fusion_model(text_input_size, image_input_size, output_size)425 426 # Compute class weights for imbalanced datasets427 class_weights = None428 if set_weights:429 class_indices = np.argmax(train_loader.labels, axis=1)430 # Compute class weights using the training labels431 weights = compute_class_weight(432 class_weight="balanced",433 classes=np.unique(class_indices),434 y=class_indices,435 )436 class_weights = {i: w for i, w in enumerate(weights)}437 438 # Choose the loss function for multi-class classification439 loss = CategoricalCrossentropy()440 441 # Choose the optimizer442 if adam:443 # Use the Adam optimizer with the specified learning rate444 optimizer = Adam(learning_rate=lr)445 else:446 # Use the SGD optimizer with the specified learning rate447 optimizer = SGD(learning_rate=lr)448 449 # Compile the model with the chosen optimizer and loss function450 model.compile(optimizer=optimizer, loss=loss, metrics=["accuracy"])451 452 # Define an early stopping callback with the specified patience453 early_stopping = EarlyStopping(454 monitor="val_loss",455 patience=patience,456 restore_best_weights=True,457 )458 459 # Train the model using the training data and validation data460 history = None461 if train_model:462 # 📌 Train the model463 history = model.fit(464 train_loader,465 validation_data=test_loader,466 epochs=num_epochs,467 class_weight=class_weights,468 callbacks=[early_stopping],469 verbose="1",470 )471 472 if test_mlp_model:473 # 📌 Test the model on the test set474 y_true, y_pred, y_prob = [], [], []475 for batch in test_loader:476 features, labels = batch477 if len(features) == 1:478 text = features["text"] if "text" in features else features["image"]479 preds = model.predict(text)480 else:481 text, image = features["text"], features["image"]482 preds = model.predict([text, image])483 y_true.extend(labels)484 y_pred.extend(np.argmax(preds, axis=1))485 y_prob.extend(preds)486 487 y_true, y_pred, y_prob = np.array(y_true), np.array(y_pred), np.array(y_prob)488 489 test_accuracy = accuracy_score(np.argmax(y_true, axis=1), y_pred)490 f1 = f1_score(np.argmax(y_true, axis=1), y_pred, average="macro")491 492 auc_scores = roc_auc_score(y_true, y_prob, average="macro", multi_class="ovr")493 macro_auc = auc_scores494 495 plt.plot(history.history["accuracy"], label="Train Accuracy")496 plt.plot(history.history["val_accuracy"], label="Validation Accuracy")497 plt.xlabel("Epoch")498 plt.ylabel("Accuracy")499 plt.legend()500 plt.show()501 502 if report:503 test_model(y_true, y_pred, y_prob, encoder=train_loader.encoder)504 505 # 📌 Store results in a dataframe and save in the results folder506 if text_input_size is not None and image_input_size is not None:507 model_type = "multimodal"508 elif text_input_size is not None:509 model_type = "text"510 elif image_input_size is not None:511 model_type = "image"512 513 if save_results:514 results = pd.DataFrame(515 {"Predictions": y_pred, "True Labels": np.argmax(y_true, axis=1)}516 )517 # create results folder if it does not exist518 os.makedirs("results", exist_ok=True)519 results.to_csv(f"results/{model_type}_results.csv", index=False)520 521 # 📌 Save the model522 models_dir = "trained_models"523 os.makedirs(models_dir, exist_ok=True)524 525 model_filename = os.path.join(models_dir, f"{model_type}_model")526 model.save(model_filename)527 print(f"✅ {model_type} model saved successfully")528 else:529 test_accuracy, f1, macro_auc = None, None, None530 531 return model, test_accuracy, f1, macro_auc532