Add application file

app.py

@@ -0,0 +1,445 @@
+import torch
+import torchvision
+from torchvision.models import convnext_base, ConvNeXt_Base_Weights
+from torchvision.models._api import WeightsEnum
+from torch.hub import load_state_dict_from_url
+from statistics import mean
+import time, os
+import torch
+from torch.utils.data import Dataset
+from torchvision import datasets
+from torchvision.transforms import ToTensor
+import matplotlib.pyplot as plt
+from torch.utils.data import DataLoader
+import gradio as gr              
+from torchvision import transforms
+import matplotlib.pyplot as plt
+
+from typing import List, Tuple
+
+from PIL import Image
+from torch.utils.data import Subset
+from torch import nn
+
+from tqdm.auto import tqdm
+from typing import Dict, List, Tuple
+
+"""
+Contains functionality for creating PyTorch DataLoaders for 
+image classification data.
+"""
+import os
+
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+
+import torch
+
+from tqdm.auto import tqdm
+from typing import Dict, List, Tuple
+import torch
+import torchvision
+from torchvision import transforms
+import matplotlib.pyplot as plt
+
+from typing import List, Tuple
+
+from PIL import Image
+
+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
+
+"""
+Contains functions for training and testing a PyTorch model.
+"""
+
+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.
+
+    Turns a target PyTorch model to training mode and then
+    runs through all of the required training steps (forward
+    pass, loss calculation, optimizer step).
+
+    Args:
+    model: A PyTorch model to be trained.
+    dataloader: A DataLoader instance for the model to be trained on.
+    loss_fn: A PyTorch loss function to minimize.
+    optimizer: A PyTorch optimizer to help minimize the loss function.
+    device: A target device to compute on (e.g. "cuda" or "cpu").
+
+    Returns:
+    A tuple of training loss and training accuracy metrics.
+    In the form (train_loss, train_accuracy). For example:
+
+    (0.1112, 0.8743)
+    """
+    # Put model in train mode
+    model.train()
+
+    # Setup train loss and train accuracy values
+    train_loss, train_acc = 0, 0
+
+    # Loop through data loader data batches
+    for batch, (X, y) in enumerate(dataloader):
+        # Send data to target device
+        X, y = X.to(device), y.to(device)
+
+        # 1. Forward pass
+        y_pred = model(X)
+
+        # 2. Calculate  and accumulate loss
+        loss = loss_fn(y_pred, y)
+        train_loss += loss.item() 
+
+        # 3. Optimizer zero grad
+        optimizer.zero_grad()
+
+        # 4. Loss backward
+        loss.backward()
+
+        # 5. Optimizer step
+        optimizer.step()
+
+        # Calculate and accumulate accuracy metric across all batches
+        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)
+        train_acc += (y_pred_class == y).sum().item()/len(y_pred)
+
+    # Adjust metrics to get average loss and accuracy per batch 
+    train_loss = train_loss / len(dataloader)
+    train_acc = 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.
+
+    Turns a target PyTorch model to "eval" mode and then performs
+    a forward pass on a testing dataset.
+
+    Args:
+    model: A PyTorch model to be tested.
+    dataloader: A DataLoader instance for the model to be tested on.
+    loss_fn: A PyTorch loss function to calculate loss on the test data.
+    device: A target device to compute on (e.g. "cuda" or "cpu").
+
+    Returns:
+    A tuple of testing loss and testing accuracy metrics.
+    In the form (test_loss, test_accuracy). For example:
+
+    (0.0223, 0.8985)
+    """
+    # Put model in eval mode
+    model.eval() 
+
+    # Setup test loss and test accuracy values
+    test_loss, test_acc = 0, 0
+
+    # Turn on inference context manager
+    with torch.inference_mode():
+        # Loop through DataLoader batches
+        for batch, (X, y) in enumerate(dataloader):
+            # Send data to target device
+            X, y = X.to(device), y.to(device)
+
+            # 1. Forward pass
+            test_pred_logits = model(X)
+
+            # 2. Calculate and accumulate loss
+            loss = loss_fn(test_pred_logits, y)
+            test_loss += loss.item()
+
+            # Calculate and accumulate accuracy
+            test_pred_labels = test_pred_logits.argmax(dim=1)
+            test_acc += ((test_pred_labels == y).sum().item()/len(test_pred_labels))
+
+    # Adjust metrics to get average loss and accuracy per batch 
+    test_loss = test_loss / len(dataloader)
+    test_acc = 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]:
+    """Trains and tests a PyTorch model.
+
+    Passes a target PyTorch models through train_step() and test_step()
+    functions for a number of epochs, training and testing the model
+    in the same epoch loop.
+
+    Calculates, prints and stores evaluation metrics throughout.
+
+    Args:
+    model: A PyTorch model to be trained and tested.
+    train_dataloader: A DataLoader instance for the model to be trained on.
+    test_dataloader: A DataLoader instance for the model to be tested on.
+    optimizer: A PyTorch optimizer to help minimize the loss function.
+    loss_fn: A PyTorch loss function to calculate loss on both datasets.
+    epochs: An integer indicating how many epochs to train for.
+    device: A target device to compute on (e.g. "cuda" or "cpu").
+
+    Returns:
+    A dictionary of training and testing loss as well as training and
+    testing accuracy metrics. Each metric has a value in a list for 
+    each epoch.
+    In the form: {train_loss: [...],
+              train_acc: [...],
+              test_loss: [...],
+              test_acc: [...]} 
+    For example if training for epochs=2: 
+             {train_loss: [2.0616, 1.0537],
+              train_acc: [0.3945, 0.3945],
+              test_loss: [1.2641, 1.5706],
+              test_acc: [0.3400, 0.2973]} 
+    """
+    # Create empty results dictionary
+    results = {"train_loss": [],
+               "train_acc": [],
+               "test_loss": [],
+               "test_acc": []
+    }
+    
+    # Make sure model on target device
+    model.to(device)
+
+    # Loop through training and testing steps for a number of epochs
+    for epoch in tqdm(range(epochs)):
+        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 out what's happening
+        print(
+          f"Epoch: {epoch+1} | "
+          f"train_loss: {train_loss:.4f} | "
+          f"train_acc: {train_acc:.4f} | "
+          f"test_loss: {test_loss:.4f} | "
+          f"test_acc: {test_acc:.4f}"
+        )
+
+        # Update results dictionary
+        results["train_loss"].append(train_loss)
+        results["train_acc"].append(train_acc)
+        results["test_loss"].append(test_loss)
+        results["test_acc"].append(test_acc)
+
+    # Return the filled results at the end of the epochs
+    return results
+
+"""
+Utility functions to make predictions.
+
+Main reference for code creation: https://www.learnpytorch.io/06_pytorch_transfer_learning/#6-make-predictions-on-images-from-the-test-set 
+"""
+
+# Set device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Predict on a target image with a target model
+# Function created in: https://www.learnpytorch.io/06_pytorch_transfer_learning/#6-make-predictions-on-images-from-the-test-set
+def pred_and_plot_image(
+    model: torch.nn.Module,
+    class_names: List[str],
+    image_path: str,
+    image_size: Tuple[int, int] = (224, 224),
+    transform: torchvision.transforms = None,
+    device: torch.device = device,
+):
+    """Predicts on a target image with a target model.
+
+    Args:
+        model (torch.nn.Module): A trained (or untrained) PyTorch model to predict on an image.
+        class_names (List[str]): A list of target classes to map predictions to.
+        image_path (str): Filepath to target image to predict on.
+        image_size (Tuple[int, int], optional): Size to transform target image to. Defaults to (224, 224).
+        transform (torchvision.transforms, optional): Transform to perform on image. Defaults to None which uses ImageNet normalization.
+        device (torch.device, optional): Target device to perform prediction on. Defaults to device.
+    """
+
+    # Open image
+    img = Image.open(image_path)
+
+    # Create transformation for image (if one doesn't exist)
+    if transform is not None:
+        image_transform = transform
+    else:
+        image_transform = transforms.Compose(
+            [
+                transforms.Resize(image_size),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+            ]
+        )
+
+    ### Predict on image ###
+
+    # Make sure the model is on the target device
+    model.to(device)
+
+    # Turn on model evaluation mode and inference mode
+    model.eval()
+    with torch.inference_mode():
+        # Transform and add an extra dimension to image (model requires samples in [batch_size, color_channels, height, width])
+        transformed_image = image_transform(img).unsqueeze(dim=0)
+
+        # Make a prediction on image with an extra dimension and send it to the target device
+        target_image_pred = model(transformed_image.to(device))
+
+    # Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)
+    target_image_pred_probs = torch.softmax(target_image_pred, dim=1)
+
+    # Convert prediction probabilities -> prediction labels
+    target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)
+
+    # Plot image with predicted label and probability
+    plt.figure()
+    plt.imshow(img)
+    plt.title(
+        f"Pred: {class_names[target_image_pred_label]} | Prob: {target_image_pred_probs.max():.3f}"
+    )
+    plt.axis(False)
+
+BATCH_SIZE = 32
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+training_datab, test_datab = torchvision.datasets.CIFAR10(root="data", train=True, download=True, transform=ConvNeXt_Base_Weights.DEFAULT.transforms()), torchvision.datasets.CIFAR10(root="data", train=False, download=True, transform=ConvNeXt_Base_Weights.DEFAULT.transforms())
+subset_train, subset_test = Subset(training_datab, indices=range(len(training_datab) // 1000)), Subset(test_datab, indices=range(len(test_datab) // 1000))  # delete here IMPORTANT!!!!!!!!!!!!!!!!!!!!!!!
+
+def get_state_dict(self, *args, **kwargs):
+    kwargs.pop("check_hash")
+    return load_state_dict_from_url(self.url, *args, **kwargs)
+WeightsEnum.get_state_dict = get_state_dict
+
+modeld = convnext_base(ConvNeXt_Base_Weights.DEFAULT)
+
+modeld.classifier = nn.Sequential(
+    nn.LayerNorm((1024, 1, 1), eps=1e-06, elementwise_affine=True),
+    nn.Flatten(start_dim=1, end_dim=-1),
+    nn.Linear(in_features=1024, out_features=10, bias=True)
+)
+
+optimizerd = torch.optim.Adam(modeld.parameters(), 0.001)
+
+loss_fn = nn.CrossEntropyLoss()
+epochs = 5
+
+train_dataloaderd, test_dataloaderd = DataLoader(subset_train, batch_size=BATCH_SIZE, shuffle=True), DataLoader(subset_test, batch_size=BATCH_SIZE, shuffle=False) # change data here IMPORTANT!!!!!!!!!!!!!!!!!!!!!!!
+
+# engine.train(modeld, train_dataloaderd, test_dataloaderd, optimizerd, loss_fn, epochs, device)
+
+def pred_image(image_path: str, model: torch.nn.Module = modeld, class_names: List[str] = training_datab.classes, image_size: Tuple[int, int] = (224, 224), transform: torchvision.transforms = ConvNeXt_Base_Weights.DEFAULT.transforms(), device: torch.device = device):
+    # Open image
+    img = Image.open(image_path)
+
+    # Create transformation for image (if one doesn't exist)
+    if transform is not None:
+        image_transform = transform
+    else:
+        image_transform = transforms.Compose(
+            [
+                transforms.Resize(image_size),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+            ]
+        )
+
+    ### Predict on image ###
+
+    # Make sure the model is on the target device
+    model.to(device)
+
+    # Turn on model evaluation mode and inference mode
+    model.eval()
+    with torch.inference_mode():
+        # Transform and add an extra dimension to image (model requires samples in [batch_size, color_channels, height, width])
+        transformed_image = image_transform(img).unsqueeze(dim=0)
+
+        # Make a prediction on image with an extra dimension and send it to the target device
+        target_image_pred = model(transformed_image.to(device))
+
+    # Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)
+    target_image_pred_probs = torch.softmax(target_image_pred, dim=1)
+
+    # Convert prediction probabilities -> prediction labels
+    target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)
+
+    return class_names[target_image_pred_label], target_image_pred_probs.max()
+
+
+demo = gr.Interface(fn=pred_image, inputs=gr.Image(type="filepath"), outputs=[gr.Textbox(label="label"), gr.Textbox(label="probability")], examples=["apple.jpg","bird.jpg","car.jpg","ocean.jpg"]) 
+
+demo.launch(share=True)