add inference script

utils/data.py

@@ -0,0 +1,341 @@
+import os
+from torch.utils.data import DataLoader, Dataset, Subset
+from torchvision.datasets import ImageFolder
+from sklearn.model_selection import train_test_split
+import torch
+import matplotlib.pyplot as plt
+import pickle
+
+CLASS_NAMES = ['Abra',
+ 'Aerodactyl',
+ 'Alakazam',
+ 'Alolan Sandslash',
+ 'Arbok',
+ 'Arcanine',
+ 'Articuno',
+ 'Beedrill',
+ 'Bellsprout',
+ 'Blastoise',
+ 'Bulbasaur',
+ 'Butterfree',
+ 'Caterpie',
+ 'Chansey',
+ 'Charizard',
+ 'Charmander',
+ 'Charmeleon',
+ 'Clefable',
+ 'Clefairy',
+ 'Cloyster',
+ 'Cubone',
+ 'Dewgong',
+ 'Diglett',
+ 'Ditto',
+ 'Dodrio',
+ 'Doduo',
+ 'Dragonair',
+ 'Dragonite',
+ 'Dratini',
+ 'Drowzee',
+ 'Dugtrio',
+ 'Eevee',
+ 'Ekans',
+ 'Electabuzz',
+ 'Electrode',
+ 'Exeggcute',
+ 'Exeggutor',
+ 'Farfetchd',
+ 'Fearow',
+ 'Flareon',
+ 'Gastly',
+ 'Gengar',
+ 'Geodude',
+ 'Gloom',
+ 'Golbat',
+ 'Goldeen',
+ 'Golduck',
+ 'Golem',
+ 'Graveler',
+ 'Grimer',
+ 'Growlithe',
+ 'Gyarados',
+ 'Haunter',
+ 'Hitmonchan',
+ 'Hitmonlee',
+ 'Horsea',
+ 'Hypno',
+ 'Ivysaur',
+ 'Jigglypuff',
+ 'Jolteon',
+ 'Jynx',
+ 'Kabuto',
+ 'Kabutops',
+ 'Kadabra',
+ 'Kakuna',
+ 'Kangaskhan',
+ 'Kingler',
+ 'Koffing',
+ 'Krabby',
+ 'Lapras',
+ 'Lickitung',
+ 'Machamp',
+ 'Machoke',
+ 'Machop',
+ 'Magikarp',
+ 'Magmar',
+ 'Magnemite',
+ 'Magneton',
+ 'Mankey',
+ 'Marowak',
+ 'Meowth',
+ 'Metapod',
+ 'Mew',
+ 'Mewtwo',
+ 'Moltres',
+ 'MrMime',
+ 'Muk',
+ 'Nidoking',
+ 'Nidoqueen',
+ 'Nidorina',
+ 'Nidorino',
+ 'Ninetales',
+ 'Oddish',
+ 'Omanyte',
+ 'Omastar',
+ 'Onix',
+ 'Paras',
+ 'Parasect',
+ 'Persian',
+ 'Pidgeot',
+ 'Pidgeotto',
+ 'Pidgey',
+ 'Pikachu',
+ 'Pinsir',
+ 'Poliwag',
+ 'Poliwhirl',
+ 'Poliwrath',
+ 'Ponyta',
+ 'Porygon',
+ 'Primeape',
+ 'Psyduck',
+ 'Raichu',
+ 'Rapidash',
+ 'Raticate',
+ 'Rattata',
+ 'Rhydon',
+ 'Rhyhorn',
+ 'Sandshrew',
+ 'Sandslash',
+ 'Scyther',
+ 'Seadra',
+ 'Seaking',
+ 'Seel',
+ 'Shellder',
+ 'Slowbro',
+ 'Slowpoke',
+ 'Snorlax',
+ 'Spearow',
+ 'Squirtle',
+ 'Starmie',
+ 'Staryu',
+ 'Tangela',
+ 'Tauros',
+ 'Tentacool',
+ 'Tentacruel',
+ 'Vaporeon',
+ 'Venomoth',
+ 'Venonat',
+ 'Venusaur',
+ 'Victreebel',
+ 'Vileplume',
+ 'Voltorb',
+ 'Vulpix',
+ 'Wartortle',
+ 'Weedle',
+ 'Weepinbell',
+ 'Weezing',
+ 'Wigglytuff',
+ 'Zapdos',
+ 'Zubat']
+
+class TransformSubset(Dataset):
+    """
+    Wrapper for applying transformations to a Subset.
+    """
+
+    def __init__(self, subset, transform):
+        self.subset = subset
+        self.transform = transform
+
+    def __getitem__(self, idx):
+        img, label = self.subset[idx]
+        if self.transform:
+            img = self.transform(img)
+        return img, label
+
+    def __len__(self):
+        return len(self.subset)
+
+
+class PokemonDataModule(Dataset):
+    def __init__(self, data_dir):
+        self.dataset = ImageFolder(root=data_dir)
+        self.class_names = self.dataset.classes
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, index):
+        image, label = self.dataset[index]
+        return image, label
+
+    def plot_examples(self, dataloader, n_rows=1, n_cols=4, stats=None):
+        """
+        Plot examples from a DataLoader.
+
+        Args:
+            dataloader (DataLoader): DataLoader object to fetch images and labels from.
+            n_rows (int): Number of rows in the plot grid.
+            n_cols (int): Number of columns in the plot grid.
+            denormalize (callable, optional): Function to reverse normalization for visualization.
+                                              Should accept a tensor and return a denormalized tensor.
+        """
+        fig, axes = plt.subplots(n_rows, n_cols, figsize=(n_cols * 3, n_rows * 3))
+        axes = axes.flatten()  # Flatten to iterate easily
+
+        # Iterate over the dataloader to get a batch of data
+        for data, labels in dataloader:
+            # Take the first n_rows * n_cols samples from the batch
+            for i, ax in enumerate(axes[: n_rows * n_cols]):
+                if i >= len(data):  # If fewer samples than the grid size, stop
+                    break
+
+                img, label = data[i], labels[i]
+
+                # Apply denormalization if provided
+                if stats:
+                    img = self._denormalize(img, stats)
+
+                # Convert CHW to HWC for plotting
+                img = img.permute(1, 2, 0).cpu().numpy()
+
+                ax.imshow(img)
+                ax.set_title(self.class_names[label.item()])
+                ax.axis("off")
+            break  # Only process the first batch
+
+        plt.tight_layout()
+        plt.show()
+
+    def _denormalize(self, img, stats):
+        """
+        Denormalize an image tensor.
+
+        Args:
+            img (Tensor): Image tensor with shape (C, H, W).
+            stats (dict): Dictionary containing 'means' and 'stds' for each channel.
+                          Example: {'means': [0.485, 0.456, 0.406], 'stds': [0.229, 0.224, 0.225]}.
+
+        Returns:
+            Tensor: Denormalized image tensor.
+        """
+        return img * stats["std"].view(-1, 1, 1) + stats["mean"].view(-1, 1, 1)
+
+    def _get_stats(self, dataset):
+        """
+        Calculate the mean and standard deviation of the dataset for standardization.
+        """
+        dataloader = DataLoader(dataset, batch_size=2048, shuffle=False)
+        total_sum, total_squared_sum, total_count = 0, 0, 0
+        with torch.cuda.device(0):
+            for data, _ in dataloader:
+                data.cuda()
+                total_sum += data.sum(dim=(0, 2, 3))
+                total_squared_sum += (data**2).sum(dim=(0, 2, 3))
+                total_count += data.size(0) * data.size(2) * data.size(3)
+
+            means = total_sum / total_count
+            stds = torch.sqrt((total_squared_sum / total_count) - (means**2))
+        return {"mean": means, "std": stds}
+
+    def prepare_data(self, indices_file="indices.pkl", get_stats=False):
+        """
+        Prepare train and test dataloaders with optional transformations.
+
+        Args:
+            indices_file (str): Path to save or load train/test indices.
+            transform (callable): Primary transformation to apply to the data.
+            additional_transforms (callable): Additional transformations to compose.
+
+        Returns:
+            tuple: trainloader, testloader
+        """
+        try:
+            with open(indices_file, "rb") as f:
+                self.train_indices, self.test_indices = pickle.load(f)
+        except (EOFError, FileNotFoundError):
+            # Generate new indices if file is empty or doesn't exist
+            self.train_indices, self.test_indices = train_test_split(
+                range(len(self.dataset)),
+                test_size=0.2,
+                stratify=self.dataset.targets,
+                random_state=42,
+            )
+
+            # Ensure directory exists before saving
+            os.makedirs(os.path.dirname(indices_file) or ".", exist_ok=True)
+
+            with open(indices_file, "wb") as f:
+                pickle.dump([self.train_indices, self.test_indices], f)
+
+        # Prepare train and test subsets
+        self.train_dataset = Subset(self.dataset, self.train_indices)
+        self.test_dataset = Subset(self.dataset, self.test_indices)
+
+        return self._get_stats(self.train_dataset) if get_stats else None
+
+    def get_dataloaders(
+        self,
+        train_transform=None,
+        test_transform=None,
+        train_batch_size=None,
+        test_batch_size=None,
+    ):
+        """
+        Prepare train and test dataloaders with optional transformations.
+
+        Args:
+            train_transform (callable): Transformation to apply to training data.
+            train_batch_size (int): Batch size for the training dataloader.
+            validation_batch_size (int): Batch size for the validation dataloader.
+
+        Returns:
+            tuple: trainloader, testloader
+        """
+        assert (
+            self.train_dataset is not None
+        ), "You need to call `prepare_data` before using `get_dataloaders`."
+
+        # Default batch sizes if not provided
+        test_batch_size = (
+            train_batch_size if test_batch_size is None else test_batch_size
+        )
+
+        # Wrap subsets in a transformed dataset if transformations are provided
+        train_dataset = (
+            TransformSubset(self.train_dataset, train_transform)
+            if train_transform
+            else self.train_dataset
+        )
+
+        test_dataset = (
+            TransformSubset(self.test_dataset, test_transform)
+            if test_transform
+            else self.test_dataset
+        )
+
+        trainloader = DataLoader(
+            train_dataset, batch_size=train_batch_size, shuffle=True
+        )
+        testloader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False)
+
+        return trainloader, testloader

