utils.py

import matplotlib.pyplot as plt
import numpy as np
from skimage.transform import pyramid_expand
from torch import backends, cat, cuda
from torch import device as torch_device
from torch import manual_seed
from torch.nn.utils.rnn import pad_sequence

# Constants
NORMALISE = True
MODEL_PARAMS = {
    "vgg16": {
        "encoder_channels": 512,
        "features_dims": 9
    },
    "resnet50": {
        "encoder_channels": 2048,
        "features_dims": 10
    },
    "inception_v3": {
        "encoder_channels": 2048,
        "features_dims": 10
    }
}

# Expectation and Standard Deviation over ImageNet
MAGIC_MU = [0.485, 0.456, 0.406]
MAGIC_SIGMA = [0.229, 0.224, 0.225]

# Set the seed for reproducibility
backends.cudnn.determinstic = True
if cuda.is_available():
    DEVICE = torch_device('cuda:0')
    cuda.manual_seed_all(42)
else:
    DEVICE = torch_device('cpu')
    manual_seed(42)


def collate(batch, pad_idx):
    """
    Form batches of data.

    :param batch: list of tuples (image, caption, image_name)
    :param pad_idx: index of padding token
    """
    images = [item[0].unsqueeze(0) for item in batch]
    images = cat(images, dim=0)

    captions = [item[1] for item in batch]
    captions = pad_sequence(captions, batch_first=True, padding_value=pad_idx)

    img_names = [item[2] for item in batch]
    return images, captions, img_names


def show_image(img, normalise, title=None):
    """
    Unormalise and show image.

    :param img: image
    :param normalise: whether to Un-normalise or not
    :param title: title of the image
    """

    img2 = np.copy(img)  # copy image to prevent changing the original one

    # Unnormalise
    if normalise:
        for i in range(3):
            img2[i] *= MAGIC_SIGMA[i]
            img2[i] += MAGIC_MU[i]

    img2 = img2.transpose((1, 2, 0))

    if title is not None:
        plt.title(title)
    plt.imshow(img2)
    plt.show()


def plot_attention(img, caption, alphas, normalise=False):
    """
    Plot the attention weights.

    :param img: image
    :param caption: caption
    :param alphas: attention weights
    :param normalise: whether to Un-normalise or not
    :param features_dims: number of features dimensions
    """
    # Unnormalise
    if normalise:
        for i in range(3):
            img[i] *= MAGIC_SIGMA[i]
            img[i] += MAGIC_MU[i]

    img = img.numpy().transpose((1, 2, 0))
    img_cpy = img

    fig = plt.figure(figsize=(15, 15))

    len_caption = len(caption)
    features_dims = np.sqrt(alphas[0].shape[1]).astype(int)

    for i in range(len_caption):
        att = alphas[i].reshape(features_dims, features_dims)
        att = pyramid_expand(att, upscale=24, sigma=8)

        ax = fig.add_subplot(len_caption // 2, len_caption // 2, i + 1)
        ax.set_title(caption[i])
        img = ax.imshow(img_cpy)
        ax.imshow(att, cmap='gray', alpha=0.7, extent=img.get_extent())

    plt.tight_layout()
    plt.show()

def plot_history(history):
    plt.figure(figsize=(10, 7))
    
    plt.plot(history)
    plt.xlabel("iterations")
    plt.ylabel("loss")
    
    plt.show()