tutorials/tutorial7_image_generation_using_mednist.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "41a16e01",
   "metadata": {},
   "source": [
    "# Tutorial 7: Image generation\n",
    "\n",
    "Synthcity supports generating synthetic images. In this tutorial, we will train a generator  based on the [MedNIST dataset](https://medmnist.com/). The Tutorial is adapted from a [MONAI example](https://github.com/Project-MONAI/tutorials/blob/main/2d_classification/mednist_tutorial.ipynb).\n",
    "\n",
    "The main components are:\n",
    " - Creating an `ImageDataloader` on top of the MedNIST dataset.\n",
    " - Training a Conditional GAN on the resulted dataloader.\n",
    " - Benchmarking the quality of the synthetic images.\n",
    " \n",
    "__Disclaimer__: The models used for the Generators and the Discriminators are not state of the art. For adding better architectures, please update the `suggest_image_generator_discriminator_arch` options from the `convnet.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae15b589",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install synthcity\n",
    "!pip uninstall -y torchaudio torchdata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "889c3c45",
   "metadata": {},
   "source": [
    "## Load MedNIST\n",
    "\n",
    "The dataset is downloaded using [MONAI](https://monai.io/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7c34bea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download MedNIST\n",
    "# stdlib\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "# third party\n",
    "import PIL\n",
    "from monai.apps import download_and_extract\n",
    "\n",
    "workspace = Path(\"workspace\")\n",
    "workspace.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "resource = \"https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/MedNIST.tar.gz\"\n",
    "md5 = \"0bc7306e7427e00ad1c5526a6677552d\"\n",
    "\n",
    "compressed_file = workspace / \"MedNIST.tar.gz\"\n",
    "data_dir = workspace / \"MedNIST\"\n",
    "\n",
    "if not data_dir.exists():\n",
    "    download_and_extract(resource, compressed_file, workspace, md5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0608986",
   "metadata": {},
   "outputs": [],
   "source": [
    "LIMIT = 1000  # samples per class\n",
    "\n",
    "class_names = sorted(\n",
    "    x for x in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, x))\n",
    ")\n",
    "num_class = len(class_names)\n",
    "image_files = [\n",
    "    [\n",
    "        os.path.join(data_dir, class_names[i], x)\n",
    "        for x in os.listdir(os.path.join(data_dir, class_names[i]))\n",
    "    ]\n",
    "    for i in range(num_class)\n",
    "]\n",
    "num_each = [len(image_files[i]) for i in range(num_class)]\n",
    "image_files_list = []\n",
    "image_class = []\n",
    "for i in range(num_class):\n",
    "    image_files_list.extend(image_files[i][:LIMIT])\n",
    "    image_class.extend([i] * min(num_each[i], LIMIT))\n",
    "num_total = len(image_class)\n",
    "image_width, image_height = PIL.Image.open(image_files_list[0]).size\n",
    "\n",
    "print(f\"Total image count: {num_total}\")\n",
    "print(f\"Image dimensions: {image_width} x {image_height}\")\n",
    "print(f\"Label names: {class_names}\")\n",
    "print(f\"Label counts: {num_each}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1453db8c",
   "metadata": {},
   "source": [
    "## Visualize random samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a7b2717",
   "metadata": {},
   "outputs": [],
   "source": [
    "# third party\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "plt.subplots(3, 3, figsize=(8, 8))\n",
    "for i, k in enumerate(np.random.randint(num_total, size=9)):\n",
    "    im = PIL.Image.open(image_files_list[k])\n",
    "    arr = np.array(im)\n",
    "    plt.subplot(3, 3, i + 1)\n",
    "    plt.xlabel(class_names[image_class[k]])\n",
    "    plt.imshow(arr, cmap=\"gray\", vmin=0, vmax=255)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5424a2c5",
   "metadata": {},
   "source": [
    "## Create the ImageDataLoader\n",
    "\n",
    "The ImageDataLoader prepares the image dataset for the `synthcity` generators.\n",
    "\n",
    "Internally, the dataloader will resize the data to the `(height, width)` parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75f1f078",
   "metadata": {},
   "outputs": [],
   "source": [
    "# third party\n",
    "import torch\n",
    "\n",
    "# synthcity absolute\n",
    "from synthcity.plugins.core.dataloader import ImageDataLoader\n",
    "\n",
    "IMG_SIZE = 64\n",
    "\n",
    "\n",
    "class MedNISTDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, image_files, labels):\n",
    "        self.image_files = image_files\n",
    "        self.image_cache = {}\n",
    "        self.labels = labels\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_files)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        if index in self.image_cache:\n",
    "            img = self.image_cache[index]\n",
    "        else:\n",
    "            img = PIL.Image.open(self.image_files[index])\n",
    "            img = np.asarray(img)\n",
    "            self.image_cache[index] = img\n",
    "\n",
    "        return img, self.labels[index]\n",
    "\n",
    "\n",
    "dataset = MedNISTDataset(image_files_list, labels=image_class)\n",
    "\n",
    "dataloader = ImageDataLoader(\n",
    "    dataset,\n",
    "    height=IMG_SIZE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ef1a423",
   "metadata": {},
   "source": [
    "## Load a generator - Conditional GAN\n",
    "\n",
    "For this experiment, we are using the `image_cgan` plugin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a95785a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# synthcity absolute\n",
    "from synthcity.plugins import Plugins\n",
    "\n",
    "generator = Plugins().get(\"image_cgan\", batch_size=100, plot_progress=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca77190b",
   "metadata": {},
   "source": [
    "## Train the generator\n",
    "\n",
    "For the training, we are using the `ImageDataLoader` object previously created.\n",
    "\n",
    "At the same time, we are using a conditional(`cond`) with the labels of the images. This way, at inference time, we can request from the generator only samples from a specific class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc6794cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "generator.fit(dataloader, cond=image_class)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6b6d5dc",
   "metadata": {},
   "source": [
    "## Generate new samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cec4b56",
   "metadata": {},
   "outputs": [],
   "source": [
    "# third party\n",
    "import torch\n",
    "\n",
    "# synthcity absolute\n",
    "from synthcity.plugins.core.models.image_gan import display_imgs\n",
    "\n",
    "syn_samples, syn_labels = generator.generate(count=5).unpack().tensors()\n",
    "\n",
    "display_imgs(syn_samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e2428fa",
   "metadata": {},
   "source": [
    "## Generate new samples using a conditional\n",
    "\n",
    "We can also generate instances from a specific class, using the conditional we used at training time(`cond`).\n",
    "\n",
    "__Disclaimer__ : Other architectures for the Generator and the Discriminator could improve the results. These architectures can be tweaked in the `suggest_image_generator_discriminator_arch` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b538951",
   "metadata": {},
   "outputs": [],
   "source": [
    "for cls_idx, cls in enumerate(class_names):\n",
    "    print(\"Class\", cls)\n",
    "    syn_samples, syn_labels = (\n",
    "        generator.generate(count=5, cond=np.ones(5) * cls_idx).unpack().tensors()\n",
    "    )\n",
    "\n",
    "    display_imgs(syn_samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62eb522b",
   "metadata": {},
   "source": [
    "## Benchmarks\n",
    "\n",
    "`synthcity` allows us to compare multiple generators on the same dataset, with a wide range of metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d895cdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# synthcity absolute\n",
    "from synthcity.benchmark import Benchmarks\n",
    "\n",
    "score = Benchmarks.evaluate(\n",
    "    [\n",
    "        (f\"test_{model}\", model, {}) for model in [\"image_cgan\", \"image_adsgan\"]\n",
    "    ],  # (testname, plugin, plugin_args) REPLACE {\"n_iter\" : 50} with {} for better performance\n",
    "    dataloader,\n",
    "    repeats=3,\n",
    "    metrics={\n",
    "        \"detection\": [\"detection_mlp\"],\n",
    "        \"performance\": [\"mlp\"],\n",
    "        \"stats\": [\"fid\"],\n",
    "    },\n",
    "    task_type=\"classification\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "243b9890",
   "metadata": {},
   "outputs": [],
   "source": [
    "Benchmarks.print(score)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "211a549e",
   "metadata": {},
   "source": [
    "## Congratulations!\n",
    "\n",
    "Congratulations on completing this notebook tutorial! If you enjoyed this and would like to join the movement towards Machine learning and AI for medicine, you can do so in the following ways!\n",
    "\n",
    "### Star [Synthcity](https://github.com/vanderschaarlab/synthcity) on GitHub\n",
    "\n",
    "- The easiest way to help our community is just by starring the Repos! This helps raise awareness of the tools we're building.\n",
    "\n",
    "\n",
    "### Checkout other projects from vanderschaarlab\n",
    "- [HyperImpute](https://github.com/vanderschaarlab/hyperimpute)\n",
    "- [AutoPrognosis](https://github.com/vanderschaarlab/autoprognosis)\n"
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}