Add application file

app.py

@@ -0,0 +1,42 @@
+import gradio as gr
+import os 
+import torch 
+
+from src.model import resnet_model
+from timeit import default_timer as timer
+from typing import Tuple,Dict
+
+class_names = ["CNV","DME","DRUSEN","NORMAL"]
+
+resnet, resnet_transforms = resnet_model(num_classes=4)
+
+
+state_dict = torch.load(f="models/model.pth", map_location=torch.device("cpu"))
+resnet.load_state_dict(state_dict, strict=False)
+
+def predict(img) -> Tuple[Dict,float]:
+    start_time = timer()
+    
+    img = resnet_transforms(img).unsqueeze(0)
+    resnet.eval()
+    with torch.inference_mode():
+        pred_probs = torch.softmax(resnet(img),dim=1)
+        
+    pred_label_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
+    pred_time = round(timer() - start_time,5)
+    return pred_label_and_probs, pred_time
+
+example_paths = [["examples/" + example] for example in os.listdir("examples")]
+
+title = "Retinal disease detection using Optical Tomography Images 👁️"
+description = " This application uses Optical Coherence Tomography (OCT) images to assist in the identification of retinal conditions such as CNV, DME, DRUSEN, and NORMAL. The tool provides predictions based on the uploaded image and displays the processing time for the analysis. Please note that this tool is intended for educational and research purposes only. It is not a substitute for professional medical advice or diagnosis. For any medical concerns, please consult a healthcare professional."
+
+gradio_interface = gr.Interface(fn=predict,
+                    inputs=gr.Image(type="pil"),
+                    outputs=[gr.Label(num_top_classes=4,label="Predictions"),
+                                gr.Number(label="prediction time: ")],
+                                title=title,
+                                examples=example_paths,
+                                description=description)
+                            
+gradio_interface.launch()

models/model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e011dfb570e65f9d2cb87c42b51b326e0093f2d9b034d7067b90ee34f46b9e3
+size 94364842

requirements.txt

@@ -0,0 +1,4 @@
+torch 
+torchvision
+loguru
+gradio 

src/__init__.py

@@ -0,0 +1 @@
+4

src/data_setup.py

@@ -0,0 +1,62 @@
+import os
+import torchvision
+
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+
+NUM_WORKERS = os.cpu_count()
+
+def create_dataloaders(
+    train_dir: str, 
+    test_dir: str, 
+    transform: transforms.Compose, 
+    batch_size: int, 
+    num_workers: int=NUM_WORKERS
+):
+  """Creates training and testing DataLoaders.
+
+  Takes in a training directory and testing directory path and turns
+  them into PyTorch Datasets and then into PyTorch DataLoaders.
+
+  Args:
+    train_dir: Path to training directory.
+    test_dir: Path to testing directory.
+    transform: torchvision transforms to perform on training and testing data.
+    batch_size: Number of samples per batch in each of the DataLoaders.
+    num_workers: An integer for number of workers per DataLoader.
+
+  Returns:
+    A tuple of (train_dataloader, test_dataloader, class_names).
+    Where class_names is a list of the target classes.
+    Example usage:
+      train_dataloader, test_dataloader, class_names = \
+        = create_dataloaders(train_dir=path/to/train_dir,
+                              test_dir=path/to/test_dir,
+                              transform=some_transform,
+                              batch_size=32,
+                              num_workers=4)
+  """
+  # Use ImageFolder to create dataset(s)
+  train_data = datasets.ImageFolder(train_dir, transform=transform)
+  test_data = datasets.ImageFolder(test_dir, transform=transform)
+
+  # Get class names
+  class_names = train_data.classes
+
+  # Turn images into data loaders
+  train_dataloader = DataLoader(
+      train_data,
+      batch_size=batch_size,
+      shuffle=True,
+      num_workers=num_workers,
+      pin_memory=True,
+  )
+  test_dataloader = DataLoader(
+      test_data,
+      batch_size=batch_size,
+      shuffle=False, 
+      num_workers=num_workers,
+      pin_memory=True,
+  )
+
+  return train_dataloader, test_dataloader, class_names

src/engine.py

