Upload 9 files

uploaded the codes and sample models

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Purushothaman
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,266 @@
----
-license: mit
----
+# Interpretable-SONAR-Image-Classifier
+
+Explainable AI for Underwater SONAR Image Classifier
+
+## Prerequisites
+
+- Python 3.6 or higher
+
+## Running the Scripts
+
+This guide will help you run the `data_loader.py`, `train.py`, `test.py`, `predict.py`, and `classify_image_and_explain.py` scripts directly from the command line or within a Python script.
+
+### Prerequisites
+
+1. **Python Installation**: Ensure you have Python installed. You can download it from [python.org](https://www.python.org/).
+
+2. **Required Packages**: Install the required packages using `requirements.txt`.
+   ```sh
+   pip install -r requirements.txt
+   ```
+
+### Script Descriptions and Usage
+
+#### 1. `data_loader.py`
+
+This script is used to load, process, split datasets (train, val, test), and augment data.
+
+**Command Line Usage:**
+
+```sh
+python data_loader.py --path <path_to_data> --target_folder <path_to_target_folder> --dim <dimension> --batch_size <batch_size> --num_workers <num_workers> [--augment_data]
+```
+
+**Arguments:**
+
+- `--path`: Path to the data.
+- `--target_folder`: Path to the target folder where processed data will be saved.
+- `--dim`: Dimension for resizing the images.
+- `--batch_size`: Batch size for data loading.
+- `--num_workers`: Number of workers for data loading.
+- `--augment_data` (optional): Flag to enable data augmentation.
+
+**Example:**
+
+```sh
+python data_loader.py --path "./dataset" --target_folder "./processed_data" --dim 224 --batch_size 32 --num_workers 4 --augment_data
+```
+
+**Dataset Structure:**
+
+```sh
+├── Dataset (Raw)
+   ├── class_name_1
+   │   └── *.jpg
+   ├── class_name_2
+   │   └── *.jpg
+   ├── class_name_3
+   │   └── *.jpg
+   └── class_name_4
+       └── *.jpg
+```
+
+#### 2. `train.py`
+
+This script is used for training and storing the models, leveraging transfer learning.
+
+**Command Line Usage:**
+
+```sh
+python train.py --base_model_names <model_names> --shape <shape> --data_path <data_path> --log_dir <log_dir> --model_dir <model_dir> --epochs <epochs> --optimizer <optimizer> --learning_rate <learning_rate> --batch_size <batch_size>
+```
+
+**Arguments:**
+
+- `--base_models`: Comma-separated list of base model names (e.g., 'VGG16, ResNet50').
+- `--shape`: Image shape (size).
+- `--data_path`: Path to the data.
+- `--log_dir`: Path to the log directory.
+- `--model_dir`: Path to the model directory.
+- `--epochs`: Number of training epochs.
+- `--optimizer`: Optimizer type ('adam' or 'sgd').
+- `--learning_rate`: Learning rate for the optimizer.
+- `--batch_size`: Batch size for training.
+- `--patience`: Patience for early stopping (to prevent overfitting).
+
+**Example:**
+
+```sh
+python train.py --base_models "VGG16" "DenseNet121" --shape 224 224 3 --data_path "./processed_data" --log_dir "./logs" --model_dir "./models" --epochs 100 --optimizer "adam" --learning_rate 0.0001 --batch_size 32
+```
+
+#### 3. `test.py`
+
+This script is used for testing and storing the test logs of the above-trained models.
+
+**Command Line Usage:**
+
+```sh
+python test.py --data_path <data_path> --base_model_name <base_model_name> --model_path <model_path> --models_folder_path <models_folder_path> --log_dir <log_dir>
+```
+
+**Arguments:**
+
+- `--models_dir` (optional): Path to the models (directory).
+- `--model_path`: Path to the model (.h5/Keras Model).
+- `--img_path`: Path to the image file.
+- `--test_dir`: Path to the test dataset (directory).
+- `--train_dir`: Path to the training data.
+- `--log_dir`: Path to the log directory.
+
+**Example:**
+
+```sh
+python test.py --model_path "./models/vgg16_model.keras" --test_dir "./test_data" --train_dir "./data/train" --log_dir "./logs"
+```
+
+#### 4. `predict.py`
+
+This script is used for making predictions on new images.
+
+**Command Line Usage:**
+
+```sh
+python predict.py --model_path <model_path> --img_path <img_path> --train_dir <train_dir>
+```
+
+**Arguments:**
+
+- `--model_path`: Path to the model file.
+- `--img_path`: Path to the image file.
+- `--train_dir`: Path to the training dataset (for the label decoder, can be replaced with a CSV file with slight code modifications).
+
+**Example:**
+
+```sh
+python predict.py --model_path "./models/vgg16_model.keras" --img_path "./images/test_image.jpg" --train_dir "./data/train"
+```
+
+#### 5. `classify_image_and_explain.py`
+
+This script is used for making predictions on new images and generating explanations using one or more explainers (LIME, SHAP, Grad-CAM). The explanations are saved in the specified output folder, with filenames indicating the method used (e.g., `lime_explanation_1.jpg`, `shap_explanation_1.jpg`, `gradcam_explanation_1.jpg`).
+
+**Command Line Usage:**
+
+```sh
+python classify_image_and_explain.py --image_path <image_path> --model_path <model_path> --train_directory <train_directory> --num_samples <num_samples> --num_features <num_features> --segmentation_alg <segmentation_alg> --kernel_size <kernel_size> --max_dist <max_dist> --ratio <ratio> --max_evals <max_evals> --batch_size <batch_size> --explainer_types <explainer_types> --output_folder <output_folder>
+```
+
+**Arguments:**
+
+- `--image_path` (required): Path to the input image.
+- `--model_path` (required): Path to the trained model.
+- `--train_directory` (required): Directory containing training images.
+- `--num_samples` (default: 300): Number of samples for LIME.
+- `--num_features` (default: 100): Number of features for LIME.
+- `--segmentation_alg` (default: `quickshift`): Segmentation algorithm for LIME (`quickshift`, `slic`).
+- `--kernel_size` (default: 4): Kernel size for the segmentation algorithm.
+- `--max_dist` (default: 200): Maximum distance for the segmentation algorithm.
+- `--ratio` (default: 0.2): Ratio for the segmentation algorithm.
+- `--max_evals` (default: 400): Maximum evaluations for SHAP.
+- `--batch_size` (default: 50): Batch size for SHAP.
+- `--explainer_types` (default: 'all'): Comma-separated list of explainers to use (`lime`, `shap`, `gradcam`). Use 'all' to include all three explainers.
+- `--output_folder` (optional): Folder to save explanation images.
+
+**Example:**
+
+```sh
+python classify_image_and_explain.py --image_path "./images/test_image.jpg" --model_path "./models/model.keras" --train_directory "./data/train" --num_samples 300 --num_features 100 --segmentation_alg "quickshift" --kernel_size 4 --max_dist 200 --ratio 0.2 --max_evals 400 --batch_size 50 --explainer_types "lime, gradcam" --output_folder "./explanations"
+```
+
+### Supported Base Models
+
+The following base models are supported for training:
+- VGG16
+- VGG19
+- ResNet50
+- ResNet101
+- InceptionV3
+- DenseNet121
+- DenseNet201
+- MobileNetV2
+- Xception
+- InceptionResNetV2
+- NASNetLarge
+- NASNetMobile
+- EfficientNetB0
+- EfficientNetB7
+
+### Running Scripts in a Python Script
+
+You can also run these scripts programmatically using Python's `subprocess` module. Here is an example of how to do this for each script:
+
+```python
+import subprocess
+
+# Run data_loader.py
+subprocess.run([
+    "python", "data_loader.py",
+    "--path", "./data",
+    "--target_folder", "./processed_data",
+    "--dim", "224",
+    "--batch_size", "32",
+    "--num_workers", "4",
+    "--augment_data"
+])
+
+# Run train.py
+subprocess.run([
+    "python", "train.py",
+    "--base_models", "VGG16,ResNet50",
+    "--shape", "224, 224, 3",
+    "--data_path", "./data",
+    "--log_dir", "./logs",
+    "--model_dir", "./models",
+    "--epochs", "100",
+    "--optimizer", "adam",
+    "--learning_rate", "0.001",
+    "--batch_size", "32",
+    "--patience", "10"
+])
+
+# Run test.py
+subprocess.run([
+    "python", "test.py",
+    "--models_dir", "./models",
+    "--img
+
+_path", "./images/test_image.jpg",
+    "--train_dir", "./data/train",
+    "--log_dir", "./logs"
+])
+
+# Run classify_image_and_explain.py
+subprocess.run([
+    "python", "classify_image_and_explain.py",
+    "--image_path", "./images/test_image.jpg",
+    "--model_path", "./models/model.h5",
+    "--train_directory", "./data/train",
+    "--num_samples", "300",
+    "--num_features", "100",
+    "--segmentation_alg", "quickshift",
+    "--kernel_size", "4",
+    "--max_dist", "200",
+    "--ratio", "0.2",
+    "--max_evals", "400",
+    "--batch_size", "50",
+    "--explainer_types", "lime,gradcam",
+    "--output_folder", "./explanations"
+])
+```
+
+## License
+This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
+
+## Citing the part of the project: Under water sonar image classifier with XAI LIME
+
+If you use our SONAR classifier or the explainer in your research, please use the following BibTeX entry.
+
+```
+@article{natarajan2024underwater,
+  title={Underwater SONAR Image Classification and Analysis using LIME-based Explainable Artificial Intelligence},
+  author={Natarajan, Purushothaman and Nambiar, Athira},
+  journal={arXiv preprint arXiv:2408.12837},
+  year={2024}
+}
+```

balanced_data_loader-1.py

@@ -0,0 +1,216 @@
+import tensorflow as tf
+import os
+import argparse
+from sklearn.model_selection import StratifiedShuffleSplit
+from tqdm import tqdm
+import uuid
+import random
+
+# Parses command line arguments
+def parse_arguments():
+    parser = argparse.ArgumentParser(description='Image Data Loader with Augmentation and Splits')
+    parser.add_argument('--path', type=str, required=True, help='Path to the folder containing images')
+    parser.add_argument('--dim', type=int, default=224, help='Required image dimension')
+    parser.add_argument('--batch_size', type=int, default=32, help='Batch size')
+    parser.add_argument('--target_folder', type=str, required=True, help='Folder to store the train, test, and val splits')
+    parser.add_argument('--augment_data', action='store_true', help='Apply data augmentation')
+    parser.add_argument('--balance', action='store_true', help='Balance the dataset')
+    parser.add_argument('--split_type', type=str, choices=['random', 'stratified'], default='random',
+                        help='Type of data split (random or stratified)')
+    return parser.parse_args()
+
+# Process the input images
+def process_image(file_path, image_size):
+    image = tf.io.read_file(file_path)
+    image = tf.image.decode_image(image, channels=3, dtype=tf.float32)
+    image = tf.image.resize(image, [image_size, image_size])
+    image = tf.clip_by_value(image, 0.0, 1.0)
+    return image
+
+# Balances the images of a specific class
+def balance_class_images(image_paths, labels, target_count, image_size, label, label_to_index, output_folder):
+    print(f"Balancing class '{label}'...")
+    label_idx = label_to_index.get(label, None)
+    if label_idx is None:
+        print(f"Label '{label}' not found in label_to_index.")
+        return [], []
+
+    image_paths = [img for img, lbl in zip(image_paths, labels) if lbl == label_idx]
+    num_images = len(image_paths)
+
+    print(f"Class '{label}' has {num_images} images before balancing.")
+
+    balanced_images = []
+    balanced_labels = []
+
+    original_count = num_images
+    synthetic_count = 0
+
+    if num_images > target_count:
+        balanced_images.extend(random.sample(image_paths, target_count))
+        balanced_labels.extend([label_idx] * target_count)
+        print(f"Removed {num_images - target_count} images from class '{label}'.")
+    elif num_images < target_count:
+        balanced_images.extend(image_paths)
+        balanced_labels.extend([label_idx] * num_images)
+
+        num_to_add = target_count - num_images
+        print(f"Class '{label}' needs {num_to_add} additional images for balancing.")
+
+        while num_to_add > 0:
+            img_path = random.choice(image_paths)
+            image = process_image(img_path, image_size)
+
+            for _ in range(min(num_to_add, 5)):  # Use up to 5 augmentations per image
+                augmented_image = augment_image(image)
+                balanced_images.append(augmented_image)
+                balanced_labels.append(label_idx)
+                num_to_add -= 1
+                synthetic_count += 1
+
+        print(f"Added {synthetic_count} augmented images to class '{label}'.")
+        print(f"Class '{label}' has {len(balanced_images)} images after balancing.")
+
+    class_folder = os.path.join(output_folder, str(label_idx))
+    if not os.path.exists(class_folder):
+        os.makedirs(class_folder)
+
+    for i, img in enumerate(balanced_images):
+        file_name = f"{uuid.uuid4()}.png"
+        file_path = os.path.join(class_folder, file_name)
+        save_image(img, file_path)
+
+    print(f"Saved {len(balanced_images)} images for class '{label}' (Original: {original_count}, Synthetic: {synthetic_count}).")
+
+    return balanced_images, balanced_labels
+
+# Saves an image to a file
+def save_image(image, file_path):
+    if isinstance(image, str):
+        image = process_image(image, image_size)
+    if isinstance(image, tf.Tensor):
+        image = tf.image.convert_image_dtype(image, dtype=tf.uint8)
+        image = tf.image.encode_png(image)
+    else:
+        raise ValueError("Expected image to be a TensorFlow tensor, but got a different type.")
+
+    tf.io.write_file(file_path, image)
+
+# Augments an image with random transformations
+def augment_image(image):
+    # Apply random augmentations using TensorFlow functions
+    image = tf.image.random_flip_left_right(image)
+    image = tf.image.random_flip_up_down(image)
+    image = tf.image.random_brightness(image, max_delta=0.1)
+    image = tf.image.random_contrast(image, lower=0.9, upper=1.1)
+    image = tf.image.random_saturation(image, lower=0.9, upper=1.1)
+    image = tf.image.random_hue(image, max_delta=0.1)
+    return image
+
+# Creates a list of data augmentation functions
+def create_datagens():
+    return [augment_image]
+
+# Balances the entire dataset by balancing each class
+def balance_data(images, labels, target_count, image_size, unique_labels, label_to_index, output_folder):
+    print(f"Balancing data: Target count per class = {target_count}")
+
+    all_balanced_images = []
+    all_balanced_labels = []
+
+    for label in tqdm(unique_labels, desc="Balancing classes"):
+        num_images = len([img for img, lbl in zip(images, labels) if lbl == label_to_index.get(label, -1)])
+        balanced_images, balanced_labels = balance_class_images(
+            images, labels, target_count, image_size, label, label_to_index, output_folder
+        )
+        all_balanced_images.extend(balanced_images)
+        all_balanced_labels.extend(balanced_labels)
+
+    total_original_images = sum(1 for img in all_balanced_images if isinstance(img, str))
+    total_synthetic_images = len(all_balanced_images) - total_original_images
+
+    print(f"\nTotal saved images: {len(all_balanced_images)} (Original: {total_original_images}, Synthetic: {total_synthetic_images})")
+
+    return all_balanced_images, all_balanced_labels
+
+# Augments an image using TensorFlow functions
+def tf_augment_image(file_path, label):
+    image = tf.image.resize(tf.image.decode_jpeg(tf.io.read_file(file_path)), [image_size, image_size])
+    image = tf.cast(image, tf.float32) / 255.0
+    augmented_image = augment_image(image)
+    return augmented_image, label
+
+
+def map_fn(file_path, label):
+    image, label = tf.py_function(tf_augment_image, [file_path, label], [tf.float32, tf.int32])
+    image.set_shape([image_size, image_size, 3])
+    label.set_shape([])
+    return image, label
+
+# Loads images, splits them into train, validation, and test sets, and saves the splits
+def load_and_save_splits(path, image_size, batch_size, balance, datagens, target_folder, split_type):
+    all_images = []
+    labels = []
+
+    for class_folder in os.listdir(path):
+        class_path = os.path.join(path, class_folder)
+        if os.path.isdir(class_path):
+            for img_file in os.listdir(class_path):
+                img_path = os.path.join(class_path, img_file)
+                all_images.append(img_path)
+                labels.append(class_folder)  # Use the folder name as the label
+
+    print(f"Loaded {len(all_images)} images across {len(set(labels))} classes.")
+    print(f"Labels found: {set(labels)}")  # Print unique labels
+
+    unique_labels = list(set(labels))
+    label_to_index = {label: idx for idx, label in enumerate(unique_labels)}
+    encoded_labels = [label_to_index[label] for label in labels]
+
+    print(f"Label to index mapping: {label_to_index}")
+
+    if split_type == 'stratified':
+        sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
+        train_indices, test_indices = next(sss.split(all_images, encoded_labels))
+    else:  # random split
+        total_images = len(all_images)
+        indices = list(range(total_images))
+        random.shuffle(indices)
+        train_indices = indices[:int(0.8 * total_images)]
+        test_indices = indices[int(0.8 * total_images):]
+
+    train_files = [all_images[i] for i in train_indices]
+    train_labels = [encoded_labels[i] for i in train_indices]
+    test_files = [all_images[i] for i in test_indices]
+    test_labels = [encoded_labels[i] for i in test_indices]
+
+    # Create validation and test sets
+    sss_val = StratifiedShuffleSplit(n_splits=1, test_size=0.5, random_state=42)
+    val_indices, test_indices = next(sss_val.split(test_files, test_labels))
+
+    val_files = [test_files[i] for i in val_indices]
+    val_labels = [test_labels[i] for i in val_indices]
+    test_files = [test_files[i] for i in test_indices]
+    test_labels = [test_labels[i] for i in test_indices]
+
+    # Save splits
+    for split_name, file_list, labels_list in [("train", train_files, train_labels), ("val", val_files, val_labels), ("test", test_files, test_labels)]:
+        split_folder = os.path.join(target_folder, split_name)
+        os.makedirs(split_folder, exist_ok=True)
+        with open(os.path.join(split_folder, f"{split_name}_list.txt"), 'w') as file_list_file:
+            for img_path, label in zip(file_list, labels_list):
+                label_folder = os.path.join(split_folder, str(label))
+                if not os.path.exists(label_folder):
+                    os.makedirs(label_folder)
+                file_list_file.write(f"{img_path}\n")
+                save_image(img_path, os.path.join(label_folder, f"{uuid.uuid4()}.png"))
+
+    print(f"Saved splits: train: {len(train_files)}, val: {len(val_files)}, test: {len(test_files)}.")
+
+# Main function to run the data loader
+def main():
+    args = parse_arguments()
+    load_and_save_splits(args.path, args.dim, args.batch_size, args.balance, create_datagens(), args.target_folder, args.split_type)
+
+if __name__ == "__main__":
+    main()

classify_image_and_explain.py

@@ -0,0 +1,256 @@
+import os
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.preprocessing.image import img_to_array, array_to_img, load_img
+from lime.lime_image import LimeImageExplainer, SegmentationAlgorithm
+import matplotlib.pyplot as plt
+from PIL import Image
+import argparse
+import shap
+import cv2
+import pickle
+
+image_counter = 0
+temp_folder = "temp_data"
+output_folder = "explanations"
+
+# Load the model and extract relevant details
+def load_model_details(model_path):
+    if model_path.endswith('.keras'):
+        print("Loading .keras format model...")
+        model = tf.keras.models.load_model(model_path, compile=False)
+    elif model_path.endswith('.h5'):
+        print("Loading .h5 format model...")
+        model = tf.keras.models.load_model(model_path, compile=False)
+    else:
+        print("Loading SavedModel using TFSMLayer...")
+        model = tf.keras.Sequential([
+            tf.keras.layers.TFSMLayer(model_path, call_endpoint='serving_default')
+        ])
+
+    input_shape = model.input_shape[1:3]
+    last_conv_layer_name = None
+    for layer in reversed(model.layers):
+        if isinstance(layer, tf.keras.layers.Conv2D):
+            last_conv_layer_name = layer.name
+            break
+
+    print(f"Model loaded with input shape: {input_shape} and last conv layer: {last_conv_layer_name}")
+    return model, last_conv_layer_name, input_shape
+
+# Load the label encoder based on the training directory
+def load_label_encoder(train_directory):
+    labels = sorted(os.listdir(train_directory))
+    label_encoder = {i: label for i, label in enumerate(labels)}
+    print(f"Label encoder created: {label_encoder}")
+    return label_encoder
+
+def load_and_preprocess_image(filename, image_size):
+    # Load and preprocess the image for model input
+    print(f"Loading and preprocessing image from: {filename}")
+    image = tf.io.read_file(filename)
+    image = tf.image.decode_image(image, channels=3)
+    
+    if not tf.executing_eagerly():
+        image.set_shape([None, None, 3])
+
+    image = tf.image.resize(image, [image_size[0], image_size[1]])
+    image = image / 255.0
+    image.set_shape([image_size[0], image_size[1], 3])
+
+    return image
+
+# Create a dataset from the training directory
+def create_dataset(data_dir, labels, image_size, batch_size):
+    print(f"Creating dataset from directory: {data_dir}")
+    image_files = []
+    image_labels = []
+
+    for label in labels:
+        label_dir = os.path.join(data_dir, label)
+        for image_file in os.listdir(label_dir):
+            image_files.append(os.path.join(label_dir, image_file))
+            image_labels.append(label)
+
+    label_map = {label: idx for idx, label in enumerate(labels)}
+    image_labels = [label_map[label] for label in image_labels]
+
+    dataset = tf.data.Dataset.from_tensor_slices((image_files, image_labels))
+    dataset = dataset.map(lambda x, y: (load_and_preprocess_image(x, image_size), y))
+    dataset = dataset.shuffle(buffer_size=len(image_files))
+    dataset = dataset.batch(batch_size).prefetch(buffer_size=tf.data.AUTOTUNE)
+
+    print("Dataset created and batched")
+    return dataset
+
+# Save preprocessed data (images and labels) to a file
+def save_preprocessed_data(X_train, y_train, file_path):
+    print(f"Saving preprocessed data to: {file_path}")
+    with open(file_path, 'wb') as file:
+        pickle.dump((X_train, y_train), file)
+
+
+def load_preprocessed_data(file_path):
+    print(f"Loading preprocessed data from: {file_path}")
+    with open(file_path, 'rb') as file:
+        return pickle.load(file)
+
+def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
+    # Generate a Grad-CAM heatmap for the given image and model
+    
+    grad_model = tf.keras.models.Model(
+        inputs=model.inputs, outputs=[model.get_layer(last_conv_layer_name).output, model.output]
+    )
+    with tf.GradientTape() as tape:
+        last_conv_layer_output, preds = grad_model(img_array)
+        preds = tf.convert_to_tensor(preds)
+        class_channel = preds[:, pred_index]
+        # if pred_index is None:
+        #     pred_index = tf.argmax(preds[0])  # Default to the class with the highest probability
+        # pred_index = tf.squeeze(pred_index)  # Ensure pred_index is a scalar tensor
+        # if tf.executing_eagerly():
+        #     pred_index = pred_index.numpy()  # Convert to a NumPy array
+        # pred_index = int(pred_index)  # Convert to a Python integer
+        # class_channel = preds[0][pred_index]
+    
+    grads = tape.gradient(class_channel, last_conv_layer_output)
+    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
+    last_conv_layer_output = last_conv_layer_output[0]
+    heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
+    heatmap = tf.squeeze(heatmap)
+    heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
+    return heatmap.numpy()
+
+def save_and_display_gradcam(array, heatmap, alpha=0.8):
+    # Save and display the Grad-CAM heatmap overlaid on the original image
+    print("Saving and displaying Grad-CAM result...")
+    heatmap = np.uint8(255 * heatmap)
+    jet = plt.cm.jet
+    jet_colors = jet(np.arange(256))[:, :3]
+    jet_heatmap = jet_colors[heatmap]
+    jet_heatmap = array_to_img(jet_heatmap)
+    jet_heatmap = jet_heatmap.resize((array.shape[1], array.shape[0]))
+    jet_heatmap = img_to_array(jet_heatmap)
+    superimposed_img = jet_heatmap * alpha + array
+    superimposed_img = array_to_img(superimposed_img)
+    return superimposed_img
+
+def generate_splime_mask_top_n(img_array, model, explainer, top_n=1, num_features=100, num_samples=300):
+    # Generate a SP-LIME mask for the given image and model
+    # Use superpixel segmentation for SP-LIME
+    segmentation_fn = SegmentationAlgorithm('quickshift', kernel_size=4, max_dist=200, ratio=0.2)
+
+    explanation_instance = explainer.explain_instance(
+        img_array, model.predict, top_labels=top_n, hide_color=0,
+        num_samples=num_samples, num_features=num_features, segmentation_fn=segmentation_fn
+    )
+    explanation_mask = explanation_instance.get_image_and_mask(
+        explanation_instance.top_labels[0], positive_only=False,
+        num_features=num_features, hide_rest=True
+    )[1]
+
+    # Ensure mask is in the same shape as the input image
+    mask = np.zeros_like(img_array)  # Create a mask of the same shape as img_array
+    mask[explanation_mask == 1] = img_array[explanation_mask == 1]  # Overlay highlighted regions
+    
+    # Set non-highlighted areas to white
+    mask = np.where(explanation_mask[:, :, np.newaxis] == 1, mask, 1.0)
+    
+    return mask, explanation_instance
+
+
+def explain_image_shap(img, model, class_names, top_prediction, max_evals=1000, batch_size=50):
+    # Generate SHAP explanations for the given image and model
+    masker = shap.maskers.Image("inpaint_telea", img[0].shape)  # Update if necessary
+
+    # Define a function to predict probabilities from the model
+    def f(X):
+        return model.predict(X)
+
+    # Create the SHAP explainer
+    explainer = shap.Explainer(f, masker, output_names=class_names)
+
+    # Get SHAP values
+    shap_values = explainer(img, max_evals=max_evals, batch_size=batch_size, outputs=shap.Explanation.argsort.flip[:1])
+
+    return shap_values
+
+def classify_image_and_explain(image_path, model_path, train_directory, num_samples, num_features, segmentation_alg, kernel_size, max_dist, ratio, max_evals, batch_size, explainer_types, output_folder):
+    # Main function to classify the image and generate explanations
+    global image_counter
+
+    if output_folder is None:
+        output_folder = "explanations"
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+    
+    model, last_conv_layer_name, input_shape = load_model_details(model_path)
+    label_encoder = load_label_encoder(train_directory)
+    labels = list(label_encoder.values())
+    
+    # Load the image
+    image = load_img(image_path, target_size=input_shape)
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    array = img_to_array(image)
+    img_array = array / 255.0
+    img_array = np.expand_dims(img_array, axis=0)
+    
+    # Predict the class of the image
+    predictions = model.predict(img_array)
+    top_prediction = np.argmax(predictions[0])
+    top_label = label_encoder[top_prediction]
+    
+    print(f"Prediction: {top_label} with probability {predictions[0][top_prediction]:.4f}")
+
+    # Generate explanations based on user-specified types
+    if 'gradcam' in explainer_types:
+        model.layers[-1].activation = None
+        heatmap = make_gradcam_heatmap(img_array, model, last_conv_layer_name)
+        gradcam_image = save_and_display_gradcam(img_to_array(image), heatmap)
+        gradcam_image.save(os.path.join(output_folder, f"gradcam_{image_counter}.png"))
+        
+    if 'lime' in explainer_types:
+        # SPLIME Explanation
+        explainer = LimeImageExplainer()
+        splime_mask, explanation_instance = generate_splime_mask_top_n(img_array[0], model, explainer, top_n=1, num_features=num_features, num_samples=num_samples)
+        # Ensure splime_mask is in [0, 1] range before saving
+        splime_mask = np.clip(splime_mask, 0, 1)
+        plt.imsave(os.path.join(output_folder, f"splime_{image_counter}.png"), splime_mask)
+
+    if 'shap' in explainer_types:
+        custom_image = img_to_array(image) / 255.0  # Preprocess image for SHAP
+        shap_values = explain_image_shap(custom_image.reshape(1, *custom_image.shape), model, labels, top_prediction, max_evals=max_evals, batch_size=batch_size)
+        shap.image_plot(shap_values[0], custom_image, labels=[top_label], show=False)
+        plt.savefig(os.path.join(output_folder, f"shap_{image_counter}.png"))
+        #plt.show()
+        plt.close()
+    
+    print("Image classification and explanation process completed.")
+    image_counter += 1
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Image classification and explanation script")
+    parser.add_argument("--image_path", type=str, required=True, help="Path to the input image")
+    parser.add_argument("--model_path", type=str, required=True, help="Path to the trained model")
+    parser.add_argument("--train_directory", type=str, required=True, help="Directory containing training images")
+    parser.add_argument("--num_samples", type=int, default=300, help="Number of samples for LIME")
+    parser.add_argument("--num_features", type=int, default=100, help="Number of features for LIME")
+    parser.add_argument("--segmentation_alg", type=str, default='quickshift', help="Segmentation algorithm for LIME (options: quickshift, slic)")
+    parser.add_argument("--kernel_size", type=int, default=4, help="Kernel size for segmentation algorithm")
+    parser.add_argument("--max_dist", type=int, default=200, help="Max distance for segmentation algorithm")
+    parser.add_argument("--ratio", type=float, default=0.2, help="Ratio for segmentation algorithm")
+    parser.add_argument("--max_evals", type=int, default=400, help="Maximum evaluations for SHAP")
+    parser.add_argument("--batch_size", type=int, default=50, help="Batch size for SHAP")
+    parser.add_argument("--explainer_types", type=str, default='all', help="Comma-separated list of explainers to use (options: lime, shap, gradcam). Use 'all' to include all three.")
+    parser.add_argument("--output_folder", type=str, default=None, help="Output folder for explanations")
+
+    args = parser.parse_args()
+    
+    explainer_types = args.explainer_types.split(',') if args.explainer_types != 'all' else ['lime', 'shap', 'gradcam']
+
+    classify_image_and_explain(
+        args.image_path, args.model_path, args.train_directory, args.num_samples,
+        args.num_features, args.segmentation_alg, args.kernel_size, args.max_dist,
+        args.ratio, args.max_evals, args.batch_size, explainer_types, args.output_folder
+    )

data_loader.py

@@ -0,0 +1,173 @@
+import tensorflow as tf
+import os
+import argparse
+from sklearn.model_selection import StratifiedShuffleSplit
+from tqdm import tqdm  # For progress display
+import sys
+import uuid  # Import uuid for unique filename generation
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description='Image Data Loader with Augmentation and Splits')
+    parser.add_argument('--path', type=str, required=True, help='Path to the folder containing images')
+    parser.add_argument('--dim', type=int, default=224, help='Required image dimension')
+    parser.add_argument('--batch_size', type=int, default=32, help='Batch size')
+    parser.add_argument('--num_workers', type=int, default=4, help='Number of workers for data loading')
+    parser.add_argument('--target_folder', type=str, required=True, help='Folder to store the train, test, and val splits')
+    parser.add_argument('--augment_data', action='store_true', help='Apply data augmentation')
+    return parser.parse_args()
+
+def create_datagens():
+    # Create a list of ImageDataGenerator objects for different augmentations
+    return [
+        ImageDataGenerator(rescale=1./255),
+        ImageDataGenerator(rotation_range=20),
+        ImageDataGenerator(width_shift_range=0.2),
+        ImageDataGenerator(height_shift_range=0.2),
+        ImageDataGenerator(horizontal_flip=True)
+    ]
+
+def process_image(file_path, image_size):
+    # Read, decode, resize, and normalize an image
+    file_path = file_path.numpy().decode('utf-8')
+    image = tf.io.read_file(file_path)
+    image = tf.image.decode_image(image, channels=3, dtype=tf.float32)
+    image = tf.image.resize(image, [image_size, image_size])
+    image = tf.clip_by_value(image, 0.0, 1.0)
+    return image
+
+def save_image(image, file_path):
+    # Convert image to uint8, encode as JPEG, and save to file
+    image = tf.image.convert_image_dtype(image, dtype=tf.uint8)
+    image = tf.image.encode_jpeg(image)
+    tf.io.write_file(file_path, image)
+
+def load_data(path, image_size, batch_size):
+    all_images = []
+    labels = []
+    # Load images and labels from the specified path
+
+    for subdir, _, files in os.walk(path):
+        label = os.path.basename(subdir)
+        for fname in files:
+            if fname.endswith(('.jpg', '.jpeg', '.png')):
+                all_images.append(os.path.join(subdir, fname))
+                labels.append(label)
+    
+    unique_labels = set(labels)
+    print(f"Found {len(all_images)} images in {path}\n")
+    print(f"Labels found ({len(unique_labels)}): {unique_labels}\n")
+    
+    # Raise an error if no images are found
+    if len(all_images) == 0:
+        raise ValueError(f"No images found in the specified path: {path}")
+
+    # Stratified splitting the dataset
+    sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
+    train_indices, test_indices = next(sss.split(all_images, labels))
+    
+    train_files = [all_images[i] for i in train_indices]
+    train_labels = [labels[i] for i in train_indices]
+    test_files = [all_images[i] for i in test_indices]
+    test_labels = [labels[i] for i in test_indices]
+    
+    sss_val = StratifiedShuffleSplit(n_splits=1, test_size=0.5, random_state=42)
+    val_indices, test_indices = next(sss_val.split(test_files, test_labels))
+    
+    val_files = [test_files[i] for i in val_indices]
+    val_labels = [test_labels[i] for i in val_indices]
+    test_files = [test_files[i] for i in test_indices]
+    test_labels = [test_labels[i] for i in test_indices]
+
+    print(f"Data split into {len(train_files)} train, {len(val_files)} validation, and {len(test_files)} test images.\n")
+
+    # Define a function to load and augment images
+    def tf_load_and_augment_image(file_path, label):
+        image = tf.py_function(func=lambda x: process_image(x, image_size), inp=[file_path], Tout=tf.float32)
+        image.set_shape([image_size, image_size, 3])
+        return image, label
+
+    train_dataset = tf.data.Dataset.from_tensor_slices((train_files, train_labels))
+    val_dataset = tf.data.Dataset.from_tensor_slices((val_files, val_labels))
+    # Create datasets from the loaded files and labels
+    test_dataset = tf.data.Dataset.from_tensor_slices((test_files, test_labels))
+
+    train_dataset = train_dataset.map(lambda x, y: tf_load_and_augment_image(x, y))
+    val_dataset = val_dataset.map(lambda x, y: tf_load_and_augment_image(x, y))
+    test_dataset = test_dataset.map(lambda x, y: tf_load_and_augment_image(x, y))
+
+    train_dataset = train_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
+    val_dataset = val_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
+    test_dataset = test_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
+
+    return train_dataset, val_dataset, test_dataset
+
+def save_datasets_to_folders(dataset, folder_path, datagens=None):
+    # Save the dataset to specified folders with optional augmentations
+    if not os.path.exists(folder_path):
+        os.makedirs(folder_path)
+    
+    count = 0
+    for batch_images, batch_labels in tqdm(dataset, desc=f"Saving to {folder_path}"):
+        for i in range(batch_images.shape[0]):
+            image = batch_images[i].numpy()
+            label = batch_labels[i].numpy().decode('utf-8')
+            label_folder = os.path.join(folder_path, label)
+            if not os.path.exists(label_folder):
+                os.makedirs(label_folder)
+
+            # Save the original image
+            file_path = os.path.join(label_folder, f"{uuid.uuid4().hex}.jpg")
+            save_image(image, file_path)
+            count += 1
+
+            # Apply augmentations if datagens are provided
+            if datagens:
+                for datagen in datagens:
+                    aug_image = datagen.random_transform(image)
+                    file_path = os.path.join(label_folder, f"{uuid.uuid4().hex}.jpg")
+                    save_image(aug_image, file_path)
+                    count += 1
+    
+    print(f"Saved {count} images to {folder_path}\n")
+    return count
+
+def main():
+    # Main function to parse arguments, load data, and save datasets
+    args = parse_arguments()
+    
+    if not os.path.exists(args.target_folder):
+        os.makedirs(args.target_folder)
+        
+    train_folder = os.path.join(args.target_folder, 'train')
+    val_folder = os.path.join(args.target_folder, 'val')
+    test_folder = os.path.join(args.target_folder, 'test')
+    
+    datagens = create_datagens() if args.augment_data else None
+
+    train_dataset, val_dataset, test_dataset = load_data(
+        args.path,
+        args.dim,
+        args.batch_size
+    )
+    
+    # Save datasets to respective folders and count images
+    train_count = save_datasets_to_folders(train_dataset, train_folder, datagens)
+    val_count = save_datasets_to_folders(val_dataset, val_folder)
+    test_count = save_datasets_to_folders(test_dataset, test_folder)
+    
+    print(f"Train dataset saved to: {train_folder}\n")
+    print(f"Validation dataset saved to: {val_folder}\n")
+    print(f"Test dataset saved to: {test_folder}\n")
+    
+    print('-'*20)
+
+    print(f"Number of images in training set: {train_count}\n")
+    print(f"Number of images in validation set: {val_count}\n")
+    print(f"Number of images in test set: {test_count}\n")
+
+if __name__ == "__main__":
+    # Redirect stdout and stderr to avoid encoding issues
+    sys.stdout = open(sys.stdout.fileno(), mode='w', encoding='utf-8', buffering=1)
+    sys.stderr = open(sys.stderr.fileno(), mode='w', encoding='utf-8', buffering=1)
+    main()

