2sVAE_UTKFace.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-07-18T10:31:11.555924Z",
     "start_time": "2024-07-18T10:31:02.700565Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import wandb\n",
    "from PIL import Image\n",
    "from torchvision.transforms import ToPILImage\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torchvision.utils as vutils\n",
    "import torchvision.datasets as datasets\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "\n",
    "from pythae.models import AutoModel\n",
    "from pythae.data.datasets import DatasetOutput\n",
    "from pythae.samplers import (\n",
    "    TwoStageVAESampler,\n",
    "    TwoStageVAESamplerConfig,\n",
    "    NormalSampler\n",
    ")\n",
    "from pythae.trainers import BaseTrainerConfig\n",
    "\n",
    "device = torch.device(\"mps\") #note this may need to be cpu for some ppl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "        # get the name of the GPU\n",
    "        gpu_name = torch.cuda.get_device_name(0)\n",
    "        print(f\"CUDA is available. GPU Name: {gpu_name}\")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "396dbabddbbb8362"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "### IMAGE DATASET pytorch class\n",
    "class UTKFaceDataset(datasets.ImageFolder):\n",
    "    def __init__(self, root, transform=None, target_transform=None):\n",
    "        super().__init__(root=root, transform=transform, target_transform=target_transform)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        path, _ = self.imgs[index]\n",
    "        X, _ = super().__getitem__(index)\n",
    "        filename = os.path.basename(path)\n",
    "        return DatasetOutput(data=X, filename=filename)\n",
    "\n",
    "transform = transforms.Compose([transforms.ToTensor(),])\n",
    "\n",
    "#Create UTKFACE dataset\n",
    "all_dataset = UTKFaceDataset(root=\"./data\", transform=transform)\n",
    "\n",
    "#Split UTKFACE dataset into train and eval sets\n",
    "total_size = len(all_dataset)\n",
    "train_size = 20000\n",
    "eval_size = total_size - train_size \n",
    "\n",
    "#split dataset randomly into train and eval datasets\n",
    "train_dataset, eval_dataset = random_split(all_dataset, [train_size, eval_size])\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "43be34f594adf761"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "indices = torch.randperm(len(all_dataset))[:25]\n",
    "images = [all_dataset[i]['data'] for i in indices]\n",
    "grid = vutils.make_grid(images, nrow=5)\n",
    "grid_pil = transforms.ToPILImage()(grid)\n",
    "grid_pil"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "376ef01bb2ce4d4b"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "wandb.init(\n",
    "    project='VAE_UTKFACE',\n",
    "    entity='charlesdoyne'\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7384562042b6ffb8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "wandb.log({\"UTKFace dataset\": [wandb.Image(grid_pil, caption=\"Images used to train main VAE\")]})"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "b77f19f9d6f40d46"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "### LOAD LATEST MODEL\n",
    "\"\"\"\n",
    "last_training = sorted(os.listdir('my_model'))[-1]\n",
    "trained_model = AutoModel.load_from_folder(os.path.join('my_model', last_training, 'final_model'))\n",
    "\"\"\"\n",
    "trained_model = AutoModel.load_from_folder(\"./my_model/bestmodel/final_model\")\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e25f59e9bc34d745"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "### GENERATION WITH NORMAL SAMPLER\n",
    "normal_sampler = NormalSampler(\n",
    "    model=trained_model\n",
    ")\n",
    "\n",
    "gen_data = normal_sampler.sample(\n",
    "    num_samples=25\n",
    ")\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "95942e0e2a38c9be"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "grid = vutils.make_grid(gen_data, nrow=5, normalize=True, scale_each=True)\n",
    "grid_np = grid.permute(1, 2, 0).numpy()  # Convert to HxWxC layout\n",
    "grid_pil = Image.fromarray((grid_np * 255).astype('uint8'), 'RGB')\n",
    "grid_pil"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "1edd89536366729a"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "wandb.log({\"Normal Sampler Generations\": wandb.Image(grid_pil, caption=\"Images generated with Normal Sampler\")})"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "9eaade2c57ace18c"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "trained_model = trained_model.to(device)\n",
    "\n",
    "reconstructions = []\n",
    "for i in range(25):\n",
    "    dataset_output = eval_dataset[i]\n",
    "    image = dataset_output.data\n",
    "    image = image.to(device) \n",
    "    reconstruction = trained_model.reconstruct(image.unsqueeze(0)).detach().cpu()\n",
    "    reconstructions.append(reconstruction)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "50a52315aa5d6d7"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "num_pairs = 8\n",
    "\n",
    "image_pairs = []\n",
    "for i in range(num_pairs):\n",
    "    original = eval_dataset[i].data\n",
    "    original = original.to(device) \n",
    "    reconstructed = reconstructions[i]\n",
    "    original_cpu = original.cpu()\n",
    "    pair = torch.cat((original_cpu, reconstructed.squeeze(0)), 2) \n",
    "    image_pairs.append(pair)\n",
    "\n",
    "n_rows = 2\n",
    "grid = vutils.make_grid(image_pairs, nrow=n_rows, padding=2)\n",
    "grid_np = grid.cpu().numpy()\n",
    "grid_np = np.transpose(grid_np, (1, 2, 0))\n",
    "grid_np = grid_np - grid_np.min()  # Normalize to the range 0 - max\n",
    "grid_np = grid_np / grid_np.max()  # Normalize to the range 0 - 1\n",
    "grid_np = (grid_np * 255).astype(np.uint8)  # Scale to range 0 - 255\n",
    "grid_pil = Image.fromarray(grid_np)\n",
    "grid_pil"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "b8c2f9fee711a969"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "wandb.log({\"VAE reconstructions\": wandb.Image(grid_pil, caption=\"How the encoder-decoder reconstructs images\")})"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "9d164f0d25c3f601"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "gender = \"Male\"\n",
    "race = \"White\"\n",
    "age_input = \"30-35\"\n",
    "\n",
    "gender_map = {\"Any\": None, \"Male\": 0, \"Female\": 1}\n",
    "race_map = {\"Any\": None, \"White\": 0, \"Black\": 1, \"Asian\": 2, \"Indian\": 3}\n",
    "\n",
    "gender_code = gender_map[gender]\n",
    "race_code = race_map[race]\n",
    "\n",
    "if '-' in age_input: \n",
    "    age_range = [int(a) for a in age_input.split('-')]\n",
    "elif age_input: \n",
    "    age_range = [int(age_input), int(age_input)]\n",
    "else:\n",
    "    age_range = [0, 116]\n",
    "\n",
    "filtered_dataset = []\n",
    "for item in all_dataset:\n",
    "    _, filename = item.data, item.filename\n",
    "    parts = filename.split('_')\n",
    "    \n",
    "    if len(parts) == 4:\n",
    "        file_age, file_gender, file_race, _ = parts\n",
    "        age_condition = age_range[0] <= int(file_age) <= age_range[1]\n",
    "        gender_condition = True if gender_code is None else int(file_gender) == gender_code\n",
    "        race_condition = True if race_code is None else int(file_race) == race_code\n",
    "\n",
    "        if age_condition and gender_condition and race_condition:\n",
    "            filtered_dataset.append(item)\n",
    "\n",
    "if filtered_dataset:\n",
    "    num_images_to_show = min(len(filtered_dataset), 25)\n",
    "    indices = torch.randperm(len(filtered_dataset))[:num_images_to_show]\n",
    "    images = [filtered_dataset[i].data for i in indices]\n",
    "    grid = vutils.make_grid(images, nrow=5)\n",
    "    attribute_data_pil = ToPILImage()(grid)\n",
    "\n",
    "    filtered_dataset_tensors = [item.data for item in filtered_dataset]  \n",
    "    filtered_dataset_tensor = torch.stack(filtered_dataset_tensors)\n",
    "\n",
    "    split_index = int(0.8 * len(filtered_dataset_tensor))\n",
    "    filtered_dataset_train = filtered_dataset_tensor[:split_index]\n",
    "    filtered_dataset_eval = filtered_dataset_tensor[split_index:]"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "74e8c96a361a0736"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "indices = torch.randperm(len(filtered_dataset))[:25]\n",
    "images = [filtered_dataset[i]['data'] for i in indices]\n",
    "grid = vutils.make_grid(images, nrow=5)\n",
    "grid_pil = transforms.ToPILImage()(grid)\n",
    "grid_pil"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "abe10d8a9b220006"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "wandb.log({\"Style Dataset\": [wandb.Image(grid_pil, caption=\"Style dataset for second VAE\")]})"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "da8bea777e9636ce"
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-12-03T21:57:24.614391Z",
     "start_time": "2023-12-03T21:57:24.145381Z"
    }
   },
   "id": "4bc32be7a442e095"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "### GENERATION WITH TWO STAGE VAE SAMPLER \n",
    "twostagevae_config = TwoStageVAESamplerConfig(\n",
    "    reconstruction_loss='mse', \n",
    ")\n",
    "\n",
    "twostage_sampler = TwoStageVAESampler(\n",
    "    model=trained_model,\n",
    "    sampler_config=twostagevae_config\n",
    ")\n",
    "\n",
    "two_stage_train_config = BaseTrainerConfig(\n",
    "    output_dir='my_model',\n",
    "    learning_rate=2e-4,\n",
    "    per_device_train_batch_size=4,\n",
    "    per_device_eval_batch_size=4,\n",
    "    num_epochs=75, # Change this to train the model a bit more\n",
    ")\n",
    "\n",
    "twostage_sampler.fit(\n",
    "    \n",
    "    train_data=filtered_dataset_train,\n",
    "    eval_data=filtered_dataset_eval,\n",
    "    training_config=two_stage_train_config,\n",
    "    \n",
    ")\n",
    "\n",
    "gen_data2 = twostage_sampler.sample(\n",
    "    num_samples=25\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "4c0349f216ab9d0e"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "grid = vutils.make_grid(gen_data2, nrow=5, normalize=True, scale_each=True)\n",
    "grid_np = grid.permute(1, 2, 0).numpy()\n",
    "grid_pil = Image.fromarray((grid_np * 255).astype('uint8'), 'RGB')\n",
    "grid_pil"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "27fee867096c85b6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "wandb.log({\"Two Stage Sampler Sampler Generations\": wandb.Image(grid_pil, caption=\"Images generated with Two Stage Sampler\")})"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "2e847f04f1f3299b"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "wandb.finish()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "da405c3627db3ee3"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}