@@ -0,0 +1,107 @@
+import torch
+import mlflow
+from tqdm.auto import tqdm
+from src.logger import global_logger as logger
+from typing import Dict, List, Tuple
+
+def train_step(model: torch.nn.Module, 
+                dataloader: torch.utils.data.DataLoader, 
+                loss_fn: torch.nn.Module, 
+                optimizer: torch.optim.Optimizer,
+                device: torch.device) -> Tuple[float, float]:
+    """Trains a PyTorch model for a single epoch."""
+    
+
+    model.train()
+    train_loss, train_acc = 0, 0
+
+    for batch, (X, y) in enumerate(dataloader):
+        X, y = X.to(device), y.to(device)
+        y_pred = model(X)
+        loss = loss_fn(y_pred, y)
+        train_loss += loss.item()
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)
+        train_acc += (y_pred_class == y).sum().item() / len(y_pred)
+
+    train_loss /= len(dataloader)
+    train_acc /= len(dataloader)
+    return train_loss, train_acc
+
+
+def test_step(model: torch.nn.Module, 
+                dataloader: torch.utils.data.DataLoader, 
+                loss_fn: torch.nn.Module,
+                device: torch.device) -> Tuple[float, float]:
+    """Tests a PyTorch model for a single epoch."""
+
+    model.eval()
+    test_loss, test_acc = 0, 0
+
+    with torch.inference_mode():
+        for batch, (X, y) in enumerate(dataloader):
+            X, y = X.to(device), y.to(device)
+            test_pred_logits = model(X)
+            loss = loss_fn(test_pred_logits, y)
+            test_loss += loss.item()
+            test_pred_labels = test_pred_logits.argmax(dim=1)
+            test_acc += (test_pred_labels == y).sum().item() / len(test_pred_labels)
+
+    test_loss /= len(dataloader)
+    test_acc /= len(dataloader)
+    return test_loss, test_acc
+
+def train(model: torch.nn.Module, 
+            train_dataloader: torch.utils.data.DataLoader, 
+            test_dataloader: torch.utils.data.DataLoader, 
+            optimizer: torch.optim.Optimizer,
+            loss_fn: torch.nn.Module,
+            epochs: int,
+            device: torch.device) -> Dict[str, List[float]]:
+    """Trains and tests a PyTorch model."""
+
+    results = {
+        "train_loss": [],
+        "train_acc": [],
+        "test_loss": [],
+        "test_acc": []
+    }
+    
+    for epoch in tqdm(range(epochs)):
+        with mlflow.start_run() as run:
+            mlflow.log_param("epoch", epoch)
+            mlflow.log_param("optimizer", optimizer.__class__.__name__)
+            mlflow.log_param("loss_fn", loss_fn.__class__.__name__)
+
+            train_loss, train_acc = train_step(model=model,
+                                                dataloader=train_dataloader,
+                                                loss_fn=loss_fn,
+                                                optimizer=optimizer,
+                                                device=device)
+            test_loss, test_acc = test_step(model=model,
+                                            dataloader=test_dataloader,
+                                            loss_fn=loss_fn,
+                                            device=device)
+
+            print(
+                f"Epoch: {epoch+1} | "
+                f"train_loss: {train_loss:.3f} | "
+                f"train_acc: {train_acc:.3f} | "
+                f"test_loss: {test_loss:.3f} | "
+                f"test_acc: {test_acc:.3f}"
+            )
+
+            results["train_loss"].append(train_loss)
+            results["train_acc"].append(train_acc)
+            results["test_loss"].append(test_loss)
+            results["test_acc"].append(test_acc)
+
+            mlflow.log_metric("train_loss", train_loss, step=epoch)
+            mlflow.log_metric("train_acc", train_acc, step=epoch)
+            mlflow.log_metric("test_loss", test_loss, step=epoch)
+            mlflow.log_metric("test_acc", test_acc, step=epoch)
+            mlflow.pytorch.log_model(model, "model")
+
+    return results

src/helper_function.py