utils/inference_utils.py

@@ -0,0 +1,193 @@
+import torch
+import matplotlib.pyplot as plt
+from torchvision import transforms
+from PIL import Image
+import os
+import random
+from utils.data import CLASS_NAMES
+
+# Function to find correctly and incorrectly classified images
+def find_images(dataloader, model, device, num_correct, num_incorrect):
+    correct_images = []
+    incorrect_images = []
+    correct_labels = []
+    incorrect_labels = []
+    correct_preds = []
+    incorrect_preds = []
+
+    model.eval()
+    with torch.no_grad():
+        for images, labels in dataloader:
+            images, labels = images.to(device), labels.to(device)
+            outputs = model(images)
+            _, preds = torch.max(outputs, 1)
+
+            for i in range(images.size(0)):
+                if preds[i] == labels[i] and len(correct_images) < num_correct:
+                    correct_images.append(images[i].cpu())
+                    correct_labels.append(labels[i].cpu())
+                    correct_preds.append(preds[i].cpu())
+                elif preds[i] != labels[i] and len(incorrect_images) < num_incorrect:
+                    incorrect_images.append(images[i].cpu())
+                    incorrect_labels.append(labels[i].cpu())
+                    incorrect_preds.append(preds[i].cpu())
+
+                if (
+                    len(correct_images) >= num_correct
+                    and len(incorrect_images) >= num_incorrect
+                ):
+                    break
+            if (
+                len(correct_images) >= num_correct
+                and len(incorrect_images) >= num_incorrect
+            ):
+                break
+
+    return (
+        correct_images,
+        correct_labels,
+        correct_preds,
+        incorrect_images,
+        incorrect_labels,
+        incorrect_preds,
+    )
+
+def find_images_from_path(data_path, model, device, num_correct=2, num_incorrect=2, label=None):
+    correct_images_paths = []
+    incorrect_images_paths = []
+    correct_labels = []
+    incorrect_labels = []
+
+    label_to_idx = {label: idx for idx, label in enumerate(CLASS_NAMES)}
+
+    model.eval()
+    # First collect available images for the specified label or all labels
+    label_images = {}
+    if label:
+        if os.path.isdir(os.path.join(data_path, label)):
+            label_path = os.path.join(data_path, label)
+            label_images[label] = [os.path.join(label_path, img) for img in os.listdir(label_path)]
+    else:
+        for label in os.listdir(data_path):
+            label_path = os.path.join(data_path, label)
+            if not os.path.isdir(label_path):
+                continue
+            label_images[label] = [os.path.join(label_path, img) for img in os.listdir(label_path)]
+
+    # Randomly process images until we have enough samples
+    with torch.no_grad():
+        while len(correct_images_paths) < num_correct or len(incorrect_images_paths) < num_incorrect:
+            # Randomly select a label that still has unprocessed images
+            available_labels = [l for l in label_images if label_images[l]]
+            if not available_labels:
+                break
+                
+            selected_label = random.choice(available_labels)
+            image_path = random.choice(label_images[selected_label])
+            label_images[selected_label].remove(image_path)  # Remove the selected image
+            
+            image = preprocess_image(image_path, (224, 224)).to(device)
+            label_idx = label_to_idx[selected_label]
+            
+            outputs = model(image)
+            _, pred = torch.max(outputs, 1)
+
+            if pred == label_idx and len(correct_images_paths) < num_correct:
+                correct_images_paths.append(image_path)
+                correct_labels.append(label_idx)
+            elif pred != label_idx and len(incorrect_images_paths) < num_incorrect:
+                incorrect_images_paths.append(image_path)
+                incorrect_labels.append(label_idx)
+
+    save_images_by_class(correct_images_paths, correct_labels, incorrect_images_paths, incorrect_labels)
+
+def save_images_by_class(correct_images_paths, correct_labels, incorrect_images_paths, incorrect_labels):
+    # Create root directories for correct and incorrect classifications
+    for class_name in CLASS_NAMES:
+        os.makedirs(os.path.join('predictions', class_name, 'correct'), exist_ok=True)
+        os.makedirs(os.path.join('predictions', class_name, 'mistake'), exist_ok=True)
+
+    # Save correctly classified images
+    for img_path, label in zip(correct_images_paths, correct_labels):
+        class_name = CLASS_NAMES[label]
+        img_name = os.path.basename(img_path)
+        destination = os.path.join('predictions', class_name, 'correct', img_name)
+        os.makedirs(os.path.dirname(destination), exist_ok=True)
+        Image.open(img_path).save(destination)
+
+    # Save incorrectly classified images
+    for img_path, label in zip(incorrect_images_paths, incorrect_labels):
+        class_name = CLASS_NAMES[label]
+        img_name = os.path.basename(img_path)
+        destination = os.path.join('predictions', class_name, 'mistake', img_name)
+        os.makedirs(os.path.dirname(destination), exist_ok=True)
+        Image.open(img_path).save(destination)
+
+def show_samples(dataloader, model, device, num_correct=3, num_incorrect=3):
+    # Get some correctly and incorrectly classified images
+    (
+        correct_images,
+        correct_labels,
+        correct_preds,
+        incorrect_images,
+        incorrect_labels,
+        incorrect_preds,
+    ) = find_images(dataloader, model, device, num_correct, num_incorrect)
+    # Display the results in a grid
+    fig, axes = plt.subplots(
+        num_correct + num_incorrect, 1, figsize=(10, (num_correct + num_incorrect) * 5)
+    )
+
+    for i in range(num_correct):
+        axes[i].imshow(correct_images[i].permute(1, 2, 0))
+        axes[i].set_title(
+            f"Correctly Classified: True Label = {correct_labels[i]}, Predicted = {correct_preds[i]}"
+        )
+        axes[i].axis("off")
+
+    for i in range(num_incorrect):
+        axes[num_correct + i].imshow(incorrect_images[i].permute(1, 2, 0))
+        axes[num_correct + i].set_title(
+            f"Incorrectly Classified: True Label = {incorrect_labels[i]}, Predicted = {incorrect_preds[i]}"
+        )
+        axes[num_correct + i].axis("off")
+
+    plt.tight_layout()
+    plt.show()
+
+
+# Function to preprocess image
+def preprocess_image(image_path, img_shape):
+
+    # Load the image using PIL
+    image = Image.open(image_path)
+
+    # Apply preprocessing transformations
+    preprocess = transforms.Compose([
+        transforms.Resize(img_shape),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    image = preprocess(image).unsqueeze(0)
+
+    return image
+
+
+# Function to predict
+def predict(model, image):
+    model.eval()
+    with torch.no_grad():
+        outputs = model(image)
+    return outputs
+
+
+# Function to get model predictions for LIME
+def batch_predict(model, images, device):
+    model.eval()
+    batch = torch.stack(
+        tuple(preprocess_image(image, (224, 224)) for image in images), dim=0
+    )
+    batch = batch.to(device)
+    logits = model(batch)
+    probs = torch.nn.functional.softmax(logits, dim=1)
+    return probs.detach().cpu().numpy()

utils/interpretability.py

@@ -0,0 +1,110 @@
+from lime import lime_image
+from skimage.segmentation import mark_boundaries
+import matplotlib.pyplot as plt
+from utils.inference_utils import predict
+import os
+import torch
+from PIL import Image
+import numpy as np
+
+
+def unnormalize(image):
+    # Make sure the image is on the correct dtype and device
+    # Convert mean and std to torch tensors with the correct dtype
+    mean = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)  # Use torch.float32
+    std = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)    # Use torch.float32
+
+    # If the image is a PyTorch tensor, ensure it has the same dtype
+    if isinstance(image, torch.Tensor):
+        image = image * std + mean
+    else:
+        image = torch.tensor(image, dtype=torch.float32) * std + mean  # Convert to torch if necessary
+
+    return image
+
+
+
+def lime_interpret_image_inference(args, model, image, device):
+    # prepare the image
+    def prepare_for_plot(image): return unnormalize(image).cpu().numpy()
+    # Remove batch dimension and Rearrange dimensions to (H, W, C)
+    image = image.squeeze(0).permute(1, 2, 0)  # From From [1, 3, 224, 224] to [224, 224, 3]
+    
+    # Convert to NumPy array
+    image_np = image.cpu().numpy()  # Ensure the tensor is on the CPU
+
+    # Initialize LIME explainer
+    explainer = lime_image.LimeImageExplainer()
+
+    # Define the prediction function
+    def predict_fn(x):
+        # Convert (B, H, W, C) to PyTorch tensor (B, C, H, W)
+        x_tensor = torch.tensor(x).permute(0, 3, 1, 2).to(device)
+        preds = model(x_tensor)
+        return preds.detach().cpu().numpy()
+
+    # Run LIME explanation
+    explanation = explainer.explain_instance(
+        image_np,
+        predict_fn,
+        top_labels=5,
+        hide_color=0,
+        num_samples=5000
+    )
+
+    # Get the mask for the top predicted class
+    temp, mask = explanation.get_image_and_mask(
+        explanation.top_labels[0],
+        positive_only=True,
+        num_features=10,
+        hide_rest=False
+    )
+
+    # Create a 2x2 subplot
+    fig, axs = plt.subplots(2, 2, figsize=(15, 15))
+
+    # Plot the original image
+    axs[0, 0].imshow(prepare_for_plot(image))
+    axs[0, 0].set_title("Original Image")
+
+    # Plot the feature that contributed the most positively
+    temp, mask = explanation.get_image_and_mask(explanation.top_labels[0], positive_only=True, num_features=5, hide_rest=True)
+    axs[0, 1].imshow(prepare_for_plot(mark_boundaries(temp, mask)))
+    axs[0, 1].set_title("Top Positive Features")
+
+    # Plot the features that contributed the most positively and negatively
+    temp, mask = explanation.get_image_and_mask(explanation.top_labels[0], positive_only=False, num_features=1000, hide_rest=False, min_weight=0.1)
+    axs[1, 0].imshow(mark_boundaries(prepare_for_plot(temp), mask))
+    axs[1, 0].set_title("Top Positive and Negative Features")
+
+    # Plot a heatmap of the features
+    ind = explanation.top_labels[0]
+    dict_heatmap = dict(explanation.local_exp[ind])
+    heatmap = np.vectorize(dict_heatmap.get)(explanation.segments)
+    im = axs[1, 1].imshow(heatmap, cmap='RdBu', vmin=-heatmap.max(), vmax=heatmap.max())
+    axs[1, 1].set_title("Feature Heatmap")
+    fig.colorbar(im, ax=axs[1, 1])
+
+    plt.tight_layout()
+    
+    # If classification mode is enabled, save in the appropriate directory
+    # check if the basename is an jpg image
+    if args.classify:
+        # Extract the class name and correctness from the image path
+        path_parts = args.image_path.split(os.sep)
+        class_name = path_parts[-3]
+        correctness = path_parts[-2] # correct or mistake
+        assert correctness in ['correct', 'mistake'], "The image path should contain 'correct' or 'mistake'"
+
+        # Create the full save path under the explanations directory
+        save_path = os.path.join('explanations', class_name, correctness, os.path.basename(args.image_path))
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+
+        # Save the explanation
+        plt.savefig(save_path, dpi=300)
+        print(f"Explanation saved at {save_path}")
+    else:
+        # make dir for storing the explanations and save it there with the same name as the image
+        os.makedirs("./explanations", exist_ok=True)
+        plt.savefig(f"./explanations/{os.path.basename(args.image_path)}")
+        print(f"Explanation saved at ./explanations/{os.path.basename(args.image_path)}")

