good (1).py

# -*- coding: utf-8 -*-
"""Good.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1AkM6wLyspo4q2ScK_pIkTAFDV-Q-JCgh
"""

!pip install torchinfo

files.upload()

from google.colab import files
import matplotlib.pyplot as plt
import torch
import torchvision

from torch import nn
from torchvision import transforms
from Helperfunction import set_seeds

device = "cuda" if torch.cuda.is_available() else "cpu"
device

# Commented out IPython magic to ensure Python compatibility.
# %%writefile predict.py
# 
# #predict
# 
# 
# """
# 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
# """
# import torch
# import torchvision
# from torchvision import transforms
# import matplotlib.pyplot as plt
# 
# from typing import List, Tuple
# 
# from PIL import Image
# 
# # 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)
#

from google.colab import drive
drive.mount('/content/drive')

# Commented out IPython magic to ensure Python compatibility.
# %%writefile model_builder.py
# 
# #model_builder
# 
# """
# Contains PyTorch model code to instantiate a TinyVGG model.
# """
# import torch
# from torch import nn
# 
# class TinyVGG(nn.Module):
#     """Creates the TinyVGG architecture.
# 
#     Replicates the TinyVGG architecture from the CNN explainer website in PyTorch.
#     See the original architecture here: https://poloclub.github.io/cnn-explainer/
# 
#     Args:
#     input_shape: An integer indicating number of input channels.
#     hidden_units: An integer indicating number of hidden units between layers.
#     output_shape: An integer indicating number of output units.
#     """
#     def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
#         super().__init__()
#         self.conv_block_1 = nn.Sequential(
#           nn.Conv2d(in_channels=input_shape,
#                     out_channels=hidden_units,
#                     kernel_size=3,
#                     stride=1,
#                     padding=0),
#           nn.ReLU(),
#           nn.Conv2d(in_channels=hidden_units,
#                     out_channels=hidden_units,
#                     kernel_size=3,
#                     stride=1,
#                     padding=0),
#           nn.ReLU(),
#           nn.MaxPool2d(kernel_size=2,
#                         stride=2)
#         )
#         self.conv_block_2 = nn.Sequential(
#           nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),
#           nn.ReLU(),
#           nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),
#           nn.ReLU(),
#           nn.MaxPool2d(2)
#         )
#         self.classifier = nn.Sequential(
#           nn.Flatten(),
#           # Where did this in_features shape come from?
#           # It's because each layer of our network compresses and changes the shape of our inputs data.
#           nn.Linear(in_features=hidden_units*13*13,
#                     out_features=output_shape)
#         )
# 
#     def forward(self, x: torch.Tensor):
#         x = self.conv_block_1(x)
#         x = self.conv_block_2(x)
#         x = self.classifier(x)
#         return x
#         # return self.classifier(self.block_2(self.block_1(x))) # <- leverage the benefits of operator fusion

# Commented out IPython magic to ensure Python compatibility.
# %%writefile utils.py
# 
# #utils.py
# 
# """
# Contains various utility functions for PyTorch model training and saving.
# """
# 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.
# 
#     Example usage:
#     save_model(model=model_0,
#                target_dir="models",
#                model_name="05_going_modular_tingvgg_model.pth")
#     """
#     # 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)

# Commented out IPython magic to ensure Python compatibility.
# %%writefile engine.py
# #engine.py
# 
# """
# Contains functions for training and testing a PyTorch model.
# """
# import torch
# 
# from tqdm.auto import tqdm
# 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.
# 
#     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

# Commented out IPython magic to ensure Python compatibility.
# %%writefile data_setup.py
# #data_setup.py
# """
# Contains functionality for creating PyTorch DataLoaders for
# image classification data.
# """
# import os
# 
# 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

