Inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use a pretrained model to do inferences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model and Dataset Selection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Select a pretrained model path and eventually a dataset and execute all cells of the notebook. Generated images will be placed in `./inferences/<model-name>/<dataset-name>/`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Type of Model Selection**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -- Takes noise vectors as input -- #\n",
    "\n",
    "# model = \"gan\"\n",
    "# model = \"dcgan\"\n",
    "\n",
    "# -- Takes images as input -- #\n",
    "\n",
    "model = \"fdcgan\"\n",
    "# model = \"resnet\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Dataset Selection**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -- Only adapted to noise vectors input -- #\n",
    "\n",
    "# dataset = \"mnist\"\n",
    "\n",
    "# -- Adapted to both noise vectors and images input -- #\n",
    "\n",
    "mother_dataset = \"apple2orange64\"\n",
    "dataset = \"apple2orange64\"\n",
    "\n",
    "# -- Opposite transformation\n",
    "\n",
    "# mother_dataset = \"apple2orange64\"\n",
    "# dataset = \"custom\"\n",
    "# dataset_path = \"inferences/cycle-gan/\"\n",
    "# source_dataset_name = \"apple2orange64\"\n",
    "# dataset_name = \"orange2apple64\"\n",
    "# my_class_A = 1\n",
    "# my_class_B = \"A\"\n",
    "\n",
    "# mother_dataset = \"apple2orange64\"\n",
    "# dataset = \"orange2apple64\"\n",
    "\n",
    "# -- Opposite transformation\n",
    "\n",
    "# mother_dataset = \"apple2orange64\"\n",
    "# dataset = \"custom\"\n",
    "# dataset_path = \"inferences/cycle-gan/\"\n",
    "# source_dataset_name = \"apple2orange64\"\n",
    "# dataset_name = \"apple2orange64\"\n",
    "# my_class_A = 0\n",
    "# my_class_B = \"B\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Model Path**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# generator_name = \"generator_199\" # no cycle-gan\n",
    "generator_name = \"G_AB_185\" # cycle-gan\n",
    "# generator_name = \"G_BA_185\" # cycle-gan\n",
    "\n",
    "# is_trained_with_cycleGAN = False\n",
    "is_trained_with_cycleGAN = True\n",
    "\n",
    "path_model = model\n",
    "if is_trained_with_cycleGAN:\n",
    "    path_model = \"cycle-gan/\" + model\n",
    "\n",
    "model_path = \"saved_models/\" + path_model + \"/\" + mother_dataset + \"/\" + generator_name + \".pth\"\n",
    "if dataset == \"custom\":\n",
    "    model_path = \"saved_models/\" + path_model + \"/\" + dataset_name + \"/\" + generator_name + \".pth\"\n",
    "    if is_trained_with_cycleGAN:\n",
    "        model_path = \"saved_models/\" + path_model + \"/\" + source_dataset_name + \"/\" + generator_name + \".pth\"\n",
    "\n",
    "# Make sure that model exists\n",
    "\n",
    "if not os.path.isfile(model_path):\n",
    "    print(f\"Model {model_path} does not exist.\")\n",
    "    exit()\n",
    "else:\n",
    "    print(\"Loaded model\", model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import datetime\n",
    "import time\n",
    "import glob\n",
    "import random\n",
    "import classifier_data\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from PIL import Image\n",
    "\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.utils import save_image, make_grid\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data import Dataset\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperparameters and Datasets Specific Configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Global Configuration**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_cpu = 8               # number of cpu threads to use during batch generation\n",
    "batch_size = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Datasets Specifications**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if dataset == \"mnist\":\n",
    "    channels = 1\n",
    "    latent_dim = 100\n",
    "    img_size = 64\n",
    "    \n",
    "elif dataset == \"apple2orange64\" or dataset == \"orange2apple64\":\n",
    "    channels = 3\n",
    "    latent_dim = 300\n",
    "    img_size = 64\n",
    "\n",
    "elif dataset == \"custom\":\n",
    "    channels = 3\n",
    "    latent_dim = 300\n",
    "    img_size = 64\n",
    "\n",
    "else:\n",
    "    raise Exception(\"Unknown dataset\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Importing Chosen Models**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model == \"gan\":\n",
    "    from models.gan import Generator\n",
    "    \n",
    "elif model == \"dcgan\":\n",
    "    from models.dcgan import Generator\n",
    "\n",
    "elif model == \"fdcgan\" or model == \"resnet\":\n",
    "    if dataset not in [\"apple2orange64\", \"orange2apple64\", \"custom\"]:\n",
    "        raise Exception(f\"Dataset {dataset} has no input image for the generator\")\n",
    "    if model == \"fdcgan\":\n",
    "        from models.fdcgan import Generator\n",
    "    elif model == \"resnet\":\n",
    "        from models.resnet import Generator\n",
    "        n_residual_blocks = 9\n",
    "\n",
    "else:\n",
    "    raise Exception(\"Unknown model\")\n",
    "\n",
    "img_shape = (channels, img_size, img_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Setting Up Cuda**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda = True if torch.cuda.is_available() else False\n",
    "device = torch.device(\"cuda\" if cuda else \"cpu\")\n",
    "print(\"Device use for training: \", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialising Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -- Initialize generator -- #\n",
    "\n",
    "if model == \"gan\" or model == \"dcgan\":\n",
    "    generator = Generator(img_shape, latent_dim)\n",
    "elif model == \"fdcgan\":\n",
    "    generator = Generator(img_shape)\n",
    "elif model == \"resnet\":\n",
    "    generator = Generator(img_shape, n_residual_blocks)\n",
    "\n",
    "# Load pretrained models\n",
    "generator.load_state_dict(torch.load(model_path))\n",
    "\n",
    "if cuda:\n",
    "    generator.cuda()\n",
    "\n",
    "Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_rgb(image):\n",
    "    rgb_image = Image.new(\"RGB\", image.size)\n",
    "    rgb_image.paste(image)\n",
    "    return rgb_image\n",
    "\n",
    "class ImageDataset(Dataset):\n",
    "    def __init__(self, root, transforms_=None, unaligned=False, mode=\"train\"):\n",
    "        self.transform = transforms.Compose(transforms_)\n",
    "        self.unaligned = unaligned\n",
    "\n",
    "        self.files_A = sorted(glob.glob(os.path.join(root, \"%s/A\" % mode) + \"/*.*\"))\n",
    "        self.files_B = sorted(glob.glob(os.path.join(root, \"%s/B\" % mode) + \"/*.*\"))\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        image_A = Image.open(self.files_A[index % len(self.files_A)])\n",
    "\n",
    "        if self.unaligned:\n",
    "            image_B = Image.open(self.files_B[random.randint(0, len(self.files_B) - 1)])\n",
    "        else:\n",
    "            image_B = Image.open(self.files_B[index % len(self.files_B)])\n",
    "\n",
    "        # Convert grayscale images to rgb\n",
    "        if image_A.mode != \"RGB\":\n",
    "            image_A = to_rgb(image_A)\n",
    "        if image_B.mode != \"RGB\":\n",
    "            image_B = to_rgb(image_B)\n",
    "\n",
    "        item_A = self.transform(image_A)\n",
    "        item_B = self.transform(image_B)\n",
    "        return {\"A\": item_A, \"B\": item_B}\n",
    "\n",
    "    def __len__(self):\n",
    "        return max(len(self.files_A), len(self.files_B))\n",
    "\n",
    "transforms_ = [\n",
    "    transforms.Resize(int(img_size * 1.12), Image.BICUBIC),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
    "]\n",
    "\n",
    "if dataset == \"mnist\":\n",
    "    \n",
    "    os.makedirs(\"./datasets/mnist\", exist_ok=True)\n",
    "\n",
    "    dataloader = torch.utils.data.DataLoader(\n",
    "        datasets.MNIST(\n",
    "            \"./datasets/mnist\",\n",
    "            train=True,\n",
    "            download=True,\n",
    "            transform=transforms.Compose(\n",
    "                [transforms.Resize(img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5]), transforms.Resize(img_size)]\n",
    "            ),\n",
    "        ),\n",
    "        batch_size=batch_size,\n",
    "        shuffle=True,\n",
    "        num_workers=n_cpu,\n",
    "    )\n",
    "\n",
    "if dataset == \"apple2orange64\":\n",
    "    my_class_A = 0\n",
    "    my_class_B = \"B\"\n",
    "\n",
    "if dataset == \"orange2apple64\":\n",
    "    my_class_A = 1\n",
    "    my_class_B = \"A\"\n",
    "\n",
    "if dataset == \"apple2orange64\" or dataset == \"orange2apple64\":\n",
    "\n",
    "    import subprocess\n",
    "    command = \"bash ./datasets/download_cyclegan_dataset.sh\"\n",
    "    subprocess.run(command, shell=True)\n",
    "\n",
    "    # Test data loader\n",
    "    input_dataset = classifier_data.BinaryClassificationImageDataset(\n",
    "        \"./datasets/apple2orange64/\",\n",
    "        transformations=transforms_,\n",
    "        mode=\"validation\"\n",
    "    )\n",
    "\n",
    "    dataloader = DataLoader(\n",
    "        input_dataset,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=False,\n",
    "        num_workers=n_cpu,\n",
    "    )\n",
    "\n",
    "if dataset == \"custom\":\n",
    "\n",
    "    input_dataset = classifier_data.BinaryClassificationImageDataset(\n",
    "        \"./inferences/\" + model,\n",
    "        transformations=transforms_,\n",
    "        mode=source_dataset_name\n",
    "    )\n",
    "    dataloader = DataLoader(\n",
    "        input_dataset,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=False,\n",
    "        num_workers=n_cpu,\n",
    "    )\n",
    "\n",
    "print(\"Number of mini-batches: \", len(dataloader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unfortunately depending on how the dataset is imported we need to do separated training loops for mnist and apple2orange but the procedure is exactly the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_inference_folder = \"inferences/%s/%s/%s\" % (path_model, dataset, my_class_B)\n",
    "if dataset == \"custom\":\n",
    "    image_inference_folder = \"inferences/%s/%s/%s\" % (path_model, dataset_name, my_class_B)\n",
    "\n",
    "os.makedirs(image_inference_folder, exist_ok=True)\n",
    "\n",
    "print(f\"Using '{model}' with '{dataset}', saving results to '{image_inference_folder}'.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Random vector noise input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**mnist:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if dataset == \"mnist\":\n",
    "#     prev_time = time.time()\n",
    "#     for i, (imgs, _) in enumerate(dataloader):\n",
    "        \n",
    "#         if (imgs.shape[0] != batch_size):\n",
    "#             continue\n",
    "\n",
    "#         # Configure input\n",
    "#         real_imgs = Variable(imgs.type(Tensor))\n",
    "\n",
    "#         # Sample noise as generator input\n",
    "#         z = Variable(Tensor(np.random.normal(0, 1, (real_imgs.shape[0], latent_dim))))\n",
    "\n",
    "#         # Generate a batch of images\n",
    "#         gen_imgs = generator(z)\n",
    "        \n",
    "#         # Loss measures generator's ability to fool the discriminator\n",
    "#         g_loss = adversarial_loss(discriminator(gen_imgs), valid)\n",
    "\n",
    "#         g_loss.backward()\n",
    "#         optimizer_G.step()\n",
    "\n",
    "#         # ---------------------\n",
    "#         #  Train Discriminator\n",
    "#         # ---------------------\n",
    "\n",
    "#         optimizer_D.zero_grad()\n",
    "\n",
    "#         # Measure discriminator's ability to classify real from generated samples\n",
    "#         real_loss = adversarial_loss(discriminator(real_imgs), valid)\n",
    "#         fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)\n",
    "#         d_loss = (real_loss + fake_loss) / 2\n",
    "\n",
    "#         d_loss.backward()\n",
    "#         optimizer_D.step()\n",
    "\n",
    "#         # --------------\n",
    "#         #  Log Progress\n",
    "#         # --------------\n",
    "        \n",
    "#         # Determine approximate time left\n",
    "#         batches_done = epoch * len(dataloader) + i\n",
    "#         batches_left = n_epochs * len(dataloader) - batches_done\n",
    "#         time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n",
    "#         prev_time = time.time()\n",
    "\n",
    "#         sys.stdout.write(\n",
    "#             \"\\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] ETA: %s\"\n",
    "#             % (epoch, n_epochs, i, len(dataloader), d_loss.item(), g_loss.item(), time_left))\n",
    "\n",
    "#         batches_done = epoch * len(dataloader) + i\n",
    "#         if batches_done % sample_interval == 0:\n",
    "#             save_image(gen_imgs.data[:batch_size], image_progress_folder + \"/%d.png\" % batches_done, nrow=5, normalize=True)\n",
    "    \n",
    "#     if checkpoint_interval != -1 and epoch % checkpoint_interval == 0:\n",
    "#         # Save model checkpoints\n",
    "#         torch.save(generator.state_dict(), \"saved_models/%s/%s/generator_%d.pth\" % (model, dataset, epoch))\n",
    "#         torch.save(discriminator.state_dict(), \"saved_models/%s/%s/generator_%d.pth\" % (model, dataset, epoch))\n",
    "                "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**orange2apple or apple2orange with gan or dcgan:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if (dataset == \"orange2apple64\" or dataset == \"apple2orange64\") and (model == \"gan\" or model == \"dcgan\"):\n",
    "#     prev_time = time.time()\n",
    "#     for epoch in range(start_epoch, n_epochs):\n",
    "#         for i, batch in enumerate(dataloader):\n",
    "\n",
    "#             # Set model input\n",
    "#             real_imgs = Variable(batch[my_class_B].type(Tensor))\n",
    "            \n",
    "#             if (real_imgs.shape[0] != batch_size):\n",
    "#                 continue\n",
    "\n",
    "#             # Adversarial ground truths\n",
    "#             valid = Variable(Tensor(np.ones((real_imgs.size(0), *discriminator.output_shape))), requires_grad=False)\n",
    "#             fake = Variable(Tensor(np.zeros((real_imgs.size(0), *discriminator.output_shape))), requires_grad=False)\n",
    "\n",
    "#             # -----------------\n",
    "#             #  Train Generator\n",
    "#             # -----------------\n",
    "\n",
    "#             optimizer_G.zero_grad()\n",
    "\n",
    "#             # Sample noise as generator input\n",
    "#             z = Variable(Tensor(np.random.normal(0, 1, (real_imgs.shape[0], latent_dim))))\n",
    "\n",
    "#             # Generate a batch of images\n",
    "#             gen_imgs = generator(z)\n",
    "            \n",
    "#             # Loss measures generator's ability to fool the discriminator\n",
    "#             g_loss = adversarial_loss(discriminator(gen_imgs), valid)\n",
    "\n",
    "#             g_loss.backward()\n",
    "#             optimizer_G.step()\n",
    "\n",
    "#             # ---------------------\n",
    "#             #  Train Discriminator\n",
    "#             # ---------------------\n",
    "\n",
    "#             optimizer_D.zero_grad()\n",
    "\n",
    "#             # Measure discriminator's ability to classify real from generated samples\n",
    "#             real_loss = adversarial_loss(discriminator(real_imgs), valid)\n",
    "#             fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)\n",
    "#             d_loss = (real_loss + fake_loss) / 2\n",
    "\n",
    "#             d_loss.backward()\n",
    "#             optimizer_D.step()\n",
    "\n",
    "#             # --------------\n",
    "#             #  Log Progress\n",
    "#             # --------------\n",
    "            \n",
    "#             # Determine approximate time left\n",
    "#             batches_done = epoch * len(dataloader) + i\n",
    "#             batches_left = n_epochs * len(dataloader) - batches_done\n",
    "#             time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n",
    "#             prev_time = time.time()\n",
    "\n",
    "#             sys.stdout.write(\n",
    "#                 \"\\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] ETA: %s\"\n",
    "#                 % (epoch, n_epochs, i, len(dataloader), d_loss.item(), g_loss.item(), time_left))\n",
    "\n",
    "#             batches_done = epoch * len(dataloader) + i\n",
    "#             if batches_done % sample_interval == 0:\n",
    "#                 save_image(gen_imgs.data[:batch_size], image_progress_folder + \"/%d.png\" % batches_done, nrow=5, normalize=True)\n",
    "        \n",
    "#         if checkpoint_interval != -1 and epoch % checkpoint_interval == 0:\n",
    "#             # Save model checkpoints\n",
    "#             torch.save(generator.state_dict(), \"saved_models/%s/%s/generator_%d.pth\" % (model, dataset, epoch))\n",
    "#             torch.save(discriminator.state_dict(), \"saved_models/%s/%s/generator_%d.pth\" % (model, dataset, epoch))\n",
    "                    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**orange2apple or apple2orange with fdcgan or resnet (improved fdgcan):**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if (dataset == \"orange2apple64\" or dataset == \"apple2orange64\" or dataset == \"custom\") and (model == \"fdcgan\" or model == \"resnet\"):\n",
    "    prev_time = time.time()\n",
    "    print(dataloader)\n",
    "    for i, batch in enumerate(dataloader):\n",
    "\n",
    "        # Set model input\n",
    "        if dataset == \"custom\":\n",
    "            x, y = batch\n",
    "            x = x.to(device)\n",
    "            real_A = Variable(x)\n",
    "        else:\n",
    "            x, y = batch\n",
    "            if (y != my_class_A):\n",
    "                continue\n",
    "            x = x.to(device)\n",
    "            real_A = Variable(x)\n",
    "\n",
    "        # ------------------\n",
    "        #  Use Generators\n",
    "        # ------------------\n",
    "\n",
    "        generator.eval()\n",
    "\n",
    "        fake_B = generator(real_A)\n",
    "\n",
    "        # --------------\n",
    "        #  Log Progress\n",
    "        # --------------\n",
    "\n",
    "        # # Determine approximate time left\n",
    "        batches_done = len(dataloader) + i\n",
    "        batches_left = len(dataloader) - batches_done\n",
    "        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n",
    "        prev_time = time.time()\n",
    "\n",
    "        # Print log\n",
    "        sys.stdout.write(\n",
    "            \"\\r[Batch %d/%d] ETA: %s\"\n",
    "            % (\n",
    "                i,\n",
    "                len(dataloader),\n",
    "                time_left,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        # --------------\n",
    "        #  Save Image\n",
    "        # --------------\n",
    "\n",
    "        image_name = image_inference_folder + \"/%d.png\" % i\n",
    "        save_image(fake_B, image_name, normalize=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "GAN_torch",
   "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.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}