predict.py

@@ -0,0 +1,65 @@
+import os
+import argparse
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.preprocessing import image
+from tensorflow.keras.models import load_model
+
+def load_and_preprocess_image(img_path, target_size):
+    # Load and preprocess the image for prediction.
+    """Load and preprocess the image for prediction."""
+    img = image.load_img(img_path, target_size=target_size)
+    img_array = image.img_to_array(img)
+    img_array = np.expand_dims(img_array, axis=0)  # Create batch axis
+    img_array = img_array / 255.0  # Normalize the image
+    return img_array
+
+def load_model_from_file(model_path):
+    # Load the pre-trained model from the specified path.
+    """Load the pre-trained model from the specified path."""
+    model = load_model(model_path)
+    print(f"Model loaded from {model_path}")
+    return model
+
+def make_predictions(model, img_array):
+    # Make predictions using the loaded model.
+    """Make predictions using the loaded model."""
+    predictions = model.predict(img_array)
+    return predictions
+
+def get_class_names(train_dir):
+    """Get class names from training directory."""
+    class_names = os.listdir(train_dir)  # Assuming subfolder names are the class labels
+    class_names.sort()  # Ensure consistent ordering
+    return class_names
+
+def main(model_path, img_path, train_dir):
+    # Main function to load model, preprocess image, make predictions, and display results.
+    # Define target image size based on model requirements
+    target_size = (224, 224)  # Adjust if needed
+
+    # Load the model
+    model = load_model_from_file(model_path)
+
+    # Get class names from train directory
+    class_names = get_class_names(train_dir)
+
+    # Load and preprocess the image
+    img_array = load_and_preprocess_image(img_path, target_size)
+
+    # Make predictions
+    predictions = make_predictions(model, img_array)
+    predicted_label_index = np.argmax(predictions, axis=1)[0]
+    predicted_label = class_names[predicted_label_index]
+    probability_score = predictions[0][predicted_label_index]
+
+    print(f"Predicted label: {predicted_label}, Probability: {probability_score:.4f}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Load a pre-trained model and make a prediction on a new image")
+    parser.add_argument('--model_path', type=str, required=True, help='Path to the saved model')
+    parser.add_argument('--img_path', type=str, required=True, help='Path to the image to be predicted')
+    parser.add_argument('--train_dir', type=str, required=True, help='Directory containing training dataset for inferring class names')
+
+    args = parser.parse_args()
+    main(args.model_path, args.img_path, args.train_dir)