utils/train_utils.py

@@ -0,0 +1,251 @@
+from torchvision import models
+import torch.nn as nn
+from tqdm import tqdm
+import torch
+import mlflow
+from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score
+from sklearn.ensemble import RandomForestClassifier
+
+
+# Define the training loop
+def train_one_epoch(model, trainloader, criterion, optimizer, device):
+    model.train()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+
+    for images, labels in tqdm(trainloader, desc="Training", leave=False):
+        images, labels = images.to(device), labels.to(device)
+
+        # Forward pass
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+
+        # Backward pass and optimization
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Track loss and accuracy
+        running_loss += loss.item()
+        _, predicted = outputs.max(1)
+        correct += predicted.eq(labels).sum().item()
+        total += labels.size(0)
+
+    epoch_loss = running_loss / len(trainloader)
+    epoch_accuracy = 100.0 * correct / total
+    return epoch_loss, epoch_accuracy
+
+
+# Define the evaluation loop
+@torch.no_grad()
+def evaluate(model, testloader, criterion, device):
+    model.eval()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+    all_labels = []
+    all_predictions = []
+
+    for images, labels in tqdm(testloader, desc="Evaluating", leave=False):
+        images, labels = images.to(device), labels.to(device)
+
+        # Forward pass
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+
+        # Track loss and accuracy
+        running_loss += loss.item()
+        _, predicted = outputs.max(1)
+        correct += predicted.eq(labels).sum().item()
+        total += labels.size(0)
+
+        # Collect all labels and predictions for metrics
+        all_labels.extend(labels.cpu().numpy())
+        all_predictions.extend(predicted.cpu().numpy())
+
+    epoch_loss = running_loss / len(testloader)
+
+    # Calculate accuracy, precision, recall, and F1-score
+    epoch_accuracy = accuracy_score(all_labels, all_predictions, normalize=True) * 100
+    precision = precision_score(all_labels, all_predictions, average="weighted")
+    recall = recall_score(all_labels, all_predictions, average="weighted")
+    f1 = f1_score(all_labels, all_predictions, average="weighted")
+
+    return epoch_loss, epoch_accuracy, precision, recall, f1
+
+
+# Define the pipeline
+def train_and_evaluate(
+    model,
+    trainloader,
+    testloader,
+    criterion,
+    optimizer,
+    device,
+    epochs,
+    use_mlflow=False,
+):
+    """
+    Train and evaluate the model.
+
+    Args:
+        model (nn.Module): The neural network model.
+        trainloader (DataLoader): DataLoader for training data.
+        testloader (DataLoader): DataLoader for test data.
+        criterion (nn.Module): Loss function.
+        optimizer (optim.Optimizer): Optimizer.
+        device (torch.device): Device to train on ('cuda' or 'cpu').
+        epochs (int): Number of epochs to train.
+
+    Returns:
+        dict: Training and evaluation statistics.
+    """
+    history = {
+        "train_loss": [],
+        "train_acc": [],
+        "test_loss": [],
+        "test_acc": [],
+        "precision": [],
+        "recall": [],
+        "f1": [],
+    }
+
+    model.to(device)
+
+    for epoch in range(epochs):
+        print(f"Epoch {epoch + 1}/{epochs}")
+
+        # Train for one epoch
+        train_loss, train_acc = train_one_epoch(
+            model, trainloader, criterion, optimizer, device
+        )
+        print(f"Train Loss: {train_loss:.4f}, Train Accuracy: {train_acc:.2f}%")
+
+        # Evaluate the model
+        test_loss, test_acc, precision, recall, f1 = evaluate(
+            model, testloader, criterion, device
+        )
+        print(f"Test Loss: {test_loss:.4f}, Test Accuracy: {test_acc:.2f}%")
+
+        # Save statistics
+        history["train_loss"].append(train_loss)
+        history["train_acc"].append(train_acc)
+        history["test_loss"].append(test_loss)
+        history["test_acc"].append(test_acc)
+        history["precision"].append(precision)
+        history["recall"].append(recall)
+        history["f1"].append(f1)
+
+        if use_mlflow:
+            mlflow.log_metric("epoch", epoch)
+            mlflow.log_metric("train_loss", train_loss)
+            mlflow.log_metric("train_acc", train_acc)
+            mlflow.log_metric("test_loss", test_loss)
+            mlflow.log_metric("test_acc", test_acc)
+            mlflow.log_metric("precision", precision)
+            mlflow.log_metric("recall", recall)
+            mlflow.log_metric("f1", f1)
+    return history
+
+
+def set_parameter_requires_grad(model, feature_extracting):
+    if feature_extracting:
+        for param in model.parameters():
+            param.requires_grad = False
+
+
+def initialize_model(
+    model_name,
+    num_classes,
+    feature_extract=True,
+    use_pretrained=True,
+    hidden_size=512,
+    image_shape=(224, 224, 3),
+):
+    # Initialize these variables which will be set in this if statement. Each of these
+    #   variables is model specific.
+    model_ft = None
+
+    if model_name == "resnet":
+        """ Resnet18
+        """
+        model_ft = models.resnet18(pretrained=use_pretrained)
+        set_parameter_requires_grad(model_ft, feature_extract)
+        num_ftrs = model_ft.fc.in_features
+        model_ft.fc = nn.Linear(num_ftrs, num_classes)
+
+    elif model_name == "alexnet":
+        """ Alexnet
+        """
+        model_ft = models.alexnet(pretrained=use_pretrained)
+        set_parameter_requires_grad(model_ft, feature_extract)
+        num_ftrs = model_ft.classifier[6].in_features
+        model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
+
+    elif model_name == "vgg":
+        """ VGG11_bn
+        """
+        model_ft = models.vgg11_bn(pretrained=use_pretrained)
+        set_parameter_requires_grad(model_ft, feature_extract)
+        num_ftrs = model_ft.classifier[6].in_features
+        model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
+
+    elif model_name == "squeezenet":
+        """ Squeezenet
+        """
+        model_ft = models.squeezenet1_0(pretrained=use_pretrained)
+        set_parameter_requires_grad(model_ft, feature_extract)
+        model_ft.classifier[1] = nn.Conv2d(
+            512, num_classes, kernel_size=(1, 1), stride=(1, 1)
+        )
+        model_ft.num_classes = num_classes
+
+    elif model_name == "densenet":
+        """ Densenet
+        """
+        model_ft = models.densenet121(pretrained=use_pretrained)
+        set_parameter_requires_grad(model_ft, feature_extract)
+        num_ftrs = model_ft.classifier.in_features
+        model_ft.classifier = nn.Linear(num_ftrs, num_classes)
+
+    elif model_name == "custom_mlp":
+        """ Custom MLP
+        """
+        model_ft = nn.Sequential(
+            nn.Linear(image_shape[0] * image_shape[1] * image_shape[2], hidden_size),
+            nn.ReLU(),
+            nn.Linear(hidden_size, hidden_size),
+            nn.ReLU(),
+            nn.Linear(hidden_size, hidden_size // 2),
+            nn.ReLU(),
+            nn.Linear(hidden_size // 2, num_classes),
+        )
+    elif model_name == "custom_cnn":
+        """ Custom CNN
+        """
+        model_ft = nn.Sequential(
+            nn.Conv2d(3, 16, 3, 1, 1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(16, 32, 3, 1, 1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(32, 64, 3, 1, 1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Flatten(),
+            nn.Linear(64 * 28 * 28, hidden_size),
+            nn.ReLU(),
+            nn.Linear(hidden_size, num_classes),
+        )
+    elif model_name == "random_forest":
+        """ Random Forest
+        """
+        model_ft = RandomForestClassifier(n_estimators=100, random_state=42)
+
+    else:
+        print("Invalid model name, exiting...")
+        exit()
+
+    return model_ft