@@ -0,0 +1,249 @@
+"""
+A series of helper functions used throughout the course.
+
+If a function gets defined once and could be used over and over, it'll go in here.
+"""
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+
+from torch import nn
+
+import os
+import zipfile
+
+from pathlib import Path
+
+import requests
+
+# Walk through an image classification directory and find out how many files (images)
+# are in each subdirectory.
+import os
+
+def walk_through_dir(dir_path):
+    """
+    Walks through dir_path returning its contents.
+    Args:
+    dir_path (str): target directory
+
+    Returns:
+    A print out of:
+      number of subdiretories in dir_path
+      number of images (files) in each subdirectory
+      name of each subdirectory
+    """
+    for dirpath, dirnames, filenames in os.walk(dir_path):
+        print(f"There are {len(dirnames)} directories and {len(filenames)} images in '{dirpath}'.")
+
+def plot_decision_boundary(model: torch.nn.Module, X: torch.Tensor, y: torch.Tensor):
+    """Plots decision boundaries of model predicting on X in comparison to y.
+
+    Source - https://madewithml.com/courses/foundations/neural-networks/ (with modifications)
+    """
+    # Put everything to CPU (works better with NumPy + Matplotlib)
+    model.to("cpu")
+    X, y = X.to("cpu"), y.to("cpu")
+
+    # Setup prediction boundaries and grid
+    x_min, x_max = X[:, 0].min() - 0.1, X[:, 0].max() + 0.1
+    y_min, y_max = X[:, 1].min() - 0.1, X[:, 1].max() + 0.1
+    xx, yy = np.meshgrid(np.linspace(x_min, x_max, 101), np.linspace(y_min, y_max, 101))
+
+    # Make features
+    X_to_pred_on = torch.from_numpy(np.column_stack((xx.ravel(), yy.ravel()))).float()
+
+    # Make predictions
+    model.eval()
+    with torch.inference_mode():
+        y_logits = model(X_to_pred_on)
+
+    # Test for multi-class or binary and adjust logits to prediction labels
+    if len(torch.unique(y)) > 2:
+        y_pred = torch.softmax(y_logits, dim=1).argmax(dim=1)  # mutli-class
+    else:
+        y_pred = torch.round(torch.sigmoid(y_logits))  # binary
+
+    # Reshape preds and plot
+    y_pred = y_pred.reshape(xx.shape).detach().numpy()
+    plt.contourf(xx, yy, y_pred, cmap=plt.cm.RdYlBu, alpha=0.7)
+    plt.scatter(X[:, 0], X[:, 1], c=y, s=40, cmap=plt.cm.RdYlBu)
+    plt.xlim(xx.min(), xx.max())
+    plt.ylim(yy.min(), yy.max())
+
+
+# Plot linear data or training and test and predictions (optional)
+def plot_predictions(
+    train_data, train_labels, test_data, test_labels, predictions=None
+):
+    """
+  Plots linear training data and test data and compares predictions.
+  """
+    plt.figure(figsize=(10, 7))
+
+    # Plot training data in blue
+    plt.scatter(train_data, train_labels, c="b", s=4, label="Training data")
+
+    # Plot test data in green
+    plt.scatter(test_data, test_labels, c="g", s=4, label="Testing data")
+
+    if predictions is not None:
+        # Plot the predictions in red (predictions were made on the test data)
+        plt.scatter(test_data, predictions, c="r", s=4, label="Predictions")
+
+    # Show the legend
+    plt.legend(prop={"size": 14})
+
+
+# Calculate accuracy (a classification metric)
+def accuracy_fn(y_true, y_pred):
+    """Calculates accuracy between truth labels and predictions.
+
+    Args:
+        y_true (torch.Tensor): Truth labels for predictions.
+        y_pred (torch.Tensor): Predictions to be compared to predictions.
+
+    Returns:
+        [torch.float]: Accuracy value between y_true and y_pred, e.g. 78.45
+    """
+    correct = torch.eq(y_true, y_pred).sum().item()
+    acc = (correct / len(y_pred)) * 100
+    return acc
+
+
+def print_train_time(start, end, device=None):
+    """Prints difference between start and end time.
+
+    Args:
+        start (float): Start time of computation (preferred in timeit format). 
+        end (float): End time of computation.
+        device ([type], optional): Device that compute is running on. Defaults to None.
+
+    Returns:
+        float: time between start and end in seconds (higher is longer).
+    """
+    total_time = end - start
+    print(f"\nTrain time on {device}: {total_time:.3f} seconds")
+    return total_time
+
+
+# Plot loss curves of a model
+def plot_loss_curves(results):
+    """Plots training curves of a results dictionary.
+
+    Args:
+        results (dict): dictionary containing list of values, e.g.
+            {"train_loss": [...],
+             "train_acc": [...],
+             "test_loss": [...],
+             "test_acc": [...]}
+    """
+    loss = results["train_loss"]
+    test_loss = results["test_loss"]
+
+    accuracy = results["train_acc"]
+    test_accuracy = results["test_acc"]
+
+    epochs = range(len(results["train_loss"]))
+
+    plt.figure(figsize=(15, 7))
+
+    # Plot loss
+    plt.subplot(1, 2, 1)
+    plt.plot(epochs, loss, label="train_loss")
+    plt.plot(epochs, test_loss, label="test_loss")
+    plt.title("Loss")
+    plt.xlabel("Epochs")
+    plt.legend()
+
+    # Plot accuracy
+    plt.subplot(1, 2, 2)
+    plt.plot(epochs, accuracy, label="train_accuracy")
+    plt.plot(epochs, test_accuracy, label="test_accuracy")
+    plt.title("Accuracy")
+    plt.xlabel("Epochs")
+    plt.legend()
+
+
+# Pred and plot image function from notebook 04
+# See creation: https://www.learnpytorch.io/04_pytorch_custom_datasets/#113-putting-custom-image-prediction-together-building-a-function
+from typing import List
+import torchvision
+
+
+def pred_and_plot_image(
+    model: torch.nn.Module,
+    image_path: str,
+    class_names: List[str] = None,
+    transform=None,
+    device: torch.device = "cuda" if torch.cuda.is_available() else "cpu",
+):
+    """Makes a prediction on a target image with a trained model and plots the image.
+
+    Args:
+        model (torch.nn.Module): trained PyTorch image classification model.
+        image_path (str): filepath to target image.
+        class_names (List[str], optional): different class names for target image. Defaults to None.
+        transform (_type_, optional): transform of target image. Defaults to None.
+        device (torch.device, optional): target device to compute on. Defaults to "cuda" if torch.cuda.is_available() else "cpu".
+    
+    Returns:
+        Matplotlib plot of target image and model prediction as title.
+
+    Example usage:
+        pred_and_plot_image(model=model,
+                            image="some_image.jpeg",
+                            class_names=["class_1", "class_2", "class_3"],
+                            transform=torchvision.transforms.ToTensor(),
+                            device=device)
+    """
+
+    # 1. Load in image and convert the tensor values to float32
+    target_image = torchvision.io.read_image(str(image_path)).type(torch.float32)
+
+    # 2. Divide the image pixel values by 255 to get them between [0, 1]
+    target_image = target_image / 255.0
+
+    # 3. Transform if necessary
+    if transform:
+        target_image = transform(target_image)
+
+    # 4. Make sure the model is on the target device
+    model.to(device)
+
+    # 5. Turn on model evaluation mode and inference mode
+    model.eval()
+    with torch.inference_mode():
+        # Add an extra dimension to the image
+        target_image = target_image.unsqueeze(dim=0)
+
+        # Make a prediction on image with an extra dimension and send it to the target device
+        target_image_pred = model(target_image.to(device))
+
+    # 6. Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)
+    target_image_pred_probs = torch.softmax(target_image_pred, dim=1)
+
+    # 7. Convert prediction probabilities -> prediction labels
+    target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)
+
+    # 8. Plot the image alongside the prediction and prediction probability
+    plt.imshow(
+        target_image.squeeze().permute(1, 2, 0)
+    )  # make sure it's the right size for matplotlib
+    if class_names:
+        title = f"Pred: {class_names[target_image_pred_label.cpu()]} | Prob: {target_image_pred_probs.max().cpu():.3f}"
+    else:
+        title = f"Pred: {target_image_pred_label} | Prob: {target_image_pred_probs.max().cpu():.3f}"
+    plt.title(title)
+    plt.axis(False)
+
+def set_seeds(seed: int=42):
+    """Sets random sets for torch operations.
+
+    Args:
+        seed (int, optional): Random seed to set. Defaults to 42.
+    """
+    # Set the seed for general torch operations
+    torch.manual_seed(seed)
+    # Set the seed for CUDA torch operations (ones that happen on the GPU)
+    torch.cuda.manual_seed(seed)
+