test.py

@@ -0,0 +1,161 @@
+import os
+import argparse
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.preprocessing import image
+from tensorflow.keras.models import load_model
+from sklearn.metrics import classification_report, confusion_matrix
+import matplotlib.pyplot as plt
+import seaborn as sns
+from tqdm import tqdm
+
+# Load and preprocess an image for prediction
+def load_and_preprocess_image(img_path, target_size):
+    """Load and preprocess the image for prediction."""
+    img = image.load_img(img_path, target_size=target_size)
+    img_array = image.img_to_array(img)
+    img_array = np.expand_dims(img_array, axis=0)  # Create batch axis
+    img_array = img_array / 255.0  # Normalize the image
+    return img_array
+
+# Load all models from a specified directory
+def load_all_models(model_dir):
+    """Load all models from the specified directory."""
+    models = {}
+    for file_name in os.listdir(model_dir):
+        if file_name.endswith('_model.keras'):
+            model_path = os.path.join(model_dir, file_name)
+            model_name = file_name.split('_model.keras')[0]  # Extract model name
+            model = load_model(model_path)
+            models[model_name] = model
+            print(f"Model loaded from {model_path}")
+    if not models:
+        raise FileNotFoundError(f"No model files found in {model_dir}.")
+    return models
+
+# Load a single model from a specified path
+def load_model_from_file(model_path):
+    """Load a single model from the specified path."""
+    model = load_model(model_path)
+    print(f"Model loaded from {model_path}")
+    return model
+
+def make_predictions(model, img_array):
+    # Make predictions using the loaded model
+    """Make predictions using the loaded model."""
+    predictions = model.predict(img_array)
+    return predictions
+
+def get_class_names(train_dir):
+    """Get class names from training directory."""
+    class_names = os.listdir(train_dir)  # Assuming subfolder names are the class labels
+    class_names.sort()  # Ensure consistent ordering
+    return class_names
+
+# Compute confusion matrix and classification report, and save to log directory
+def compute_confusion_matrix_and_report(true_labels, predicted_labels, class_names, log_dir, model_name):
+    """Compute confusion matrix and classification report, and save to log directory."""
+    # Compute confusion matrix
+    conf_matrix = confusion_matrix(true_labels, predicted_labels, labels=class_names)
+    report = classification_report(true_labels, predicted_labels, target_names=class_names)
+
+    # Print the classification report
+    print(f"Model: {model_name}")
+    print(report)
+
+    # Plot the confusion matrix
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues', xticklabels=class_names, yticklabels=class_names)
+    plt.xlabel('Predicted Label')
+    plt.ylabel('True Label')
+    plt.title(f'Confusion Matrix - {model_name}')
+    
+    # Save plot
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
+    
+    conf_matrix_plot_file = os.path.join(log_dir, f'confusion_matrix_{model_name}.png')
+    plt.savefig(conf_matrix_plot_file)
+    plt.close()
+
+    # Save results to log directory
+    conf_matrix_file = os.path.join(log_dir, f'confusion_matrix_{model_name}.txt')
+    report_file = os.path.join(log_dir, f'classification_report_{model_name}.txt')
+
+    np.savetxt(conf_matrix_file, conf_matrix, fmt='%d', delimiter=',', header=','.join(class_names))
+    with open(report_file, 'w') as f:
+        f.write(report)
+
+    print(f"Confusion matrix and classification report saved to {log_dir} with model name: {model_name}")
+
+# Main function to load models, make predictions, and evaluate performance
+def main(model_path, model_dir, img_path, test_dir, train_dir, log_dir):
+    # Define target image size based on model requirements
+    target_size = (224, 224)  # Adjust if needed
+
+    if model_path:
+        # Load a single model
+        model = load_model_from_file(model_path)
+        models = {os.path.basename(model_path): model}
+    elif model_dir:
+        # Load all models from a directory
+        models = load_all_models(model_dir)
+    else:
+        raise ValueError("Either --model_path or --model_dir must be provided.")
+    
+    # Get class names from train directory
+    class_names = get_class_names(train_dir)
+    num_classes = len(class_names)
+
+    # If an image path is provided, perform prediction on that image
+    if img_path:
+        img_array = load_and_preprocess_image(img_path, target_size)
+        for model_name, model in models.items():
+            print(f"Model: {model_name}")
+            predictions = make_predictions(model, img_array)
+            predicted_label_index = np.argmax(predictions, axis=1)[0]
+            if predicted_label_index >= num_classes:
+                raise ValueError(f"Predicted label index {predicted_label_index} is out of range for class names list.")
+            predicted_label = class_names[predicted_label_index]
+            probability_score = predictions[0][predicted_label_index]
+            print('-'*20)
+            print(f"Predicted label: {predicted_label}, Probability: {probability_score:.4f}")
+            print('-'*20)
+
+    # If a test directory is provided, perform batch predictions and evaluation
+    if test_dir:
+        files = [os.path.join(root, file) for root, _, files in os.walk(test_dir) for file in files if file.endswith(('png', 'jpg', 'jpeg'))]
+        
+        for model_name, model in models.items():
+            true_labels = []
+            predicted_labels = []
+
+            for img_path in tqdm(files, desc=f"Processing images with {model_name}"):
+                img_array = load_and_preprocess_image(img_path, target_size)
+                predictions = make_predictions(model, img_array)
+                predicted_label_index = np.argmax(predictions, axis=1)[0]
+                if predicted_label_index >= num_classes:
+                    raise ValueError(f"Predicted label index {predicted_label_index} is out of range for class names list.")
+                predicted_label = class_names[predicted_label_index]
+                
+                true_label = os.path.basename(os.path.dirname(img_path))  # Assuming folder name is the label
+                if true_label not in class_names:
+                    raise ValueError(f"True label {true_label} is not in class names list.")
+                
+                true_labels.append(true_label)
+                predicted_labels.append(predicted_label)
+            
+            # Compute and save confusion matrix and classification report
+            compute_confusion_matrix_and_report(true_labels, predicted_labels, class_names, log_dir, model_name)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Load models and make predictions on new images or a test dataset")
+    parser.add_argument('--model_path', type=str, help='Path to a single saved model')
+    parser.add_argument('--model_dir', type=str, help='Directory containing saved models (loads all models in the folder)')
+    parser.add_argument('--img_path', type=str, help='Path to the image to be predicted')
+    parser.add_argument('--test_dir', type=str, help='Directory containing test dataset for batch predictions')
+    parser.add_argument('--train_dir', type=str, required=True, help='Directory containing training dataset for inferring class names')
+    parser.add_argument('--log_dir', type=str, required=True, help='Directory to save prediction results')
+
+    args = parser.parse_args()
+    main(args.model_path, args.model_dir, args.img_path, args.test_dir, args.train_dir, args.log_dir)

