Added PyTorch Lightning utilities and model pipeline

.gitignore

@@ -0,0 +1,174 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Misc
+*.DS_Store
+*/.DS_Store
+extras/data/*
+
+# Data files
+data/10_whole_foods*
+data/FashionMNIST
+data/10_whole_foods*
+data/food*
+data/pizza_steak_sushi/*
+data/pizza_steak_sushi_20_percent/
+10_whole_foods.zip
+going_modular/data/
+data/cifar-*
+logs/
+runs/
+lightning_logs/
+extras/cifar-10*
+
+# Notebooks
+09_pytorch_model_deployment-*
+
+# Models
+models/03_pytorch_computer_vision_model_2.pth
+models/04_pytorch_custom_datasets_tinyvgg.pth
+models/07_effnetb0_data_10_percent_10_epochs.pth
+models/07_effnetb0_data_10_percent_5_epochs.pth
+models/07_effnetb0_data_20_percent_10_epochs.pth
+models/07_effnetb0_data_20_percent_5_epochs.pth
+models/07_effnetb2_data_10_percent_10_epochs.pth
+models/07_effnetb2_data_10_percent_5_epochs.pth
+models/07_effnetb2_data_20_percent_5_epochs.pth
+models/08_*
+models/09_*
+
+# Demos
+demos/foodvision_big/
+demos/foodvision_mini/
+flagged/
+
+# Docs
+.cache
+mkdocs-material-insiders

lightning_utils.py

@@ -0,0 +1,179 @@
+'''
+Collection of boilerplate and utility functions for PyTorch Lightning.
+'''
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, random_split
+import torchvision as tv
+# from torch.optim.optimizer import ParamsT # Could use instead of nn.Module as optimiser_factory argument
+#     # I.e. ParamsT: TypeAlias = Union[Iterable[torch.Tensor], Iterable[Dict[str, Any]]]
+
+import pytorch_lightning as L
+
+import os
+from pathlib import Path
+
+from typing import Callable
+
+
+class Strike(L.LightningModule):
+    '''As in "Lightning Strike", to make a PyTorch Module a LightningModule'''
+    def __init__(self, model: nn.Module,
+                 loss_fn: Callable[[torch.Tensor], torch.Tensor], metric_name_and_fn: tuple[str, Callable[[torch.Tensor, torch.Tensor], torch.tensor]],
+                 optimiser_factory: Callable[[nn.Module], torch.optim.Optimizer],
+                 prediction_fn: Callable[[torch.Tensor], torch.Tensor],
+                 learning_rate = 0.001, log_at_every_step = False):
+        '''Class for turning a nn.Module into a LightningModule (a lightning strike of sorts).
+        The optimiser_factory argument is a callable taking in the module, from which it extracts .parameters() and .learning_rate to produce an optimiser.'''
+        super().__init__()
+
+        self.model = model
+            # If the model form were known then its layers could be moved to this object's level rather than a nested one (not necessary but neater)
+            # The procedural versions pf this are not useful since a nested nn.Sequential still exists, i.e. any of
+            #   self.model = nn.Sequential(target._modules) # Preserves layer names
+            #   self.model = nn.Sequential(*source.children()) # *source.modules() would return the larger container as well
+        
+        self.loss_fn = loss_fn
+        self.metric_name, self.metric_fn = metric_name_and_fn
+        self.optimiser_factory = optimiser_factory
+        self.prediction_fn = prediction_fn
+
+        self.learning_rate = learning_rate
+        self.log_at_every_step = log_at_every_step
+        self.train_step_outputs, self.validation_step_outputs, self.test_step_outputs = dict(), dict(), dict()
+    
+    def forward(self, x):
+        return self.model(x)
+    
+    # No need to override these two hooks
+    # def backward(self, trainer, loss, optimizer, optimizer_idx):
+    #     loss.backward()
+    # def optimizer_step(self, epoch, batch_idx, optimiser, optimizer_idx):
+    #     optimiser.step()
+
+    def training_step(self, batch, batch_idx):
+        loss, metric, x_hat, y = self._common_step(batch, batch_idx)
+        self.train_step_outputs = dict(prefix = 'train', loss = loss, metric = metric)
+        return loss
+    
+    def on_train_epoch_end(self):
+        self._common_epoch_end_step(self.train_step_outputs)
+    
+    def validation_step(self, batch, batch_idx):
+        loss, metric, x_hat, y = self._common_step(batch, batch_idx)
+        self.validation_step_outputs = dict(prefix = 'val', loss = loss, metric = metric)
+        return loss
+    
+    def on_validation_epoch_end(self):
+        self._common_epoch_end_step(self.validation_step_outputs)
+
+    def test_step(self, batch, batch_idx):
+        loss, metric, x_hat, y = self._common_step(batch, batch_idx)
+        self.test_step_outputs = dict(prefix = 'test', loss = loss, metric = metric)
+        return loss
+    
+    def on_test_epoch_end(self):
+        self._common_epoch_end_step(self.test_step_outputs)
+
+    def _common_step(self, batch, batch_idx):
+        x, y = batch 
+        x_hat = self.forward(x)
+        loss = self.loss_fn(x_hat, y)
+        metric = self.metric_fn(x_hat, y)
+        return loss, metric, x_hat, y
+    
+    def _common_epoch_end_step(self, outputs):
+        self.log_dict({f'{outputs["prefix"]}_loss': outputs['loss'], f'{outputs["prefix"]}_{self.metric_name}': outputs['metric']}, prog_bar = True, on_step = self.log_at_every_step, on_epoch = True)
+        outputs.clear() # Freeing memory is suggested in the docs, though it is trivial in this class
+
+    def predict_step(self, batch, batch_idx):
+        x, y = batch 
+        x_hat = self.forward(x)
+        preds = self.prediction_fn(x_hat)
+        return preds
+
+    def configure_optimizers(self):
+        return self.optimiser_factory(self)
+
+
+
+class IteratedLearningRateFinder(L.callbacks.LearningRateFinder):
+    def __init__(self, at_epochs: list[int], *args, **kwargs):
+        '''CURRENTLY FAILS AT THE 2ND OCCURRENCE (despite being directly from the docs: https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.LearningRateFinder.html)
+        The lr finding tuns at epoch 0 regardless of whether 0 is in at_epochs.
+        E.g. for periodic lr adjustments pass [e for e in range(epochs) if e % period == 0]'''
+        super().__init__(*args, **kwargs)
+        self.at_epochs = at_epochs
+
+    def on_fit_start(self, *args, **kwargs):
+        return
+
+    def on_train_epoch_start(self, trainer, pl_module):
+        if trainer.current_epoch in self.at_epochs or trainer.current_epoch == 0:
+            self.lr_find(trainer, pl_module)
+
+
+
+class LocalImageDataModule(L.LightningDataModule):
+    def __init__(self, folders: str | Path | dict[str, str | Path], transform: tv.transforms.Compose,
+                 batch_size: int, num_workers: int = os.cpu_count(), split: tuple[float, float, float] = (0.7, 0.2, 0.1)):
+        super().__init__()
+        '''Return a LightningDataModule for a local image folder (or folders) for classification purposes.
+        Images are expected to be in subfolders named by their classes.
+        In the str or Path folders cases, the folder content is checked for subfolders called train, test and valid (yes, in this order for consistency), and if any is present they are treated as the list input,
+        however, if none is present, then the split argument is required, containing a tuple of proportions to allocate to training, validation and testing datasets.
+        In the dict folders case the keys are expected to be in ['train', 'valid', 'test'].
+        The class names are from the first folder and assumed to be consistent across the others.
+        '''
+
+        ########### Could relax requirement to train and test and then produce a validate dataset from the training one #########
+
+        self.prefixes = ['train', 'valid', 'test']
+
+        data_path = None
+        if isinstance(folders, (str, Path)):
+            data_path = Path(folders)
+            folders = {sub: full_sub for sub in self.prefixes if (full_sub := data_path / sub).is_dir()}
+        elif not isinstance(folders, dict): raise ValueError('Please provide a folders argument of types str | Path | dict[str, str | Path].')
+
+        assert set(folders.keys()).issubset(self.prefixes), f'Exactly the {self.prefixes} folders are required; {folders.keys()} were provided.'
+        if len(folders) == 3: folders = folders
+        elif len(folders) == 0 and data_path is not None:
+            assert sum(split) == 1
+            folders = (data_path, dict(zip(self.prefixes, split)))
+        else: raise ValueError(f'All of {self.prefixes} subfolders are required for the single-folder folders argument; only {folders.keys()} were provided.')
+
+        self.folders = folders
+        self.transform = transform
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+
+        self.train_ds, self.val_ds, self.test_ds = None, None, None
+        self.classes = None
+
+    # def prepare_data(self):
+    #     '''Not currently implemented. Mostly meant for downloading data.'''
+    #     pass
+
+    def setup(self, stage):
+        if isinstance(self.folders, tuple):
+            all_data = tv.datasets.ImageFolder(self.folders[0], transform = self.transform)
+            self.classes = all_data.classes
+            self.train_ds, self.val_ds, self.test_ds = random_split(all_data, self.folders[1])
+        else:
+            if stage == 'fit':
+                self.train_ds, self.val_ds = [tv.datasets.ImageFolder(self.folders[k], transform = self.transform) for k in self.prefixes[:-1]]
+                self.classes = self.train_ds.classes
+            if stage == 'test':
+                self.test_ds = tv.datasets.ImageFolder(self.folders[self.prefixes[-1]], transform = self.transform)
+
+    def train_dataloader(self):
+        return DataLoader(self.train_ds, batch_size = self.batch_size, shuffle = True, num_workers = self.num_workers, pin_memory = True, persistent_workers = True)
+
+    def val_dataloader(self):
+        return DataLoader(self.val_ds, batch_size = self.batch_size, shuffle = False, num_workers = self.num_workers, pin_memory = True, persistent_workers = True)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_ds, batch_size = self.batch_size, shuffle = False, num_workers = self.num_workers, pin_memory = True, persistent_workers = True)
+
+

pytorch_utils.py

@@ -7,7 +7,6 @@ from torch import nn
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
 
-# Here just to be exported
 from torchinfo import summary
 import torchmetrics
 
@@ -41,7 +40,7 @@ def set_seeds(seed: int = 42):
 def train_combinations(combinations: dict[str, tuple[str, str, str, int, str]],
                        model_factories: dict[str, Callable[[], nn.Module]], train_dataloaders: dict[str, DataLoader],
                        optimiser_factories: dict[str, Callable[[nn.Module], torch.optim.Optimizer]],
-                       test_dataloader: DataLoader, loss_fn: nn.Module, metric_name_and_fn: tuple[str, Callable[[torch.tensor, torch.tensor], torch.tensor]],
+                       test_dataloader: DataLoader, loss_fn: nn.Module, metric_name_and_fn: tuple[str, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]],
                        reset_seed: int = 42, device: torch.device = 'cuda' if torch.cuda.is_available() else 'cpu', show_progress_bar = True):
     '''Run a series of modelling tasks by defining combinations of models, dataloaders, optimisers and epochs, as well as an optional previously-fit combination
     to start from (e.g. for a combination which is the same as a previous one but with more epochs or different training data).
@@ -96,7 +95,7 @@ def train_combinations(combinations: dict[str, tuple[str, str, str, int, str]],
 
 
 def fit(model: nn.Module, train_dataloader: DataLoader, test_dataloader: DataLoader, 
-        optimiser: torch.optim.Optimizer, loss_fn: nn.Module, metric_name_and_fn: tuple[str, Callable[[torch.tensor, torch.tensor], torch.tensor]],
+        optimiser: torch.optim.Optimizer, loss_fn: nn.Module, metric_name_and_fn: tuple[str, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]],
         epochs: int, writer: torch.utils.tensorboard.writer.SummaryWriter,
         device: torch.device = 'cuda' if torch.cuda.is_available() else 'cpu', show_progress_bar = True, model_name: str = None) -> dict[str, list]:
     '''Trains and tests a PyTorch model.
@@ -148,7 +147,7 @@ def fit(model: nn.Module, train_dataloader: DataLoader, test_dataloader: DataLoa
 
 
 def training_step(model: nn.Module, dataloader: DataLoader,
-                  loss_fn: nn.Module, metric_fn: Callable[[torch.tensor, torch.tensor], torch.tensor], optimiser: torch.optim.Optimizer,
+                  loss_fn: nn.Module, metric_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], optimiser: torch.optim.Optimizer,
                   device: torch.device = 'cuda' if torch.cuda.is_available() else 'cpu', show_progress_bar = True, epoch: int = None) -> tuple[float, float]:
     '''Trains a PyTorch model for a single epoch.
 
@@ -189,7 +188,7 @@ def training_step(model: nn.Module, dataloader: DataLoader,
 
 
 def testing_step(model: nn.Module, dataloader: DataLoader, 
-                 loss_fn: nn.Module, metric_fn: Callable[[torch.tensor, torch.tensor], torch.tensor],
+                 loss_fn: nn.Module, metric_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
                  device: torch.device = 'cuda' if torch.cuda.is_available() else 'cpu', show_progress_bar = True, epoch: int = None) -> tuple[float, float]:
     '''Tests a PyTorch model for a single epoch.
 
@@ -302,6 +301,14 @@ def download_unzip(source: str, destination: str, remove_source: bool = True) ->
 
 
 
+#### Info Functions ####
+
+def summ(model: nn.Module, input_size: tuple):
+    '''Shorthand for typical summary specification'''
+    return summary(model = model, input_size = (32, 3, 224, 224),
+                   col_names = ['input_size', 'output_size', 'num_params', 'trainable'], col_width = 20, row_settings = ['var_names'])
+
+
 
 #### Plotting Functions ####
 

pytorch_vision_utils.py

@@ -0,0 +1,168 @@
+'''Torchvision and related utility functions'''
+
+import torch
+import torchvision as tv
+from torch.utils.data import DataLoader
+import timm # Here just to be exported
+
+import numpy as np
+import pandas as pd
+from tqdm.auto import tqdm
+import base64
+import altair as alt
+import matplotlib.pyplot as plt # REMOVE IN FAVOUR OF ALTAIR
+
+import os
+import io
+from pathlib import Path
+from PIL import Image
+# from itertools import batched # in Python>=3.12
+
+
+def image_dataloaders(folders: str | Path | list[str | Path], transform: tv.transforms.Compose, batch_size: int, num_workers: int = os.cpu_count()) -> tuple[list[DataLoader], list[str]]:
+    '''Return PyTorch DataLoaders and class names for the given folder or list of folders (with expected subfolders named by class).
+    In the non-list folders case, the folder content is checked for subfolders called train, test and valid (yes, in this order for consistency), and if any is present they are treated as the list input.
+    The first folder is assumed to be the training data and will therefore produce a shuffling dataloader, while the others will not.
+    The class names are from the first folder and assumed to be consistent across the others.
+    '''
+    if isinstance(folders, (str, Path)):
+        data_path = Path(folders)
+        folders = subfolders if (subfolders := [full_sub for sub in ['train', 'valid', 'test'] if (full_sub := data_path / sub).is_dir()]) else [folders]
+
+    datasets = [tv.datasets.ImageFolder(folder, transform = transform) for folder in folders]
+    dataloaders = [DataLoader(ds, batch_size = batch_size, shuffle = i == 0, num_workers = num_workers, pin_memory = True, persistent_workers = True) for i, ds in enumerate(datasets)]
+
+    return dataloaders, datasets[0].classes
+
+
+def plot_img_preds(model: torch.nn.Module, image_path: str, class_names: list[str], transform: tv.transforms, device: torch.device = 'cuda' if torch.cuda.is_available() else 'cpu'):
+    '''Plot one image with its prediction and probability as the title.
+    '''
+    img = Image.open(image_path)
+
+    model.to(device)
+    model.eval()
+    with torch.inference_mode(): pred_logit = model(transform(img).unsqueeze(dim = 0).to(device)) # Prepend "batch" dimension (-> [batch_size, color_channels, height, width])
+    pred_prob = torch.softmax(pred_logit, dim = 1)
+    pred_label = torch.argmax(pred_prob, dim = 1)
+
+    plt.figure()
+    plt.imshow(img)
+    plt.title(f"Pred: {class_names[pred_label]} | Prob: {pred_prob.max():.3f}")
+    plt.axis(False)
+
+    # Change text colour based on correctness?
+
+
+def record_image_preds(image_paths: str | list[str], model: torch.nn.Module, transform: tv.transforms.Compose, class_names: list[str],
+                       sort_by_correctness = True, device: torch.device = 'cuda' if torch.cuda.is_available() else 'cpu'):
+    '''Generate a dataframe of paths, true classes, (single) predicted classes and their confidence.
+    Column names: path, true_class, pred_class, pred_prob, correct.
+    If sort_by_correctness, then the dataframe is sorted by increasing correctness and confidence, i.e. first by prediction correctness and then by its probability,
+    with wrong predictions first, and both wrong and right by decreasing confidence.
+    If a single string is given as image_paths, then all */*.jpg and */*.png matches from it are used instead.
+    '''
+    true_classes, pred_classes, pred_probs, correctness, image_data = [], [], [], [], []
+
+    if isinstance(image_paths, str): image_paths = list(Path(image_paths).glob('*/*.jpg')) + list(Path(image_paths).glob('*/*.png'))
+
+    for path in tqdm(image_paths):
+        img = Image.open(path)
+
+        model.eval()
+        with torch.inference_mode(): pred_logit = model(transform(img).unsqueeze(0).to(device)) # Prepend "batch" dimension (-> [batch_size, color_channels, height, width])
+        pred_prob = torch.softmax(pred_logit, dim = 1)
+        pred_label = torch.argmax(pred_prob, dim = 1)
+
+        true_classes.append(class_name := path.parent.stem)
+        pred_classes.append(pred_class := class_names[pred_label.cpu()])
+        pred_probs.append(pred_prob.unsqueeze(0).max().cpu().item())
+        correctness.append(class_name == pred_class)
+
+
+    res = pd.DataFrame(dict(path = [str(p) for p in image_paths], true_class = true_classes, pred_class = pred_classes, pred_prob = pred_probs, correct = correctness))
+    return res.sort_values(by = ['correct', 'pred_prob'], ascending = [True, False]) if sort_by_correctness else res
+
+
+def base64_image_formatter(image_or_path: Image.Image | str) -> str:
+    '''Generate a base64-encoded string representation of the given image (or path).
+    Example usecase: a dataframe meant for Altair contains PIL images (or their paths) in a column, in which case pass this temporary dataframe to the alt.Chart:
+        `df.assign(image = df.image_or_path_column.apply(base64_image_formatter))`
+    '''
+    if isinstance(image_or_path, str): image_or_path = Image.open(image_or_path)
+    with io.BytesIO() as buffer: # Docs: https://altair-viz.github.io/user_guide/marks/image.html#use-local-images-as-image-marks
+        image_or_path.save(buffer, format = 'PNG')
+        data = base64.b64encode(buffer.getvalue()).decode('utf-8')
+        return f'data:image/png;base64,{data}'
+
+
+def image_pred_grid(image_df: pd.DataFrame, ncols = 4, img_width = 200, img_height = 200, allow_1_col_reduction = True):
+    '''Create an Altair plot displaying a grid of images and their predicted classes, highlighting incorrect predictions.
+    image_df is expected to have the columns: path, true_class, pred_class, pred_prob, correct.
+    If allow_1_col_reduction and the last row (by the given ncols) is at least half empty and using ncols-1 would not increase rows, then ncols-1 is used instead.
+    '''
+    # Docs: https://altair-viz.github.io/user_guide/compound_charts.html
+    # Opened issue on making it easier through alt.Facet: https://github.com/altair-viz/altair/issues/3398
+
+    ncols = min(ncols, len(image_df))
+    nrows = 1 + len(image_df) // ncols
+    # If the last row is at least half empty and could reduce columns without increasing rows, do so
+    if allow_1_col_reduction and nrows > 1 and len(image_df) % ncols <= ncols / 2 and 1 + len(image_df) // (ncols - 1) == nrows: ncols -= 1
+
+    expanded_df = image_df.assign(
+        image = image_df.path.apply(base64_image_formatter),
+        title = image_df.pred_class + ' - ' + image_df.pred_prob.map(lambda p: f'{p:.2f}'),
+        index = image_df.index
+    )
+
+    base = alt.Chart(expanded_df).mark_image(width = img_width, height = img_height).encode(url = 'image:N')
+    chart = alt.vconcat()
+    for row_indices in (expanded_df.index[i:i + ncols] for i in range(0, len(expanded_df), ncols)): # itertools.batched(expanded_df.index, ncols) in Python>=3.12
+        row_chart = alt.hconcat()
+        for index in row_indices:
+            row_chart |= base.transform_filter(alt.datum.index == index).properties(
+                title = alt.Title(expanded_df.title[index], fontSize = 17, color = 'green' if expanded_df.correct[index] else 'red'))
+        chart &= row_chart
+
+    ## Version with no subplots (but no titles)
+    # chart = alt.Chart(image_df.assign( # vv cannot trust the df index since it might not be ordered
+    #     row = np.arange(len(image_df)) // ncols, col = np.arange(len(image_df)) % ncols # Could use the transform_compose block for this, but no // in the alt.expr language
+    # )).mark_image(width = img_width, height = img_height).encode(
+    #     alt.X('col:O', title = None, axis = None), alt.Y('row:O', title = None, axis = None), url = 'image:N'
+    # ).properties(
+    #     width = img_width * 1.1 * ncols, height = img_height * 1.1 * nrows
+    # )
+
+    ## Version with faceting (but not coloured titles (no titles in fact, but non-coloured headers))
+    # chart = alt.Chart(image_df.assign(
+    #     image = image_df.path.apply(base64_image_formatter),
+    #     title = image_df.pred_class + ' - ' + image_df.pred_prob.map(lambda p: f'{p:.2f}')
+    # )).mark_image(width = img_width, height = img_height).encode(url = 'image:N'
+    # ).facet( # Header fields: https://altair-viz.github.io/user_guide/generated/core/altair.Header.html
+    #     alt.Facet('title:N', header = alt.Header(labelFontSize = 17, labelColor = 'red')).title('Prediction and Confidence'), columns = ncols, title = 'Hi'
+    # )
+
+    return chart
+
+
+
+
+
+# import torchvision
+# import matplotlib.pyplot as plt
+# # Plot the top 5 most wrong images
+# for row in top_5_most_wrong.iterrows():
+#   row = row[1]
+#   image_path = row[0]
+#   true_label = row[1]
+#   pred_prob = row[2]
+#   pred_class = row[3]
+#   # Plot the image and various details
+#   img = torchvision.io.read_image(str(image_path)) # get image as tensor
+#   plt.figure()
+#   plt.imshow(img.permute(1, 2, 0)) # matplotlib likes images in [height, width, color_channels]
+#   plt.title(f"True: {true_label} | Pred: {pred_class} | Prob: {pred_prob:.3f}")
+#   plt.axis(False);
+
+
+