src/logger.py

@@ -0,0 +1,17 @@
+from loguru import logger
+import os
+
+class LoguruLogger:
+    def __init__(self, log_folder="logs", log_file="training.log"):
+        os.makedirs(log_folder, exist_ok=True)
+        log_path = os.path.join(log_folder, log_file)
+
+        logger.remove()  # Remove default logger
+        logger.add(log_path, level="INFO", format="{time} - {name} - {level} - {message}")
+        logger.add(lambda msg: print(msg, end=""), level="INFO", format="{time} - {name} - {level} - {message}")
+
+    def get_logger(self):
+        return logger
+
+# Create a global logger instance
+global_logger = LoguruLogger().get_logger()

src/model.py

@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+import torchvision
+from src.logger import global_logger as logger
+from torchvision.models import resnet50, ResNet50_Weights
+
+def resnet_model(num_classes: int = 4, seed: int = 42):
+    # Load pretrained ResNet18 model
+    weights = ResNet50_Weights.DEFAULT
+    model = resnet50(weights=weights)
+
+    # Freeze the parameters of the pretrained model
+    for param in model.parameters():
+        param.requires_grad = False
+
+    #logger.info("Model initialized with frozen ResNet18 backbone and new fully connected layers.")
+
+    # Replace the final fully connected layer with a new one
+    torch.manual_seed(seed)
+    model.fc = nn.Sequential(
+        nn.Dropout(p=0.3, inplace=True),
+        nn.Linear(in_features=model.fc.in_features, out_features=num_classes),
+    )
+
+    # Define the transforms using the predefined transforms from weights
+    transforms = weights.transforms()
+
+    return model, transforms
+
+# Example usage
+model, transforms = resnet_model(num_classes=4)

