code/utils/vision-transform.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📝 **Author:** Amirhossein Heydari - 📧 **Email:** <amirhosseinheydari78@gmail.com> - 📍 **Origin:** [mr-pylin/pytorch-workshop](https://github.com/mr-pylin/pytorch-workshop)\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Table of contents**<a id='toc0_'></a>    \n",
    "- [Dependencies](#toc1_)    \n",
    "- [Load CIFAR-10 Dataset](#toc2_)    \n",
    "- [Transforms](#toc3_)    \n",
    "  - [Built-in Transforms](#toc3_1_)    \n",
    "    - [Geometry](#toc3_1_1_)    \n",
    "      - [Resize](#toc3_1_1_1_)    \n",
    "      - [Cropping](#toc3_1_1_2_)    \n",
    "      - [Others](#toc3_1_1_3_)    \n",
    "    - [Color](#toc3_1_2_)    \n",
    "    - [Composition](#toc3_1_3_)    \n",
    "    - [Miscellaneous](#toc3_1_4_)    \n",
    "    - [Conversion](#toc3_1_5_)    \n",
    "    - [Auto-Augmentation](#toc3_1_6_)    \n",
    "  - [Custom Transforms](#toc3_2_)    \n",
    "    - [Approach 1: Using nn.Module](#toc3_2_1_)    \n",
    "    - [Approach 2: Using v2.Transform](#toc3_2_2_)    \n",
    "  - [A Typical Transform Pipeline](#toc3_3_)    \n",
    "\n",
    "<!-- vscode-jupyter-toc-config\n",
    "\tnumbering=false\n",
    "\tanchor=true\n",
    "\tflat=false\n",
    "\tminLevel=1\n",
    "\tmaxLevel=6\n",
    "\t/vscode-jupyter-toc-config -->\n",
    "<!-- THIS CELL WILL BE REPLACED ON TOC UPDATE. DO NOT WRITE YOUR TEXT IN THIS CELL -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc1_'></a>[Dependencies](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.datasets import CIFAR10\n",
    "from torchvision.transforms import v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set a seed for deterministic results\n",
    "seed = 42\n",
    "torch.manual_seed(seed)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# update paths as needed based on your project structure\n",
    "DATASET_DIR = r\"../../datasets\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot(x1: list[torch.Tensor], x2: list[torch.Tensor], y: list[int], transform: v2.Transform) -> None:\n",
    "    fig, axs = plt.subplots(nrows=2, ncols=5, figsize=(16, 8), layout=\"compressed\")\n",
    "    for i, (img1, img2, label) in enumerate(zip(x1, x2, y)):\n",
    "        axs[0, i].imshow(img1.permute(1, 2, 0))\n",
    "        axs[0, i].set(title=\"Original\")\n",
    "        axs[1, i].imshow(img2.permute(1, 2, 0))\n",
    "        axs[1, i].set(title=str(transform).split(\"(\")[0])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc2_'></a>[Load CIFAR-10 Dataset](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainset = CIFAR10(DATASET_DIR, train=True, transform=v2.ToImage(), download=False)\n",
    "\n",
    "x = [trainset[i][0] for i in range(5)]\n",
    "y = [trainset[i][1] for i in range(5)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot\n",
    "fig, axs = plt.subplots(nrows=1, ncols=5, figsize=(16, 4), layout=\"compressed\")\n",
    "for i, (img, label) in enumerate(zip(x, y)):\n",
    "    axs[i].imshow(img.permute(1, 2, 0))\n",
    "    axs[i].set(title=label)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc3_'></a>[Transforms](#toc0_)\n",
    "\n",
    "- Torchvision supports common **computer vision** transformations in the `torchvision.transforms` and `torchvision.transforms.v2` modules.\n",
    "- Transforms can be used to **transform** or **augment** data for **training** or **inference** of different tasks (image classification, detection, segmentation, video classification).\n",
    "\n",
    "📝 **Docs**:\n",
    "\n",
    "- Transforming and augmenting images: [pytorch.org/vision/stable/transforms.html](https://pytorch.org/vision/stable/transforms.html)\n",
    "- Illustration of transforms: [pytorch.org/vision/stable/auto_examples/transforms/plot_transforms_illustrations.html](https://pytorch.org/vision/stable/auto_examples/transforms/plot_transforms_illustrations.html)\n",
    "\n",
    "✍️ **Performance considerations**:\n",
    "\n",
    "- Rely on the v2 transforms from `torchvision.transforms.v2`.\n",
    "- Use tensors instead of `PIL` images.\n",
    "- Use `torch.uint8` dtype, especially for resizing.\n",
    "- Resize with **bilinear** or **bicubic** mode.\n",
    "- Consider `num_workers > 0` when creating `torch.utils.data.DataLoader`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc3_1_'></a>[Built-in Transforms](#toc0_)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc3_1_1_'></a>[Geometry](#toc0_)\n",
    "\n",
    "<table style=\"margin: 0 auto;\">\n",
    "  <thead>\n",
    "    <tr>\n",
    "      <th>Resizing</th>\n",
    "      <th>Cropping</th>\n",
    "      <th>Other</th>\n",
    "    </tr>\n",
    "  </thead>\n",
    "  <tbody style=\"font-family: monospace\">\n",
    "    <tr>\n",
    "      <td>v2.Resize()</td>\n",
    "      <td>v2.RandomCrop()</td>\n",
    "      <td>v2.RandomHorizontalFlip()</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>v2.RandomResize()</td>\n",
    "      <td>v2.RandomResizedCrop()</td>\n",
    "      <td>v2.RandomVerticalFlip()</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td></td>\n",
    "      <td>v2.RandomIoUCrop()</td>\n",
    "      <td>v2.Pad()</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td></td>\n",
    "      <td>v2.CenterCrop()</td>\n",
    "      <td>v2.RandomZoomOut()</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td></td>\n",
    "      <td></td>\n",
    "      <td>v2.RandomRotation()</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td></td>\n",
    "      <td></td>\n",
    "      <td>v2.RandomAffine()</td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <a id='toc3_1_1_1_'></a>[Resize](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_resize = v2.Resize(size=(16, 16), interpolation=v2.InterpolationMode.NEAREST)\n",
    "x_resize = t_resize(x)\n",
    "\n",
    "# plot\n",
    "plot(x, x_resize, y, t_resize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_random_resize = v2.RandomResize(min_size=8, max_size=64)\n",
    "x_random_resize = [t_random_resize(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_random_resize, y, t_random_resize)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <a id='toc3_1_1_2_'></a>[Cropping](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_random_crop = v2.RandomCrop(size=(24, 24))\n",
    "x_random_crop = [t_random_crop(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_random_crop, y, t_random_crop)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_center_crop = v2.CenterCrop(size=(16, 16))\n",
    "x_center_crop = t_center_crop(x)\n",
    "\n",
    "# plot\n",
    "plot(x, x_center_crop, y, t_center_crop)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <a id='toc3_1_1_3_'></a>[Others](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_random_horizontal_flip = v2.RandomHorizontalFlip()\n",
    "x_random_horizontal_flip = [t_random_horizontal_flip(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_random_horizontal_flip, y, t_random_horizontal_flip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_pad = v2.Pad(padding=(1, 2, 3, 4), fill=0, padding_mode=\"constant\")\n",
    "x_pad = t_pad(x)\n",
    "\n",
    "# plot\n",
    "plot(x, x_pad, y, t_pad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_random_zoomout = v2.RandomZoomOut(fill=0)\n",
    "x_random_zoomout = [t_random_zoomout(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_random_zoomout, y, t_random_zoomout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_random_rotation = v2.RandomRotation(degrees=45, expand=True)\n",
    "x_random_rotation = [t_random_rotation(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_random_rotation, y, t_random_rotation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_random_affine = v2.RandomAffine(degrees=45, translate=(0.1, 0.1), scale=(0.8, 1.2), shear=2)\n",
    "x_random_affine = [t_random_affine(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_random_affine, y, t_random_affine)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc3_1_2_'></a>[Color](#toc0_)\n",
    "\n",
    "<table style=\"margin: 0 auto;\">\n",
    "  <tbody style=\"font-family: monospace\">\n",
    "    <tr>\n",
    "      <td>v2.ColorJitter()</td>\n",
    "      <td>v2.Grayscale()</td>\n",
    "      <td>v2.RGB()</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>v2.RandomGrayscale()</td>\n",
    "      <td>v2.GaussianBlur()</td>\n",
    "      <td>v2.GaussianNoise()</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>v2.RandomInvert()</td>\n",
    "      <td>v2.RandomPosterize()</td>\n",
    "      <td>v2.RandomSolarize()</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>v2.RandomEqualize()</td>\n",
    "      <td></td>\n",
    "      <td></td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_color_jitter = v2.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5)\n",
    "x_color_jitter = [t_color_jitter(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_color_jitter, y, t_color_jitter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_gray_scale = v2.Grayscale(num_output_channels=3)\n",
    "x_gray_scale = t_gray_scale(x)\n",
    "\n",
    "# plot\n",
    "plot(x, x_gray_scale, y, t_gray_scale)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_random_solarize = v2.RandomSolarize(threshold=100)\n",
    "x_random_solarize = [t_random_solarize(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_random_solarize, y, t_random_solarize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_gaussian_noise = v2.GaussianNoise(mean=0, sigma=0.1)\n",
    "x_gaussian_noise = [t_gaussian_noise(img.to(torch.float32) / 255) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_gaussian_noise, y, t_gaussian_noise)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc3_1_3_'></a>[Composition](#toc0_)\n",
    "\n",
    "<table style=\"margin: 0 auto;\">\n",
    "  <thead>\n",
    "    <tr>\n",
    "      <th>Transform</th>\n",
    "      <th>Always Applied?</th>\n",
    "      <th>Applies All?</th>\n",
    "      <th>Order Fixed?</th>\n",
    "      <th>Use Case</th>\n",
    "    </tr>\n",
    "  </thead>\n",
    "  <tbody>\n",
    "    <tr>\n",
    "      <td style=\"font-family: monospace;\">v2.Compose()</td>\n",
    "      <td>✅ Yes</td>\n",
    "      <td>✅ Yes</td>\n",
    "      <td>✅ Yes</td>\n",
    "      <td>Fixed pipeline of transformations</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td style=\"font-family: monospace;\">v2.RandomApply()</td>\n",
    "      <td>❌ No (probabilistic)</td>\n",
    "      <td>✅ Yes (if applied)</td>\n",
    "      <td>✅ Yes</td>\n",
    "      <td>Conditional application</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td style=\"font-family: monospace;\">v2.RandomChoice()</td>\n",
    "      <td>✅ Yes</td>\n",
    "      <td>❌ No (only one)</td>\n",
    "      <td>✅ Yes</td>\n",
    "      <td>Random selection of one transform</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td style=\"font-family: monospace;\">v2.RandomOrder()</td>\n",
    "      <td>✅ Yes</td>\n",
    "      <td>✅ Yes</td>\n",
    "      <td>❌ No</td>\n",
    "      <td>Random order of all transforms</td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_compose = v2.Compose(\n",
    "    [\n",
    "        v2.Resize((64, 64)),\n",
    "        v2.RandomHorizontalFlip(p=0.5),\n",
    "        v2.ColorJitter(brightness=0.2, contrast=0.2),\n",
    "    ]\n",
    ")\n",
    "\n",
    "x_compose = [t_compose(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_compose, y, t_compose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_random_apply = v2.RandomApply(\n",
    "    [\n",
    "        v2.Resize((64, 64)),\n",
    "        v2.Grayscale(num_output_channels=3),\n",
    "    ],\n",
    "    p=0.5,\n",
    ")\n",
    "\n",
    "x_random_apply = [t_random_apply(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_random_apply, y, t_random_apply)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_random_choice = v2.RandomChoice(\n",
    "    [\n",
    "        v2.RandomHorizontalFlip(p=1.0),\n",
    "        v2.RandomVerticalFlip(p=1.0),\n",
    "    ],\n",
    ")\n",
    "\n",
    "x_random_choice = [t_random_choice(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_random_choice, y, t_random_choice)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_random_order = v2.RandomOrder(\n",
    "    [\n",
    "        v2.RandomRotation(degrees=30),\n",
    "        v2.RandomHorizontalFlip(p=1.0),\n",
    "        v2.ColorJitter(0.2, 0.2, 0.2, 0.2),\n",
    "    ],\n",
    ")\n",
    "\n",
    "x_random_order = [t_random_order(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_random_order, y, t_random_order)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# combining compositions\n",
    "t_combine = v2.Compose(\n",
    "    [\n",
    "        v2.Resize((224, 224)),\n",
    "        v2.RandomApply([v2.ColorJitter(brightness=0.5, contrast=0.5)], p=0.3),\n",
    "        v2.RandomChoice([v2.GaussianBlur(kernel_size=3), v2.RandomRotation(degrees=30)]),\n",
    "        v2.RandomOrder(\n",
    "            [v2.RandomHorizontalFlip(p=0.5), v2.RandomVerticalFlip(p=0.5), v2.Grayscale(num_output_channels=3)]\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "\n",
    "x_combine = [t_combine(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_combine, y, t_combine)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc3_1_4_'></a>[Miscellaneous](#toc0_)\n",
    "\n",
    "<table style=\"margin: 0 auto;\">\n",
    "  <tbody style=\"font-family: monospace\">\n",
    "    <tr>\n",
    "      <td>v2.Normalize()</td>\n",
    "      <td>v2.RandomErasing()</td>\n",
    "      <td>v2.Lambda()</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>v2.JPEG()</td>\n",
    "      <td></td>\n",
    "      <td></td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_normalize = v2.Normalize(mean=(0.4914, 0.4822, 0.4465), std=(0.2470, 0.2435, 0.2616))\n",
    "x_normalize = t_normalize([img.to(torch.float32) / 255 for img in x])\n",
    "\n",
    "# plot\n",
    "plot(x, x_normalize, y, t_normalize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_random_erasing = v2.RandomErasing(p=0.7)\n",
    "x_random_erasing = [t_random_erasing(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_random_erasing, y, t_random_erasing)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc3_1_5_'></a>[Conversion](#toc0_)\n",
    "\n",
    "<table style=\"margin: 0 auto;\">\n",
    "  <tbody style=\"font-family: monospace\">\n",
    "    <tr>\n",
    "      <td>v2.ToImage()</td>\n",
    "      <td>v2.ToDtype()</td>\n",
    "      <td>v2.ToPILImage()</td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n",
    "\n",
    "✍️ **Note**:\n",
    "\n",
    "- These transforms are **deprecated**:\n",
    "  - `v2.ToTensor()`\n",
    "  - `v2.ConvertImageDtype()`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_to_image = v2.ToImage()\n",
    "x_to_image = t_to_image(x)\n",
    "\n",
    "# plot\n",
    "plot(x, x_to_image, y, t_to_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_to_dtype = v2.ToDtype(torch.float32, scale=True)\n",
    "x_to_dtype = t_to_dtype(x)\n",
    "\n",
    "# plot\n",
    "plot(x, x_to_dtype, y, t_to_dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc3_1_6_'></a>[Auto-Augmentation](#toc0_)\n",
    "\n",
    "<table style=\"margin: 0 auto\">\n",
    "  <thead>\n",
    "    <tr>\n",
    "      <th>Method</th>\n",
    "      <th>Selection</th>\n",
    "      <th>Applied Transforms</th>\n",
    "      <th>Strength Control</th>\n",
    "      <th>Blending</th>\n",
    "    </tr>\n",
    "  </thead>\n",
    "  <tbody>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">v2.AutoAugment()</span></td>\n",
    "      <td>Predefined policies</td>\n",
    "      <td>Fixed sequence</td>\n",
    "      <td>❌</td>\n",
    "      <td>❌</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">v2.RandAugment()</span></td>\n",
    "      <td>Random (<span style=\"font-family: monospace;\">num_ops</span>)</td>\n",
    "      <td>Sequential</td>\n",
    "      <td>✅ (<span style=\"font-family: monospace;\">magnitude</span>)</td>\n",
    "      <td>❌</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">v2.TrivialAugmentWide()</span></td>\n",
    "      <td>Random (1 transform)</td>\n",
    "      <td>Single transform</td>\n",
    "      <td>❌</td>\n",
    "      <td>❌</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">v2.AugMix()</span></td>\n",
    "      <td>Random (multiple)</td>\n",
    "      <td>Mixed (random order)</td>\n",
    "      <td>✅ (<span style=\"font-family: monospace;\">severity</span>)</td>\n",
    "      <td>✅ (weighted mix)</td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_auto_augment = v2.AutoAugment(policy=v2.AutoAugmentPolicy.CIFAR10)\n",
    "x_auto_augment = [t_auto_augment(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_auto_augment, y, t_auto_augment)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_rand_augment = v2.RandAugment(num_ops=2, magnitude=9)\n",
    "x_rand_augment = [t_rand_augment(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_rand_augment, y, t_rand_augment)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_trivial_augment_wide = v2.TrivialAugmentWide()\n",
    "x_trivial_augment_wide = [t_trivial_augment_wide(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_trivial_augment_wide, y, t_trivial_augment_wide)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_aug_mix = v2.AugMix(severity=10, mixture_width=10, chain_depth=-1)\n",
    "x_aug_mix = [t_aug_mix(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_aug_mix, y, t_aug_mix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc3_2_'></a>[Custom Transforms](#toc0_)\n",
    "\n",
    "- You can define custom transforms in PyTorch using either `torch.nn.Module` or `torchvision.transforms.v2.Transform`.\n",
    "- To create a custom transform:\n",
    "  - Extend `torch.nn.Module` and implement the `forward` method for **simple** transforms.\n",
    "  - Extend `torchvision.transforms.v2.Transform` and implement the `_transform` method for **advanced** transforms that support **arbitrary** input structures (e.g., **images**, **bounding boxes**, **segmentation masks**).\n",
    "\n",
    "📝 **Docs**:\n",
    "\n",
    "- `nn.Module`: [pytorch.org/docs/stable/generated/torch.nn.Module.html](https://pytorch.org/docs/stable/generated/torch.nn.Module.html)\n",
    "- `v2.Transform`: [pytorch.org/vision/stable/generated/torchvision.transforms.v2.Transform.html](https://pytorch.org/vision/stable/generated/torchvision.transforms.v2.Transform.html)\n",
    "- How to write your own v2 transforms: [pytorch.org/vision/stable/auto_examples/transforms/plot_custom_transforms.html](https://pytorch.org/vision/stable/auto_examples/transforms/plot_custom_transforms.html)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc3_2_1_'></a>[Approach 1: Using nn.Module](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomRandomColorInversion1(nn.Module):\n",
    "    def __init__(self, p=0.5):\n",
    "        super().__init__()\n",
    "        self.p = p  # probability of applying the transform\n",
    "\n",
    "    def forward(self, img: torch.Tensor) -> torch.Tensor:\n",
    "        # check for unsupported dtypes\n",
    "        if img.dtype not in (torch.uint8, torch.float16, torch.float32, torch.float64):\n",
    "            raise ValueError(f\"Unsupported dtype: {img.dtype}. Expected uint8 or float.\")\n",
    "\n",
    "        # check if float image is in range [0, 1]\n",
    "        if img.dtype.is_floating_point:\n",
    "            if img.min() < 0 or img.max() > 1:\n",
    "                raise ValueError(f\"Float image must be in range [0, 1]. Found range [{img.min()}, {img.max()}].\")\n",
    "\n",
    "        # Apply the transform\n",
    "        if torch.rand(1) < self.p:\n",
    "            if img.dtype == torch.uint8:\n",
    "                img = 255 - img  # invert for uint8\n",
    "            else:\n",
    "                img = 1.0 - img  # invert for float\n",
    "\n",
    "        return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_custom_random_color_inversion_1 = CustomRandomColorInversion1()\n",
    "x_custom_random_color_inversion_1 = [t_custom_random_color_inversion_1(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_custom_random_color_inversion_1, y, t_custom_random_color_inversion_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc3_2_2_'></a>[Approach 2: Using v2.Transform](#toc0_)\n",
    "\n",
    "- The `_transform` method takes two arguments:\n",
    "  - `inpt`: The input to transform (e.g., an **image** tensor).\n",
    "  - `params`: A **dictionary** of **parameters** generated by the `_get_params` method (if implemented). This is useful for transforms that require dynamic or random parameters.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomRandomColorInversion2(v2.Transform):\n",
    "    def __init__(self, p=0.5):\n",
    "        super().__init__()\n",
    "        self.p = p  # probability of applying the transform\n",
    "\n",
    "    def _transform(self, img: torch.Tensor, params: dict) -> torch.Tensor:\n",
    "        # check for unsupported dtypes\n",
    "        if img.dtype not in (torch.uint8, torch.float16, torch.float32, torch.float64):\n",
    "            raise ValueError(f\"Unsupported dtype: {img.dtype}. Expected uint8 or float.\")\n",
    "\n",
    "        # check if float image is in range [0, 1]\n",
    "        if img.dtype.is_floating_point:\n",
    "            if img.min() < 0 or img.max() > 1:\n",
    "                raise ValueError(f\"Float image must be in range [0, 1]. Found range [{img.min()}, {img.max()}].\")\n",
    "\n",
    "        # Apply the transform\n",
    "        if torch.rand(1) < self.p:\n",
    "            if img.dtype == torch.uint8:\n",
    "                img = 255 - img  # invert for uint8\n",
    "            else:\n",
    "                img = 1.0 - img  # invert for float\n",
    "\n",
    "        return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_custom_random_color_inversion_2 = CustomRandomColorInversion2()\n",
    "x_custom_random_color_inversion_2 = [t_custom_random_color_inversion_2(img) for img in x]\n",
    "\n",
    "# plot\n",
    "plot(x, x_custom_random_color_inversion_2, y, t_custom_random_color_inversion_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc3_3_'></a>[A Typical Transform Pipeline](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_typical = v2.Compose(\n",
    "    [\n",
    "        # 1. convert to Tensor (if input is a PIL image or numpy ndarray)\n",
    "        # only needed if your input is not already a tensor\n",
    "        v2.ToImage(),\n",
    "        # 2. ensure the image is in uint8 format (optional)\n",
    "        # converts to uint8 and scales values to [0, 255]\n",
    "        # most inputs are already uint8 at this point, so this is optional\n",
    "        v2.ToDtype(torch.uint8, scale=True),\n",
    "        # 3. Data augmentation: Randomly resize and crop the image\n",
    "        # Randomly crops and resizes the image to (224, 224)\n",
    "        # `antialias=True` improves resizing quality (increases computations)\n",
    "        v2.RandomResizedCrop(size=(224, 224), antialias=True),\n",
    "        # 4. Data augmentation: Randomly flip the image horizontally\n",
    "        # Flips the image horizontally with a probability of 0.5\n",
    "        v2.RandomHorizontalFlip(p=0.5),\n",
    "        # 5. Convert to float32 and scale to [0, 1]\n",
    "        # Converts to float32 and scales values to [0, 1]\n",
    "        # Required for normalization\n",
    "        v2.ToDtype(torch.float32, scale=True),\n",
    "        # 6. Normalize the image using mean and std\n",
    "        # Normalizes the image\n",
    "        # These values are standard for ImageNet\n",
    "        v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainset = CIFAR10(DATASET_DIR, train=True, transform=t_typical, download=False)\n",
    "train_loader = DataLoader(trainset, batch_size=5)\n",
    "\n",
    "# plot\n",
    "for x_transformed, y in train_loader:\n",
    "    plot(x, x_transformed, y, \"t_typical\")\n",
    "    break"
   ]
  }
 ],
 "metadata": {
  "author_email": "AmirhosseinHeydari78@gmail.com",
  "author_github": "https://github.com/mr-pylin",
  "author_name": "Amirhossein Heydari",
  "kernelspec": {
   "display_name": "pytorch-workshop-U_zYfVTd-py3.12",
   "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.12.8"
  },
  "origin_repo": "https://github.com/mr-pylin/pytorch-workshop"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}