latentspacevis.py

import itertools
import os
import random

import lightning as L
import matplotlib.ticker as ticker
import numpy as np
import pandas as pd
import seaborn as sns
import torch
from matplotlib import pyplot as plt

from sisiae import AutoencoderNetwork
from utils.data import get_batch, load_data
from utils.names import naming, rename_algs

random.seed(0)
torch.manual_seed(0)
np.random.seed(0)

fabric = L.Fabric(accelerator="cpu")
fabric.launch()


ax_labels = ["latent dimension 1", "latent dimension 2"]
ax_labels_fontsize = 8


def ax_ticks_formatter(x, pos):
    x *= 100
    x = int(x) if int(x) == x else float(x)
    x /= 100
    return x


if fabric.global_rank == 0:
    model_path = "./models/"
    models = ["Eigenimages", "SISSIAE-v22-lin", "SISSIAE-v23-arctan"]

    datasets = {
        "PK4Particles": 11 * 11,
        "MNIST": 28 * 28,
        "FashionMNIST": 28 * 28,
    }

    dfs = []
    for test_case in models:
        if os.path.exists(f"test_results/{test_case}.txt"):
            with open(f"test_results/{test_case}.txt") as f:
                results = f.read()
            results = results.split("\n")
            results = [line.split(", ") for line in results]
            if test_case == "Eigenimages":
                head = ["date", "algorithm", "image data"]
            else:
                head = ["model file", "date", "algorithm", "image data"]
            for col in results[0][len(head) :]:
                head.append(col.split(": ")[0])
            df = pd.DataFrame(results, columns=head)
            df = df.dropna(axis=0)
            for h in head:
                df[h] = df[h].map(lambda x: x.lstrip(f"{h}: "))
            dfs.append(df)

    dfs.append(df)
    df = pd.concat(dfs, ignore_index=True)
    df["LQ"] = df["LQ"].fillna("None")
    df["MI"] = df["MI"].fillna("None")

    df = df.drop(df[(df["MI"] == "data") | (df["MI"] == "pca")].index)

    df["MSE"] = df["MSE"].astype(float)
    df["PSNR"] = df["PSNR"].astype(float)
    df["SSIM"] = df["SSIM"].astype(float)
    df["C"] = df["C"].astype(int)

    naming = {
        "MSE": "mean squared error",
        "PSNR": "peak signal to noise ratio",
        "SSIM": "structural similarity measure",
        "CR": "compression ratio",
        "LQ": "latent quantization",
        "C": "latent space dimensionality",
    }

    df = df.rename(columns=naming)

    for image_data in datasets.keys():
        train_dataset, test_dataset, _ = load_data(image_data)

        # Initialize Eigenimages algorithm for comparison
        data_loader = torch.utils.data.DataLoader(
            dataset=train_dataset,
            batch_size=min(len(train_dataset), 10000 if image_data == "Omniglot" else 60000),
            shuffle=True,
        )

        data_iter = iter(data_loader)
        images = get_batch(data_iter, fabric)

        mean_image = torch.mean(images, dim=0)
        caricatures = images - mean_image
        (_, eigenvalues, eigenimages) = torch.pca_lowrank(
            caricatures.view(caricatures.size(0), -1),
            q=caricatures.view(caricatures.size(0), -1).shape[1],
        )
        eigenimages = eigenimages.T.view(1, *eigenimages.T.shape)

        # Data Loader (Input Pipeline)
        data_loader = torch.utils.data.DataLoader(
            dataset=test_dataset, batch_size=len(test_dataset), shuffle=True
        )

        data_loader = fabric.setup_dataloaders(data_loader)

        data_iter = iter(data_loader)
        images, labels = get_batch(data_iter, fabric, labels=True)
        if labels is None:
            labels = np.zeros(images.shape[0])
        else:
            labels = labels.detach().numpy()

        _df = df[df["image data"] == image_data]
        _df = _df[_df[naming["C"]] == 2]

        # num_models = _df['algorithm'].unique().shape[0]
        num_models = 5

        fig, axes = plt.subplots(1, num_models, figsize=(10, 2.3), tight_layout=True)
        # fig.suptitle(f'Data Set: {image_data}', y=0.93, fontsize=14)

        # Eigenimages for comparison
        mean_image = torch.mean(images, dim=0)
        caricatures = images - mean_image

        components = 2
        latent_space = torch.inner(
            caricatures.view(caricatures.size(0), -1), eigenimages[:, :components, :]
        )

        eigen_norm = torch.norm(eigenimages[:, :components, :], p=1, dim=-1).view(
            1, eigenimages[:, :components, :].size(1)
        )
        # print(eigen_norm.shape)
        latent_space = latent_space / eigen_norm

        data = np.concatenate([latent_space.squeeze().T, np.expand_dims(labels, axis=1).T]).T
        df_latentspace = pd.DataFrame(data=data, columns=[ax_labels[0], ax_labels[1], "label"])
        df_latentspace["label"] = df_latentspace["label"].astype(int)
        # print(df_latentspace)

        index = 0
        axes[index].set_title("Eigenimages", fontsize=10)
        g = sns.scatterplot(
            ax=axes[index],
            data=df_latentspace,
            x=ax_labels[0],
            y=ax_labels[1],
            hue="label",
            size="label",
            sizes=(0.3, 0.3),
            legend=False,  # (np.unique(labels).shape[0] > 1),
            palette="gist_rainbow",
        )
        axes[index].axis("equal")
        if index >= 1:
            g.set(ylabel=None)
        else:
            axes[index].set_ylabel(ax_labels[1], fontsize=ax_labels_fontsize)
        axes[index].set_xlabel(ax_labels[0], fontsize=ax_labels_fontsize)
        g.set_yticklabels(g.get_yticks(), size=5)
        axes[index].yaxis.set_major_formatter(ticker.FuncFormatter(ax_ticks_formatter))
        g.set_xticklabels(g.get_xticks(), size=5)
        axes[index].xaxis.set_major_formatter(ticker.FuncFormatter(ax_ticks_formatter))

        for i, (activation, costs) in enumerate(
            itertools.product(["linear", "arctan"], ["mse", "ssim"])
        ):
            index = i + 1
            print(index)
            for _, row in _df.iterrows():
                if activation in row["algorithm"] and costs in row["algorithm"]:
                    Autoencoder = AutoencoderNetwork(
                        images,
                        row[naming["C"]],
                        mode="",
                        latent_quantization=row[naming["LQ"]],
                        activation=activation,
                        costs=costs,
                    )
                    Autoencoder = fabric.setup_module(Autoencoder)

                    path = model_path + row["model file"]
                    state_dict = torch.load(path)
                    if activation == "linear":
                        state_dict.pop("sigmoid_scaling", None)
                        state_dict.pop("sigmoid_offset", None)
                        state_dict.pop("decoder.sigmoid_scaling", None)
                        state_dict.pop("decoder.sigmoid_offset", None)
                    Autoencoder.load_state_dict(state_dict)
                    Autoencoder.eval()

                    latent_space = Autoencoder.encoder(images).detach().numpy()

                    # print(latent_space.shape, np.expand_dims(labels, axis=1).shape)

                    data = np.concatenate([latent_space.T, np.expand_dims(labels, axis=1).T]).T
                    df_latentspace = pd.DataFrame(
                        data=data, columns=[ax_labels[0], ax_labels[1], "label"]
                    )
                    df_latentspace["label"] = df_latentspace["label"].astype(int)
                    # print(df_latentspace)

                    axes[index].set_title(rename_algs[row["algorithm"]], fontsize=10)
                    g = sns.scatterplot(
                        ax=axes[index],
                        data=df_latentspace,
                        x=ax_labels[0],
                        y=ax_labels[1],
                        hue="label",
                        size="label",
                        sizes=(0.3, 0.3),
                        legend=False,  # (np.unique(labels).shape[0] > 1),
                        palette="gist_rainbow",
                    )
                    axes[index].axis("equal")
                    if index >= 1:
                        g.set(ylabel=None)
                    else:
                        axes[index].set_ylabel(ax_labels[1], fontsize=ax_labels_fontsize)
                    axes[index].set_xlabel(ax_labels[0], fontsize=ax_labels_fontsize)
                    g.set_yticklabels(g.get_yticks(), size=5)
                    axes[index].yaxis.set_major_formatter(ticker.FuncFormatter(ax_ticks_formatter))
                    g.set_xticklabels(g.get_xticks(), size=5)
                    axes[index].xaxis.set_major_formatter(ticker.FuncFormatter(ax_ticks_formatter))

        fig.tight_layout()
        plt.savefig(f"paper_images/latentspace/{image_data}.jpeg", format="jpeg", dpi=800)
        plt.clf()


torch.cuda.empty_cache()