# Commented out IPython magic to ensure Python compatibility.
# %%writefile train.py
# #train.py only in this cell
# 
# """
# Trains a PyTorch image classification model using device-agnostic code.
# """
# 
# import os
# import torch
# #import data_setup, engine, model_builder, utils
# 
# from torchvision import transforms
# 
# # Setup hyperparameters
# NUM_EPOCHS = 5
# BATCH_SIZE = 32
# HIDDEN_UNITS = 10
# LEARNING_RATE = 0.001
# 
# # Setup directories
# train_dir = "data/pizza_steak_sushi/train"
# test_dir = "data/pizza_steak_sushi/test"
# 
# # Setup target device
# device = "cuda" if torch.cuda.is_available() else "cpu"
# 
# # Create transforms
# data_transform = transforms.Compose([
#   transforms.Resize((64, 64)),
#   transforms.ToTensor()
# ])
# 
# # Create DataLoaders with help from data_setup.py
# 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
# )
# 
# # Create model with help from model_builder.py
# model = model_builder.TinyVGG(
#     input_shape=3,
#     hidden_units=HIDDEN_UNITS,
#     output_shape=len(class_names)
# ).to(device)
# 
# # Set loss and optimizer
# loss_fn = torch.nn.CrossEntropyLoss()
# optimizer = torch.optim.Adam(model.parameters(),
#                              lr=LEARNING_RATE)
# 
# # Start training with help from engine.py
# engine.train(model=model,
#              train_dataloader=train_dataloader,
#              test_dataloader=test_dataloader,
#              loss_fn=loss_fn,
#              optimizer=optimizer,
#              epochs=NUM_EPOCHS,
#              device=device)
# 
# # Save the model with help from utils.py
# utils.save_model(model=model,
#                  target_dir="models",
#                  model_name="05_going_modular_script_mode_tinyvgg_model.pth")
# 
# 
# 
#

!python /content/data_setup.py/train.py --batch_size 64 --learning_rate 0.001 --num_epochs 25

# 1. Get pretrained weights for ViT-Base
pretrained_vit_weights = torchvision.models.ViT_B_16_Weights.DEFAULT

# 2. Setup a ViT model instance with pretrained weights
pretrained_vit = torchvision.models.vit_b_16(weights=pretrained_vit_weights).to(device)

# 3. Freeze the base parameters
for parameter in pretrained_vit.parameters():
    parameter.requires_grad = False

# 4. Change the classifier head
class_names = ['Bad_tire','Good_tire']

set_seeds()
pretrained_vit.heads = nn.Linear(in_features=768, out_features=len(class_names)).to(device)
# pretrained_vit # uncomment for model output

from torchinfo import summary

# Print a summary using torchinfo (uncomment for actual output)
summary(model=pretrained_vit,
        input_size=(32, 3, 224, 224), # (batch_size, color_channels, height, width)
        #col_names=["input_size"], # uncomment for smaller output
        col_names=["input_size", "output_size", "num_params", "trainable"],
        col_width=20,
        row_settings=["var_names"]
)

# Setup directory paths to train and test images
train_dir = '/content/drive/MyDrive/Test/test'
test_dir = '/content/drive/MyDrive/Train/train'

# Get automatic transforms from pretrained ViT weights
pretrained_vit_transforms = pretrained_vit_weights.transforms()
print(pretrained_vit_transforms)

import os

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
):

  # 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

# Setup dataloaders
train_dataloader_pretrained, test_dataloader_pretrained, class_names = create_dataloaders(
                                                                                            train_dir=train_dir,
                                                                                            test_dir=test_dir,
                                                                                            transform=pretrained_vit_transforms,
                                                                                            batch_size=32) # Could increase if we had more samples, such as here: https://arxiv.org/abs/2205.01580 (there are other improvements there too...)

#import data_setup.py

!python train.py --batch_size 64 --learning_rate 0.001 --num_epochs 25

import engine

# Create optimizer and loss function
optimizer = torch.optim.Adam(params=pretrained_vit.parameters(),
                             lr=1e-3)
loss_fn = torch.nn.CrossEntropyLoss()

# Train the classifier head of the pretrained ViT feature extractor model
set_seeds()
pretrained_vit_results = engine.train(model=pretrained_vit,
                                      train_dataloader=train_dataloader_pretrained,
                                      test_dataloader=test_dataloader_pretrained,
                                      optimizer=optimizer,
                                      loss_fn=loss_fn,
                                      epochs=10,
                                      device=device)

# Commented out IPython magic to ensure Python compatibility.
# %%writefile helper_functions.py
# 
# # helper_functions.py
# 
# """
# 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
# import os
# 
# 
# 
# # 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)
#

# Plot the loss curves
from helper_functions import plot_loss_curves

plot_loss_curves(pretrained_vit_results)

import requests

# Import function to make predictions on images and plot them
from predict import pred_and_plot_image

# Setup custom image path
custom_image_path = "/content/drive/MyDrive/validation/Bad_Tire (3).jpg"

# Predict on custom image
pred_and_plot_image(model=pretrained_vit,
                    image_path=custom_image_path,
                    class_names=class_names)

# Import function to make predictions on images and plot them
from predict import pred_and_plot_image

# Setup custom image path
custom_image_path = "/content/drive/MyDrive/validation/Good_Tire (4).jpg"

# Predict on custom image
pred_and_plot_image(model=pretrained_vit,
                    image_path=custom_image_path,
                    class_names=class_names)