src/train.py

@@ -0,0 +1,62 @@
+import os
+import torch
+import src.data_setup as data_setup
+import src.engine as engine
+import src.utils as utils
+from src.logger import global_logger as logger
+from torchvision import transforms
+import src.model as model_module
+
+def main():
+
+    NUM_EPOCHS = 20
+    BATCH_SIZE = 32
+    LEARNING_RATE = 0.001
+
+    train_dir = "data\\retinal_oct\\train"
+    test_dir = "data\\retinal_oct\\test"
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Use the transformations required by ResNet50
+    data_transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+
+    train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(
+        train_dir=train_dir,
+        test_dir=test_dir,
+        transform=data_transform,
+        batch_size=BATCH_SIZE
+    )
+    logger.info("Data transformed successfully.")
+
+    # Initialize the ResNet50 model
+    model, _ = model_module.resnet_model(num_classes=len(class_names))
+    model = model.to(device)
+
+    loss_fn = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
+
+    engine.train(
+        model=model,
+        train_dataloader=train_dataloader,
+        test_dataloader=test_dataloader,
+        loss_fn=loss_fn,
+        optimizer=optimizer,
+        epochs=NUM_EPOCHS,
+        device=device
+    )
+    
+    utils.save_model(
+        model=model,
+        target_dir="models",
+        model_name="model.pth"
+    )
+    logger.info("Model trained successfully.")
+    logger.info("Model saved to models folder.")
+
+if __name__ == '__main__':
+    main()

src/utils.py

@@ -0,0 +1,28 @@
+import torch
+from pathlib import Path
+
+def save_model(model: torch.nn.Module,
+                target_dir: str,
+                model_name: str):
+  """Saves a PyTorch model to a target directory.
+
+  Args:
+    model: A target PyTorch model to save.
+    target_dir: A directory for saving the model to.
+    model_name: A filename for the saved model. Should include
+      either ".pth" or ".pt" as the file extension.
+
+  """
+  # Create target directory
+  target_dir_path = Path(target_dir)
+  target_dir_path.mkdir(parents=True,
+                        exist_ok=True)
+
+  # Create model save path
+  assert model_name.endswith(".pth") or model_name.endswith(".pt"), "model_name should end with '.pt' or '.pth'"
+  model_save_path = target_dir_path / model_name
+
+  # Save the model state_dict()
+  print(f"[INFO] Saving model to: {model_save_path}")
+  torch.save(obj=model.state_dict(),
+              f=model_save_path)