train.py

@@ -0,0 +1,176 @@
+import os
+import argparse
+import tensorflow as tf
+from tensorflow.keras.models import Model
+from tensorflow.keras.applications import (VGG16, VGG19, ResNet50, ResNet101, InceptionV3,
+                                           DenseNet121, DenseNet201, MobileNetV2, Xception, InceptionResNetV2,
+                                           NASNetLarge, NASNetMobile, EfficientNetB0, EfficientNetB7)
+from tensorflow.keras.layers import Dense, Flatten, Dropout, BatchNormalization
+from tensorflow.keras.optimizers import Adam, SGD
+from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping
+import numpy as np
+
+def load_and_preprocess_image(filename, label, image_size):
+    # Load image
+    image = tf.io.read_file(filename)
+    image = tf.image.decode_image(image, channels=3)
+    
+    # Ensure the image tensor has shape
+    if not tf.executing_eagerly():
+        image.set_shape([None, None, 3])
+
+    # Resize image to the specified size
+    image = tf.image.resize(image, [image_size[0], image_size[1]])  # Use height and width from the tuple
+    
+    # Normalize image to [0, 1]
+    image = image / 255.0
+    image.set_shape([image_size[0], image_size[1], 3])
+
+    return image, label
+
+def create_dataset(data_dir, labels, image_size, batch_size):
+    image_files = []
+    image_labels = []
+
+    for label in labels:
+        label_dir = os.path.join(data_dir, label)
+        for image_file in os.listdir(label_dir):
+            image_files.append(os.path.join(label_dir, image_file))
+            image_labels.append(label)
+
+    # Create a mapping from labels to indices
+    label_map = {label: idx for idx, label in enumerate(labels)}
+    image_labels = [label_map[label] for label in image_labels]
+
+    # Convert to TensorFlow datasets
+    dataset = tf.data.Dataset.from_tensor_slices((image_files, image_labels))
+    dataset = dataset.map(lambda x, y: load_and_preprocess_image(x, y, image_size))
+    dataset = dataset.shuffle(buffer_size=len(image_files))
+    dataset = dataset.batch(batch_size).prefetch(buffer_size=tf.data.AUTOTUNE)
+
+    return dataset
+
+def create_and_train_model(base_model, model_name, shape, X_train, X_val, num_classes, labels, log_dir, model_dir,
+                           epochs, optimizer_name, learning_rate, step_gamma, alpha, batch_size, patience):
+    # Freeze the base model layers
+    for layer in base_model.layers:
+        layer.trainable = False
+
+    # Add custom layers on top
+    x = base_model.output
+    x = Flatten()(x)
+    x = Dense(1024, activation='relu')(x)
+    x = Dropout(0.25)(x)
+
+    x = Dense(512, activation='relu')(x)
+    x = Dropout(0.25)(x)
+
+    x = Dense(256, activation='relu')(x)
+    x = BatchNormalization()(x)
+    x = Dropout(0.25)(x)
+
+    predictions = Dense(num_classes, activation='softmax')(x)  # Use the number of classes
+    model = Model(inputs=base_model.input, outputs=predictions)
+
+    # Learning rate schedule
+    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
+        initial_learning_rate=learning_rate,
+        decay_steps=1000,  # Adjust this according to your needs
+        decay_rate=step_gamma
+    )
+
+    # Select the optimizer
+    if optimizer_name.lower() == 'adam':
+        optimizer = Adam(learning_rate=lr_schedule)
+    elif optimizer_name.lower() == 'sgd':
+        optimizer = SGD(learning_rate=lr_schedule, momentum=alpha)  # Example settings for SGD
+    else:
+        raise ValueError(f"Unsupported optimizer: {optimizer_name}")
+
+    # Compile the model
+    model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])
+
+    # Set up callbacks
+    checkpoint = ModelCheckpoint(os.path.join(model_dir, f'{model_name}_best_model.keras'), 
+                                 monitor='val_accuracy', save_best_only=True, save_weights_only=False, 
+                                 mode='max', verbose=1)
+    early_stopping = EarlyStopping(monitor='val_accuracy', patience=patience, verbose=1)
+
+    # Train the model
+    history = model.fit(X_train, epochs=epochs, validation_data=X_val, batch_size=batch_size,
+                        callbacks=[checkpoint, early_stopping])
+
+    # Save training logs
+    with open(os.path.join(log_dir, f'{model_name}_training.log'), 'w') as f:
+        num_epochs = len(history.history['loss'])  # Get the actual number of epochs completed
+        for epoch in range(num_epochs):
+            f.write(f"Epoch {epoch + 1}, "
+                    f"Train Loss: {history.history['loss'][epoch]:.4f}, "
+                    f"Train Accuracy: {history.history['accuracy'][epoch]:.4f}, "
+                    f"Val Loss: {history.history['val_loss'][epoch]:.4f}, "
+                    f"Val Accuracy: {history.history['val_accuracy'][epoch]:.4f}\n")
+
+    # Save labels in the model directory
+    with open(os.path.join(model_dir, 'labels.txt'), 'w') as f:
+        f.write('\n'.join(labels))
+
+    # Evaluate the model
+    test_loss, test_accuracy = model.evaluate(X_val)
+    print(f'Test Accuracy for {model_name}: {test_accuracy:.4f}')
+    print(f'Test Loss for {model_name}: {test_loss:.4f}')
+
+    # Optionally, save the trained model
+    model.save(os.path.join(model_dir, f'{model_name}_final_model.keras'))
+
+def main(base_model_names, shape, data_path, log_dir, model_dir, epochs, optimizer, learning_rate, step_gamma, alpha, batch_size, patience):
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir)
+
+    # Extract labels from folder names
+    labels = sorted([d for d in os.listdir(os.path.join(data_path, 'train')) if os.path.isdir(os.path.join(data_path, 'train', d))])
+    num_classes = len(labels)
+
+    # Load data
+    X_train = create_dataset(os.path.join(data_path, 'train'), labels, shape, batch_size)
+    X_val = create_dataset(os.path.join(data_path, 'val'), labels, shape, batch_size)
+
+    if not base_model_names:
+        print("No base models specified. Exiting.")
+        return
+
+    # Define base models
+    base_models_dict = {
+        model_name: globals()[model_name](weights='imagenet', include_top=False, input_shape=shape)
+        for model_name in base_model_names
+    }
+
+    for model_name in base_model_names:
+        print(f'Training {model_name}...')
+        base_model = base_models_dict.get(model_name)
+        if base_model is None:
+            print(f"Model {model_name} not supported.")
+            continue
+        create_and_train_model(base_model, model_name, shape, X_train, X_val, num_classes, labels, log_dir, model_dir,
+                               epochs, optimizer, learning_rate, step_gamma, alpha, batch_size, patience)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Train models using transfer learning")
+    parser.add_argument('--base_models', type=str, nargs='+', default=[],
+                        help='List of base models to use for training. Leave empty to skip model training.')
+    parser.add_argument('--shape', type=int, nargs=3, default=(224, 224, 3), help='Input shape of the images')
+    parser.add_argument('--data_path', type=str, required=True, help='Path to the image data')
+    parser.add_argument('--log_dir', type=str, required=True, help='Directory to save logs')
+    parser.add_argument('--model_dir', type=str, required=True, help='Directory to save models')
+    parser.add_argument('--epochs', type=int, default=100, help='Number of epochs to train')
+    parser.add_argument('--optimizer', type=str, default='adam', help='Optimizer to use (adam or sgd)')
+    parser.add_argument('--learning_rate', type=float, default=0.001, help='Initial learning rate')
+    parser.add_argument('--step_gamma', type=float, default=0.96, help='Gamma value for step learning rate schedule')
+    parser.add_argument('--alpha', type=float, default=0.9, help='Alpha for the optimizer (used for SGD)')
+    parser.add_argument('--batch_size', type=int, default=32, help='Batch size for training')
+    parser.add_argument('--patience', type=int, default=10, help='Patience for early stopping')
+
+    args = parser.parse_args()
+    main(args.base_models, tuple(args.shape), args.data_path, args.log_dir, args.model_dir,
+         args.epochs, args.optimizer, args.learning_rate, args.step_gamma, args.alpha, args.batch_size, args.patience)