betaVAE_train_v1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "548f5c85-34f4-4bb2-a299-982a1bbae698",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import numpy as np\n",
    "import imageio\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader, Subset\n",
    "from sklearn.preprocessing import KBinsDiscretizer\n",
    "from sklearn.metrics import r2_score\n",
    "from tqdm import tqdm\n",
    "from termcolor import colored"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f002a265-8e02-49bf-a0e5-8ff67ada0c75",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Custom dataset class for dSprites dataset\n",
    "class DspritesDataset(Dataset):\n",
    "    def __init__(self, npz_file):\n",
    "        \"\"\"\n",
    "        Initialize the dSprites dataset by loading the image data and ground truth factors.\n",
    "\n",
    "        Args:\n",
    "            npz_file (str): Path to the .npz file containing the dataset.\n",
    "        \"\"\"\n",
    "        self.data = np.load(npz_file)\n",
    "        self.images = np.expand_dims(self.data['imgs'], axis=1)\n",
    "        self.latent_classes = self.data['latents_classes'][:, 1:]  # Skip the 'color' factor (always 1)\n",
    "        self.latent_values = self.data['latents_values'][:, 1:]  # Skip the 'color' factor (always 1)\n",
    "\n",
    "    def __len__(self):\n",
    "        \"\"\"Returns the total number of images.\"\"\"\n",
    "        return len(self.images)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"\n",
    "        Fetch an image and its corresponding ground truth factors.\n",
    "\n",
    "        Args:\n",
    "            idx (int): Index of the data sample.\n",
    "\n",
    "        Returns:\n",
    "            Tuple: Image as torch.Tensor, latent classes, and latent values.\n",
    "        \"\"\"\n",
    "        image = torch.tensor(self.images[idx], dtype=torch.float32)  # Convert image to float32 tensor\n",
    "        latent_classes = torch.tensor(self.latent_classes[idx], dtype=torch.int)  # Discrete factors\n",
    "        latent_values = torch.tensor(self.latent_values[idx], dtype=torch.float32)  # Continuous factors\n",
    "\n",
    "        return image, latent_classes, latent_values\n",
    "\n",
    "# Encoder architecture\n",
    "class ConvEncoder(nn.Module):\n",
    "    def __init__(self, latent_dim):\n",
    "        \"\"\"\n",
    "        Encoder for beta-VAE using convolutional layers.\n",
    "\n",
    "        Args:\n",
    "            latent_dim (int): Dimensionality of the latent space.\n",
    "        \"\"\"\n",
    "        super(ConvEncoder, self).__init__()\n",
    "        # Convolutional layers\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, kernel_size=4, stride=2, padding=1),  # 64x64 -> 32x32\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),  # 32x32 -> 16x16\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(64, 64, kernel_size=4, stride=2, padding=1),  # 16x16 -> 8x8\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(64, 64, kernel_size=4, stride=2, padding=1),  # 8x8 -> 4x4\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        # Fully connected layers for mu and logvar\n",
    "        self.fc1 = nn.Linear(64 * 4 * 4, 256)  # Fully connected layer with 256 units\n",
    "        self.fc_mu = nn.Linear(256, latent_dim)  # Latent mean\n",
    "        self.fc_logvar = nn.Linear(256, latent_dim)  # Latent log-variance\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Forward pass for the encoder to produce latent space parameters.\n",
    "\n",
    "        Args:\n",
    "            x (torch.Tensor): Input image.\n",
    "\n",
    "        Returns:\n",
    "            mu (torch.Tensor): Mean of the latent space.\n",
    "            logvar (torch.Tensor): Log-variance of the latent space.\n",
    "        \"\"\"\n",
    "        h = self.encoder(x)\n",
    "        h = h.view(h.size(0), -1)  # Flatten the output of the convolution layers\n",
    "        h = self.fc1(h)  # Apply fully connected layer\n",
    "        mu = self.fc_mu(h)\n",
    "        logvar = self.fc_logvar(h)\n",
    "        return mu, logvar\n",
    "\n",
    "# Decoder architecture\n",
    "class ConvDecoder(nn.Module):\n",
    "    def __init__(self, latent_dim, num_channels=1):\n",
    "        \"\"\"\n",
    "        Decoder for beta-VAE using transpose convolutional layers.\n",
    "\n",
    "        Args:\n",
    "            latent_dim (int): Dimensionality of the latent space.\n",
    "            num_channels (int): Number of output image channels (1 for grayscale images).\n",
    "        \"\"\"\n",
    "        super(ConvDecoder, self).__init__()\n",
    "        \n",
    "        # Fully connected layers\n",
    "        self.fc1 = nn.Linear(latent_dim, 256)  # Latent dim -> 256\n",
    "        self.fc2 = nn.Linear(256, 64 * 4 * 4)  # 256 -> 4x4x64\n",
    "        \n",
    "        # Transpose convolutional layers with adjusted padding\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1),  # 4x4 -> 8x8\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),  # 8x8 -> 16x16\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(32, 32, kernel_size=4, stride=2, padding=1),  # 16x16 -> 32x32\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(32, num_channels, kernel_size=4, stride=2, padding=1),  # 32x32 -> 64x64\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "    \n",
    "    def forward(self, z):\n",
    "        \"\"\"\n",
    "        Forward pass for the decoder to reconstruct the image from the latent space.\n",
    "\n",
    "        Args:\n",
    "            z (torch.Tensor): Latent vector.\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: Reconstructed image.\n",
    "        \"\"\"\n",
    "        h = self.fc1(z)\n",
    "        h = self.fc2(h)\n",
    "        h = h.view(h.size(0), 64, 4, 4)  # Reshape to match convolutional layers\n",
    "        return self.decoder(h)\n",
    "\n",
    "# beta-VAE model\n",
    "class BetaVAE(nn.Module):\n",
    "    def __init__(self, latent_dim):\n",
    "        \"\"\"\n",
    "        beta-VAE model with encoder and decoder networks.\n",
    "\n",
    "        Args:\n",
    "            latent_dim (int): Dimensionality of the latent space.\n",
    "        \"\"\"\n",
    "        super(BetaVAE, self).__init__()\n",
    "        self.encoder = ConvEncoder(latent_dim)\n",
    "        self.decoder = ConvDecoder(latent_dim)\n",
    "\n",
    "    def reparameterize(self, mu, logvar):\n",
    "        \"\"\"\n",
    "        Reparameterization trick to sample from a Gaussian distribution.\n",
    "\n",
    "        Args:\n",
    "            mu (torch.Tensor): Mean of the latent space.\n",
    "            logvar (torch.Tensor): Log-variance of the latent space.\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: Reparameterized latent vector.\n",
    "        \"\"\"\n",
    "        std = torch.exp(0.5 * logvar)\n",
    "        eps = torch.randn_like(std)  # Random normal noise\n",
    "        return mu + eps * std\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Forward pass for the entire VAE model.\n",
    "\n",
    "        Args:\n",
    "            x (torch.Tensor): Input image.\n",
    "\n",
    "        Returns:\n",
    "            recon_x (torch.Tensor): Reconstructed image.\n",
    "            mu (torch.Tensor): Mean of the latent space.\n",
    "            logvar (torch.Tensor): Log-variance of the latent space.\n",
    "        \"\"\"\n",
    "        mu, logvar = self.encoder(x)\n",
    "        z = self.reparameterize(mu, logvar)\n",
    "        recon_x = self.decoder(z)\n",
    "        return recon_x, mu, logvar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88ce69ae-b1e4-4152-849b-fd3a43058a15",
   "metadata": {},
   "outputs": [],
   "source": [
    "def beta_vae_loss(X_hat, X, Z_mu, Z_logvar, beta):\n",
    "    \"\"\"\n",
    "    Compute the total loss for beta-VAE.\n",
    "\n",
    "    Args:\n",
    "        X_hat (torch.Tensor): Reconstructed input (output of the decoder).\n",
    "        X (torch.Tensor): Original input.\n",
    "        Z_mu (torch.Tensor): Mean of the latent distribution.\n",
    "        Z_logvar (torch.Tensor): Log-variance of the latent distribution.\n",
    "        beta (float): Weight for the KL divergence term.\n",
    "\n",
    "    Returns:\n",
    "        bce_loss (torch.Tensor): Binary cross-entropy loss (summed over the batch).\n",
    "        kl_loss (torch.Tensor): KL divergence loss (summed over the batch).\n",
    "        total_loss (torch.Tensor): Total loss (bce_loss + beta * kl_loss).\n",
    "    \"\"\"\n",
    "    # Binary cross-entropy loss\n",
    "    bce_loss = F.binary_cross_entropy(X_hat, X, reduction='sum')\n",
    "\n",
    "    # KL divergence loss\n",
    "    kl_loss = -0.5 * torch.sum(1 + Z_logvar - Z_mu.pow(2) - Z_logvar.exp())\n",
    "\n",
    "    # Total loss\n",
    "    total_loss = bce_loss + beta * kl_loss\n",
    "\n",
    "    return bce_loss, kl_loss, total_loss\n",
    "\n",
    "def train_beta_vae(model, train_loader, val_loader, optimizer, loss_fn, beta, device, num_epochs):\n",
    "    \"\"\"\n",
    "    Train the beta-VAE model.\n",
    "    \"\"\"\n",
    "    # Initialize loss histories\n",
    "    train_loss_hist, val_loss_hist = [], []\n",
    "    kl_loss_hist, val_kl_loss_hist = [], []\n",
    "    bce_loss_hist, val_bce_loss_hist = [], []\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        # --- Training Step ---\n",
    "        model.train()\n",
    "        train_loss, kl_loss, bce_loss = 0, 0, 0\n",
    "\n",
    "        for batch, _, _ in tqdm(train_loader, desc=f'Epoch {epoch+1}/{num_epochs} - Training'):\n",
    "            batch = batch.to(device)\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # Forward pass\n",
    "            X_hat, Z_mu, Z_logvar = model(batch)\n",
    "\n",
    "            # Compute loss\n",
    "            bce_loss_batch, kl_loss_batch, total_loss_batch = loss_fn(X_hat, batch, Z_mu, Z_logvar, beta)\n",
    "\n",
    "            # Backward pass\n",
    "            total_loss_batch.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            # Accumulate batch losses\n",
    "            train_loss += total_loss_batch.item()\n",
    "            kl_loss += kl_loss_batch.item()\n",
    "            bce_loss += bce_loss_batch.item()\n",
    "\n",
    "        num_samples = len(train_loader.dataset)\n",
    "        train_loss_hist.append(train_loss / num_samples)\n",
    "        kl_loss_hist.append(kl_loss / num_samples)\n",
    "        bce_loss_hist.append(bce_loss / num_samples)\n",
    "\n",
    "        # --- Validation Step ---\n",
    "        model.eval()\n",
    "        val_loss, val_kl_loss, val_bce_loss = 0, 0, 0\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for batch, _, _ in tqdm(val_loader, desc=f'Epoch {epoch+1}/{num_epochs} - Validation'):\n",
    "                batch = batch.to(device)\n",
    "\n",
    "                # Forward pass\n",
    "                X_hat, Z_mu, Z_logvar = model(batch)\n",
    "\n",
    "                # Compute loss\n",
    "                bce_loss_batch, kl_loss_batch, total_loss_batch = loss_fn(X_hat, batch, Z_mu, Z_logvar, beta)\n",
    "\n",
    "                # Accumulate batch losses\n",
    "                val_loss += total_loss_batch.item()\n",
    "                val_kl_loss += kl_loss_batch.item()\n",
    "                val_bce_loss += bce_loss_batch.item()\n",
    "\n",
    "        num_samples = len(val_loader.dataset)\n",
    "        val_loss_hist.append(val_loss / num_samples)\n",
    "        val_kl_loss_hist.append(val_kl_loss / num_samples)\n",
    "        val_bce_loss_hist.append(val_bce_loss / num_samples)\n",
    "\n",
    "        # --- Logging the losses ---\n",
    "        print(f\"Epoch [{epoch+1}/{num_epochs}] Train BCE: {bce_loss_hist[-1]:.4f}, KL: {kl_loss_hist[-1]:.4f}, Total: {train_loss_hist[-1]:.4f}\")\n",
    "        print(f\"Epoch [{epoch+1}/{num_epochs}] Val BCE: {val_bce_loss_hist[-1]:.4f}, KL: {val_kl_loss_hist[-1]:.4f}, Total: {val_loss_hist[-1]:.4f}\")\n",
    "\n",
    "    return {\n",
    "        'train_loss': train_loss_hist,\n",
    "        'val_loss': val_loss_hist,\n",
    "        'kl_loss': kl_loss_hist,\n",
    "        'val_kl_loss': val_kl_loss_hist,\n",
    "        'bce_loss': bce_loss_hist,\n",
    "        'val_bce_loss': val_bce_loss_hist\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e8bc48e-27af-4e36-8177-704f288dff7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to select exactly 3 distinct shape indices\n",
    "def select_3_distinct_shape_indices(dataset, split_indices):\n",
    "    \"\"\"\n",
    "    Select one index for each distinct shape (square, ellipse, heart).\n",
    "\n",
    "    Args:\n",
    "    - dataset: The dataset containing images and latent classes.\n",
    "    - split_indices: The indices corresponding to the data split (train, validation, etc.).\n",
    "\n",
    "    Returns:\n",
    "    - distinct_shape_indices: A list containing one randomly selected index for each distinct shape.\n",
    "    \"\"\"\n",
    "    # Dictionary to store indices for each shape\n",
    "    shape_indices = {0: [], 1: [], 2: []}\n",
    "    \n",
    "    # Iterate over the split indices to collect indices for each shape\n",
    "    for idx in split_indices:\n",
    "        _, latent_classes, _ = dataset[idx]\n",
    "        shape_class = latent_classes[0].item()  # Shape class: 0 = square, 1 = ellipse, 2 = heart\n",
    "        \n",
    "        # Add the index to the corresponding shape list\n",
    "        shape_indices[shape_class].append(idx)\n",
    "    \n",
    "    # Print the number of images for each shape\n",
    "    shape_names = {0: 'square', 1: 'ellipse', 2: 'heart'}\n",
    "    for shape_class, indices in shape_indices.items():\n",
    "        print(f\"Number of {shape_names[shape_class]} images: {len(indices)}\")\n",
    "    \n",
    "    # Randomly select one index from each shape class\n",
    "    distinct_shape_indices = [random.choice(shape_indices[shape]) for shape in range(3)]\n",
    "    \n",
    "    return distinct_shape_indices\n",
    "\n",
    "def visualize_3_images_from_shapes(dataset, distinct_shape_indices):\n",
    "    \"\"\"\n",
    "    Visualize 3 images corresponding to distinct shapes (square, ellipse, heart) and save them.\n",
    "\n",
    "    Args:\n",
    "    - dataset: The dataset containing images and latent classes.\n",
    "    - distinct_shape_indices: A list of 3 indices corresponding to distinct shape images.\n",
    "    \"\"\"\n",
    "    # Ensure exactly 3 indices are provided\n",
    "    if len(distinct_shape_indices) != 3:\n",
    "        raise ValueError(\"distinct_shape_indices must contain exactly 3 indices.\")\n",
    "    \n",
    "    # Define the shape names\n",
    "    shape_names = ['square', 'ellipse', 'heart']\n",
    "    \n",
    "    # Create a figure with 3 subplots\n",
    "    fig, axes = plt.subplots(1, 3, figsize=(12, 4))\n",
    "\n",
    "    # Iterate through the indices, display, and save the images\n",
    "    for i, (idx, shape_name) in enumerate(zip(distinct_shape_indices, shape_names)):\n",
    "        image, _, _ = dataset[idx]  # Retrieve the image at the specified index\n",
    "        \n",
    "        # Display the image in the subplot\n",
    "        axes[i].imshow(image.squeeze(), cmap='gray')\n",
    "        axes[i].axis('off')\n",
    "        axes[i].set_title(f\"{shape_name}\", fontsize=16)\n",
    "\n",
    "        # Save each image with its corresponding shape name\n",
    "        image_save_path = os.path.join(SAVE_PATH, f\"{shape_name}.png\")\n",
    "        plt.imsave(image_save_path, image.squeeze(), cmap='gray')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9e708e3-e8e9-4d68-9623-f97c2cd42942",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to plot losses for both training and validation over epochs\n",
    "def plot_beta_vae_losses(train_loss_hist, val_loss_hist, kl_loss_hist, val_kl_loss_hist, recon_loss_hist, val_recon_loss_hist, beta):\n",
    "    \"\"\"\n",
    "    Plot the total loss, KL divergence, and reconstruction error for both training and validation sets.\n",
    "\n",
    "    Args:\n",
    "    - train_loss_hist (list): History of training total losses (-ELBO).\n",
    "    - val_loss_hist (list): History of validation total losses (-ELBO).\n",
    "    - kl_loss_hist (list): History of training KL divergence losses.\n",
    "    - val_kl_loss_hist (list): History of validation KL divergence losses.\n",
    "    - recon_loss_hist (list): History of training reconstruction errors.\n",
    "    - val_recon_loss_hist (list): History of validation reconstruction errors.\n",
    "    - beta (float): The beta value used for the loss scaling in beta-VAE.\n",
    "    \"\"\"\n",
    "    epochs = len(train_loss_hist)\n",
    "\n",
    "    plt.figure(figsize=(15, 5))\n",
    "\n",
    "    # Subplot for Total Loss (-ELBO)\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.plot(range(epochs), train_loss_hist, label=\"Training\", linestyle='-')\n",
    "    plt.plot(range(epochs), val_loss_hist, label=\"Validation\", linestyle='--')\n",
    "    plt.xlabel('Epochs', fontsize=14)\n",
    "    plt.ylabel('Total Loss', fontsize=14)\n",
    "    plt.yticks(fontsize=12)\n",
    "    plt.legend(loc='best', fontsize=12)\n",
    "\n",
    "    # Subplot for KL Divergence\n",
    "    plt.subplot(1, 3, 2)\n",
    "    plt.plot(range(epochs), kl_loss_hist, label=\"Training\", linestyle='-')\n",
    "    plt.plot(range(epochs), val_kl_loss_hist, label=\"Validation\", linestyle='--')\n",
    "    plt.xlabel('Epochs', fontsize=14)\n",
    "    plt.ylabel('KL Divergence', fontsize=14)\n",
    "    plt.yticks(fontsize=12)\n",
    "    plt.legend(loc='best', fontsize=12)\n",
    "\n",
    "    # Subplot for Reconstruction Error\n",
    "    plt.subplot(1, 3, 3)\n",
    "    plt.plot(range(epochs), recon_loss_hist, label=\"Training\", linestyle='-')\n",
    "    plt.plot(range(epochs), val_recon_loss_hist, label=\"Validation\", linestyle='--')\n",
    "    plt.xlabel('Epochs', fontsize=14)\n",
    "    plt.ylabel('BCE Loss', fontsize=14)\n",
    "    plt.yticks(fontsize=12)\n",
    "    plt.legend(loc='best', fontsize=12)\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.subplots_adjust(left=0.1, right=0.95, top=0.95, bottom=0.1, wspace=0.3)\n",
    "    plt.savefig(f\"{SAVE_PATH}beta_vae_loss_plots_{beta}.png\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ccf43b6-e41a-4ea8-b741-15859fdf0d30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualization of original and reconstructed images in separate grids\n",
    "def visualize_images_grid(model, data_loader, device, beta, num_images=64):\n",
    "    \"\"\"\n",
    "    Visualize and save original and reconstructed images from a batch of data.\n",
    "    \n",
    "    Args:\n",
    "    - model (torch.nn.Module): Trained beta-VAE model.\n",
    "    - data_loader (torch.utils.data.DataLoader): DataLoader for the dataset.\n",
    "    - device (torch.device): Device to run the model on (CPU or GPU).\n",
    "    - beta (float): The beta value used for the loss scaling in beta-VAE.\n",
    "    - num_images (int): Number of images to visualize.\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "\n",
    "    # Retrieve one batch of images from the data loader\n",
    "    images, _, _ = next(iter(data_loader))\n",
    "    images = images[:num_images]\n",
    "\n",
    "    # Disable gradient calculations\n",
    "    with torch.no_grad():\n",
    "        images = images.to(device)\n",
    "        mu, logvar = model.encoder(images)\n",
    "        \n",
    "        # Sample z using the reparameterization trick\n",
    "        std = torch.exp(0.5 * logvar)\n",
    "        epsilon = torch.randn_like(std)\n",
    "        z = mu + std * epsilon\n",
    "\n",
    "        recon_images = model.decoder(z)\n",
    "\n",
    "    img_dim = images[0].cpu().squeeze().shape[0]\n",
    "\n",
    "    # Create separate images for original and reconstructed images\n",
    "    original_image_grid = np.ones((8 * img_dim, 8 * img_dim))  # 8 rows, 8 columns grid\n",
    "    reconstructed_image_grid = np.ones((8 * img_dim, 8 * img_dim))  # 8 rows, 8 columns grid\n",
    "\n",
    "    # Fill in the grids with images\n",
    "    for i in range(8):\n",
    "        for j in range(8):\n",
    "            idx = i * 8 + j\n",
    "            if idx < num_images:\n",
    "                # Fill the original image grid\n",
    "                original_image_grid[i * img_dim:(i + 1) * img_dim, j * img_dim:(j + 1) * img_dim] = images[idx].cpu().squeeze()\n",
    "\n",
    "                # Fill the reconstructed image grid\n",
    "                reconstructed_image_grid[i * img_dim:(i + 1) * img_dim, j * img_dim:(j + 1) * img_dim] = recon_images[idx].cpu().squeeze()\n",
    "\n",
    "    # Save the original images only for the first beta = 1\n",
    "    if beta == 1:\n",
    "        plt.figure(figsize=(10, 10))\n",
    "        plt.imshow(original_image_grid, cmap='gray')\n",
    "        plt.axis('off')\n",
    "        plt.tight_layout()\n",
    "        plt.savefig(f\"{SAVE_PATH}/original_images.png\", bbox_inches='tight', pad_inches=0)\n",
    "        plt.close()\n",
    "\n",
    "    # Save the reconstructed images for all betas\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(reconstructed_image_grid, cmap='gray')\n",
    "    plt.axis('off')\n",
    "    plt.tight_layout()\n",
    "\n",
    "    plt.savefig(f\"{SAVE_PATH}/reconstructed_images_beta_{beta}.png\", bbox_inches='tight', pad_inches=0)    \n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c24d781-61c5-4625-af82-63b7f91b680c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_latent_traversal_combined(model, device, gif_path, distinct_shape_indices, dataset, latent_dim=10, num_traversals=11):\n",
    "    \"\"\"\n",
    "    Visualize latent space traversal for dSprites dataset, creating both PNG and GIF outputs.\n",
    "\n",
    "    Args:\n",
    "    - model (torch.nn.Module): Trained Beta-VAE model.\n",
    "    - device (torch.device): Device to run the model on (CPU or GPU).\n",
    "    - gif_path (str): Path to save the generated GIF and PNG files.\n",
    "    - distinct_shape_indices (list): List of indices representing distinct shapes in the dataset.\n",
    "    - dataset (Dataset): dSprites dataset object.\n",
    "    - latent_dim (int): The number of latent dimensions (default=10).\n",
    "    - num_traversals (int): Number of traversal steps between -2 and 2 (default=11).\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "\n",
    "    # Define the interpolation range (from -2 to 2)\n",
    "    interpolation = torch.linspace(-2, 2, num_traversals)\n",
    "\n",
    "    # Prepare for storing frames for the final combined GIF (3 rows x num_traversals columns)\n",
    "    combined_gif_frames_all_shapes = [[] for _ in range(num_traversals)]\n",
    "\n",
    "    # Loop through each distinct shape (3 rows in the final GIF)\n",
    "    for row, idx in enumerate(distinct_shape_indices):\n",
    "        image, _, _ = dataset[idx]\n",
    "        image = image.unsqueeze(0).to(device)\n",
    "\n",
    "        mu, logvar = model.encoder(image)\n",
    "        std = torch.exp(0.5 * logvar)\n",
    "        eps = torch.randn_like(std)\n",
    "        z = mu + eps * std      \n",
    "\n",
    "        # Prepare for storing frames for the current shape's latent traversal\n",
    "        combined_gif_frames = [[] for _ in range(num_traversals)]\n",
    "\n",
    "        # Perform traversal for each latent dimension\n",
    "        for latent_idx in range(latent_dim):\n",
    "            z_ori = z.clone()  # Clone the latent vector\n",
    "\n",
    "            # Perform interpolation across the specified latent dimension\n",
    "            for i, alpha in enumerate(interpolation):\n",
    "                z_copy = z_ori.clone()\n",
    "                z_copy[:, latent_idx] = alpha  # Set this latent dimension to the interpolation value\n",
    "\n",
    "                # Generate the reconstructed image using the decoder\n",
    "                with torch.no_grad():\n",
    "                    generated_image = model.decoder(z_copy)\n",
    "\n",
    "                # Convert to numpy array and append the frame for saving\n",
    "                gen_img = generated_image.squeeze().cpu().numpy()\n",
    "                combined_gif_frames[i].append(np.uint8(255 * gen_img))\n",
    "\n",
    "        # Create grid for latent traversals (one shape, 10 latent dimensions, 11 traversal steps)\n",
    "        gif_frames_with_10_columns = [np.hstack(frame_list) for frame_list in combined_gif_frames]\n",
    "\n",
    "        # Store frames for all shapes into the combined GIF structure\n",
    "        for i in range(num_traversals):\n",
    "            if row == 0:\n",
    "                combined_gif_frames_all_shapes[i] = [gif_frames_with_10_columns[i]]\n",
    "            else:\n",
    "                combined_gif_frames_all_shapes[i].append(gif_frames_with_10_columns[i])\n",
    "\n",
    "        # Generate and save PNG for the current shape's latent traversals\n",
    "        png_output_path = f\"{gif_path}_image_{row + 1}.png\"\n",
    "        save_png_latent_traversal_only(combined_gif_frames, png_output_path, latent_dim, num_traversals)\n",
    "\n",
    "    # Combine frames of all shapes for the final GIF (3 rows x num_traversals columns)\n",
    "    final_gif_frames = [np.vstack(frame_list) for frame_list in combined_gif_frames_all_shapes]\n",
    "\n",
    "    # Save the final GIF combining all shapes and latent traversals\n",
    "    gif_output_path = f\"{gif_path}_combined.gif\"\n",
    "    save_combined_gif(final_gif_frames, gif_output_path)\n",
    "\n",
    "# Helper function to save the final combined GIF (3 rows and 10 columns) with column titles\n",
    "def save_combined_gif(frames, output_path):\n",
    "    \"\"\"\n",
    "    Save the combined GIF that shows the latent traversal for multiple shapes.\n",
    "\n",
    "    Args:\n",
    "    - frames (list): List of frames to include in the GIF.\n",
    "    - output_path (str): Path to save the GIF.\n",
    "    \"\"\"\n",
    "    gif_frames_with_labels = []\n",
    "\n",
    "    for frame in frames:\n",
    "        if frame is not None and frame.size > 0:\n",
    "            fig, ax = plt.subplots(figsize=(20, 6))\n",
    "            ax.imshow(frame, cmap='gray', vmin=0, vmax=255)  # Display the frame in grayscale\n",
    "            ax.axis('off')\n",
    "\n",
    "            # Add column labels (x_1 to x_10) for each latent dimension\n",
    "            column_labels = [f'$x_{{{i + 1}}}$' for i in range(10)]\n",
    "            for i, label in enumerate(column_labels):\n",
    "                x_pos = (frame.shape[1] / 10) * (i + 0.5)  # Calculate label position\n",
    "                ax.text(x_pos, -10, label, fontsize=14, ha='center', va='bottom', color='black')\n",
    "\n",
    "            # Ensure the figure is drawn before converting to an image\n",
    "            fig.canvas.draw()\n",
    "\n",
    "            # Convert the figure to an image array and append it to the list of frames\n",
    "            gif_frame = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8).reshape(fig.canvas.get_width_height()[::-1] + (4,))\n",
    "            gif_frames_with_labels.append(gif_frame)\n",
    "            plt.close(fig)\n",
    "\n",
    "    # Save the GIF\n",
    "    if gif_frames_with_labels:\n",
    "        # Convert images to uint8\n",
    "        gif_frames_with_labels = [np.uint8(frame) for frame in gif_frames_with_labels]\n",
    "        imageio.mimsave(output_path, gif_frames_with_labels, fps=5, loop=0)  # loop=0 for infinite looping\n",
    "    else:\n",
    "        print(\"No valid frames to save for GIF.\")\n",
    "\n",
    "# Function to save PNG for latent traversals with row labels for each latent dimension\n",
    "def save_png_latent_traversal_only(traversal_images, output_path, latent_dim=10, num_traversals=11):\n",
    "    \"\"\"\n",
    "    Save a PNG file for visualizing the latent traversals across all dimensions.\n",
    "\n",
    "    Args:\n",
    "    - traversal_images (list): List of images generated during the traversal.\n",
    "    - output_path (str): Path to save the PNG.\n",
    "    - latent_dim (int): Number of latent dimensions.\n",
    "    - num_traversals (int): Number of interpolation steps.\n",
    "    \"\"\"\n",
    "    # Create a figure with tight-fitting subplots\n",
    "    fig, axes = plt.subplots(latent_dim, num_traversals, figsize=(num_traversals * 2, latent_dim * 2))\n",
    "\n",
    "    # Define the interpolation values for the column labels from -2 to 2\n",
    "    interpolation_values = np.linspace(-2, 2, num_traversals)\n",
    "\n",
    "    # Plot the traversal images in a grid for all latent dimensions\n",
    "    for latent_idx in range(latent_dim):\n",
    "        for frame_idx in range(num_traversals):\n",
    "            axes[latent_idx, frame_idx].imshow(traversal_images[frame_idx][latent_idx], cmap='gray', aspect='auto')\n",
    "            axes[latent_idx, frame_idx].axis('off')\n",
    "\n",
    "    # Add row labels for each latent dimension (x_1 through x_10)\n",
    "    row_labels = [f'$x_{{{i + 1}}}$' for i in range(latent_dim)]\n",
    "    for i, label in enumerate(row_labels):\n",
    "        axes[i, 0].annotate(label, xy=(-0.05, 0.5), xycoords='axes fraction', xytext=(-30, 0), textcoords='offset points',\n",
    "                            ha='center', va='center', fontsize=18, color='black')\n",
    "\n",
    "    # Add column labels for x-axis (-2 to +2)\n",
    "    for frame_idx in range(num_traversals):\n",
    "        axes[0, frame_idx].set_title(f'{interpolation_values[frame_idx]:.1f}', fontsize=18, color='black')\n",
    "\n",
    "    plt.subplots_adjust(left=0.05, right=0.99, top=0.9, bottom=0.1, wspace=0, hspace=0)\n",
    "\n",
    "    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fced51a9-1ec3-4589-8116-fd97a24d67af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to extract latent representations and ground truth factors\n",
    "def extract_latents_and_factors(model, data_loader, device, flag):\n",
    "    \"\"\"\n",
    "    Extract latent representations and ground truth factors.\n",
    "\n",
    "    Args:\n",
    "    - model (torch.nn.Module): Trained beta-VAE model.\n",
    "    - data_loader (torch.utils.data.DataLoader): DataLoader object to load the dataset.\n",
    "    - device (torch.device): Device to run the model on (CPU or GPU).\n",
    "    - flag (int): 0 for training data, 1 for validation data.\n",
    "    \n",
    "    Returns:\n",
    "    - latent_classes (numpy.ndarray): Discrete ground truth factors for the data.\n",
    "    - latent_values (numpy.ndarray): Continuous ground truth factors for the data.\n",
    "    - mu_list (numpy.ndarray): List of mean values for the latent variables.\n",
    "    - std_list (numpy.ndarray): List of standard deviations for the latent variables.\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    latent_classes_list, latent_values_list = [], []\n",
    "    mu_list, std_list = [], []\n",
    "\n",
    "    desc = 'Training' if flag == 0 else 'Validation'\n",
    "    \n",
    "    # No gradient computation needed during evaluation\n",
    "    with torch.no_grad():\n",
    "        for image, latent_classes, latent_values in tqdm(data_loader, desc=desc):\n",
    "            image = image.to(device)\n",
    "            \n",
    "            mu, logvar = model.encoder(image)\n",
    "            sigma = torch.exp(0.5 * logvar)\n",
    "            \n",
    "            latent_classes_list.append(latent_classes)\n",
    "            latent_values_list.append(latent_values)\n",
    "            mu_list.append(mu.cpu())\n",
    "            std_list.append(sigma.cpu())\n",
    "\n",
    "    latent_classes = torch.cat(latent_classes_list, dim=0).numpy()\n",
    "    latent_values = torch.cat(latent_values_list, dim=0).numpy()\n",
    "    mu_list = torch.cat(mu_list, dim=0).numpy()\n",
    "    std_list = torch.cat(std_list, dim=0).numpy()\n",
    "    \n",
    "    return latent_classes, latent_values, mu_list, std_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aca34b1e-324f-4945-85a1-de1cafa14117",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    for beta in BETAS:\n",
    "        \n",
    "        # Define model and paths\n",
    "        pretrained_model_path = os.path.join(SAVE_PATH, f'beta_vae_{beta}.pth')\n",
    "        model = BetaVAE(latent_dim=LATENT_DIM).to(DEVICE)\n",
    "        \n",
    "        # Load pre-trained model if available\n",
    "        if os.path.exists(pretrained_model_path):\n",
    "            print(colored(f\"Loading pre-trained model for beta = {beta}...\", \"red\", attrs=['bold']))\n",
    "            model.load_state_dict(torch.load(pretrained_model_path, map_location=DEVICE))\n",
    "        else:\n",
    "            print(colored(f\"\\nNo pre-trained model found. Training beta-VAE with beta = {beta} from scratch...\", \"red\", attrs=['bold']))\n",
    "        \n",
    "        new_model_path = os.path.join(SAVE_PATH, f'beta_vae_{beta}_v1.pth')  # Save new model\n",
    "        optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\n",
    "            \n",
    "        print(colored(f\"\\nProcessing with beta = {beta}...\", \"green\", attrs=['bold']))\n",
    "\n",
    "        # Train the model and track the losses\n",
    "        train_results = train_beta_vae(model, train_loader, val_loader, optimizer, beta_vae_loss, beta, DEVICE, NUM_EPOCHS)\n",
    "\n",
    "        # Unpack the training results\n",
    "        train_loss_hist = train_results['train_loss']\n",
    "        val_loss_hist = train_results['val_loss']\n",
    "        kl_loss_hist = train_results['kl_loss']\n",
    "        val_kl_loss_hist = train_results['val_kl_loss']\n",
    "        bce_loss_hist = train_results['bce_loss']\n",
    "        val_bce_loss_hist = train_results['val_bce_loss']\n",
    "\n",
    "        np.save(os.path.join(SAVE_PATH, f'train_loss_hist_beta_{beta}.npy'), train_loss_hist)\n",
    "        np.save(os.path.join(SAVE_PATH, f'val_loss_hist_beta_{beta}.npy'), val_loss_hist)\n",
    "        np.save(os.path.join(SAVE_PATH, f'kl_loss_hist_beta_{beta}.npy'), kl_loss_hist)\n",
    "        np.save(os.path.join(SAVE_PATH, f'val_kl_loss_hist_beta_{beta}.npy'), val_kl_loss_hist)\n",
    "        np.save(os.path.join(SAVE_PATH, f'bce_loss_hist_beta_{beta}.npy'), bce_loss_hist)\n",
    "        np.save(os.path.join(SAVE_PATH, f'val_bce_loss_hist_beta_{beta}.npy'), val_bce_loss_hist)\n",
    "        \n",
    "        # Save the newly trained model\n",
    "        torch.save(model.state_dict(), new_model_path)\n",
    "\n",
    "        # Extract latent and ground truth factors for both training and validation datasets\n",
    "        print(colored(f\"\\nExtracting latent vectors and ground truth factors for beta = {beta}...\", \"blue\", attrs=['bold']))\n",
    "        latent_classes_train, latent_values_train, mu_list_train, std_list_train = extract_latents_and_factors(model, train_loader, DEVICE, 0)\n",
    "        latent_classes_val, latent_values_val, mu_list_val, std_list_val = extract_latents_and_factors(model, val_loader, DEVICE, 1)\n",
    "\n",
    "        np.save(os.path.join(SAVE_PATH, f'train_latent_classes_beta_{beta}.npy'), latent_classes_train)\n",
    "        np.save(os.path.join(SAVE_PATH, f'train_latent_values_beta_{beta}.npy'), latent_values_train)\n",
    "        np.save(os.path.join(SAVE_PATH, f'train_mu_beta_{beta}.npy'), mu_list_train)\n",
    "        np.save(os.path.join(SAVE_PATH, f'train_std_beta_{beta}.npy'), std_list_train)\n",
    "        np.save(os.path.join(SAVE_PATH, f'val_latent_classes_beta_{beta}.npy'), latent_classes_val)\n",
    "        np.save(os.path.join(SAVE_PATH, f'val_latent_values_beta_{beta}.npy'), latent_values_val)\n",
    "        np.save(os.path.join(SAVE_PATH, f'val_mu_beta_{beta}.npy'), mu_list_val)\n",
    "        np.save(os.path.join(SAVE_PATH, f'val_std_beta_{beta}.npy'), std_list_val)\n",
    "        \n",
    "        # Plot training and validation losses\n",
    "        print(colored(f\"\\nPlotting training and validation losses for beta = {beta}...\", \"blue\", attrs=['bold']))\n",
    "        plot_beta_vae_losses(train_loss_hist, val_loss_hist, kl_loss_hist, val_kl_loss_hist, bce_loss_hist, val_bce_loss_hist, beta)\n",
    "\n",
    "        # Visualization functions\n",
    "        visualize_images_grid(model, val_loader, DEVICE, beta)\n",
    "\n",
    "        gif_path = os.path.join(SAVE_PATH, f'latent_traversal_beta_{beta}')\n",
    "        visualize_latent_traversal_combined(model, DEVICE, gif_path, distinct_shape_indices, val_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2698cdbd-46d4-4bdb-a4cd-27c3619511d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set random seed for reproducibility\n",
    "SEED = 42\n",
    "torch.manual_seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "random.seed(SEED)\n",
    "\n",
    "# Hyperparameters and constants\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "NPZ_FILE = \"dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz\"\n",
    "SAVE_PATH = \"./betaVAE_results_v1/\"\n",
    "LATENT_DIM = 10  # Dimensionality of the latent space\n",
    "BATCH_SIZE = 64  # Batch size for training and validation\n",
    "LEARNING_RATE = 1e-4  # Learning rate for the optimizer\n",
    "NUM_EPOCHS = 50  # Number of training epochs\n",
    "BETAS = [0.1, 1, 2, 4, 8, 16]  # Different beta values for beta-VAE\n",
    "\n",
    "# Load dataset\n",
    "dataset = DspritesDataset(NPZ_FILE)\n",
    "\n",
    "# Ensure the results directory exists\n",
    "if not os.path.exists(SAVE_PATH):\n",
    "    os.makedirs(SAVE_PATH)\n",
    "\n",
    "# Define file paths for the indices\n",
    "train_indices_file = os.path.join(SAVE_PATH, \"train_indices.npy\")\n",
    "val_indices_file = os.path.join(SAVE_PATH, \"val_indices.npy\")\n",
    "\n",
    "# Split the dataset\n",
    "train_size = int(0.9 * len(dataset))\n",
    "val_size = len(dataset) - train_size\n",
    "\n",
    "# Check if saved indices exist for a consistent split across runs\n",
    "if os.path.exists(train_indices_file) and os.path.exists(val_indices_file):\n",
    "    # Load saved indices using numpy\n",
    "    train_indices = np.load(train_indices_file)\n",
    "    val_indices = np.load(val_indices_file)\n",
    "else:\n",
    "    # Save the indices for reproducibility\n",
    "    train_indices, val_indices = torch.utils.data.random_split(\n",
    "        range(len(dataset)), [train_size, val_size], generator=torch.Generator().manual_seed(SEED)\n",
    "    )\n",
    "\n",
    "    # Save indices using numpy\n",
    "    np.save(train_indices_file, train_indices)\n",
    "    np.save(val_indices_file, val_indices)\n",
    "\n",
    "# Create subsets using saved indices\n",
    "train_dataset = Subset(dataset, train_indices)\n",
    "val_dataset = Subset(dataset, val_indices)\n",
    "\n",
    "# Data loaders\n",
    "train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
    "\n",
    "# Select distinct shape indices (square, ellipse, heart) from the validation dataset\n",
    "distinct_shape_indices = select_3_distinct_shape_indices(val_dataset, list(range(val_size)))\n",
    "visualize_3_images_from_shapes(val_dataset, distinct_shape_indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df5dd559-9c63-43cd-a5bb-0d8c50181661",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Main script execution\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}