notebooks/migrating_from_torchvision_to_albumentations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Migrating from torchvision to Albumentations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook shows how you can use Albumentations instead of torchvision to perform data augmentation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import the required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import cv2\n",
    "import numpy as np\n",
    "from torch.utils.data import Dataset\n",
    "from torchvision import transforms\n",
    "\n",
    "import albumentations as A\n",
    "from albumentations.pytorch import ToTensorV2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# An example pipeline that uses torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TorchvisionDataset(Dataset):\n",
    "    def __init__(self, file_paths, labels, transform=None):\n",
    "        self.file_paths = file_paths\n",
    "        self.labels = labels\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.file_paths)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        label = self.labels[idx]\n",
    "        file_path = self.file_paths[idx]\n",
    "\n",
    "        # Read an image with PIL\n",
    "        image = Image.open(file_path)\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "        return image, label\n",
    "\n",
    "\n",
    "torchvision_transform = transforms.Compose([\n",
    "    transforms.Resize((256, 256)),\n",
    "    transforms.RandomCrop(224),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(\n",
    "        mean=[0.485, 0.456, 0.406],\n",
    "        std=[0.229, 0.224, 0.225],\n",
    "    )\n",
    "])\n",
    "\n",
    "\n",
    "torchvision_dataset = TorchvisionDataset(\n",
    "    file_paths=['./images/image_1.jpg', './images/image_2.jpg', './images/image_3.jpg'],\n",
    "    labels=[1, 2, 3],\n",
    "    transform=torchvision_transform,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The same pipeline with Albumentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AlbumentationsDataset(Dataset):\n",
    "    \"\"\"__init__ and __len__ functions are the same as in TorchvisionDataset\"\"\"\n",
    "    def __init__(self, file_paths, labels, transform=None):\n",
    "        self.file_paths = file_paths\n",
    "        self.labels = labels\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.file_paths)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        label = self.labels[idx]\n",
    "        file_path = self.file_paths[idx]\n",
    "\n",
    "        # Read an image with OpenCV\n",
    "        image = cv2.imread(file_path)\n",
    "\n",
    "        # By default OpenCV uses BGR color space for color images,\n",
    "        # so we need to convert the image to RGB color space.\n",
    "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "        if self.transform:\n",
    "            augmented = self.transform(image=image)\n",
    "            image = augmented['image']\n",
    "        return image, label\n",
    "\n",
    "\n",
    "albumentations_transform = A.Compose([\n",
    "    A.Resize(256, 256),\n",
    "    A.RandomCrop(224, 224),\n",
    "    A.HorizontalFlip(),\n",
    "    A.Normalize(\n",
    "        mean=[0.485, 0.456, 0.406],\n",
    "        std=[0.229, 0.224, 0.225],\n",
    "    ),\n",
    "    ToTensorV2()\n",
    "])\n",
    "\n",
    "\n",
    "albumentations_dataset = AlbumentationsDataset(\n",
    "    file_paths=['./images/image_1.jpg', './images/image_2.jpg', './images/image_3.jpg'],\n",
    "    labels=[1, 2, 3],\n",
    "    transform=albumentations_transform,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using albumentations with PIL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use PIL instead of OpenCV while working with Albumentations, but in that case, you need to convert a PIL image to a NumPy array before applying transformations. Them you need to convert the augmented image back from a NumPy array to a PIL image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AlbumentationsPilDataset(Dataset):\n",
    "    \"\"\"__init__ and __len__ functions are the same as in TorchvisionDataset\"\"\"\n",
    "    def __init__(self, file_paths, labels, transform=None):\n",
    "        self.file_paths = file_paths\n",
    "        self.labels = labels\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.file_paths)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        label = self.labels[idx]\n",
    "        file_path = self.file_paths[idx]\n",
    "\n",
    "        image = Image.open(file_path)\n",
    "\n",
    "        if self.transform:\n",
    "            # Convert PIL image to numpy array\n",
    "            image_np = np.array(image)\n",
    "            # Apply transformations\n",
    "            augmented = self.transform(image=image_np)\n",
    "            # Convert numpy array to PIL Image\n",
    "            image = Image.fromarray(augmented['image'])\n",
    "        return image, label\n",
    "\n",
    "\n",
    "albumentations_pil_transform = A.Compose([\n",
    "    A.Resize(256, 256),\n",
    "    A.RandomCrop(224, 224),\n",
    "    A.HorizontalFlip(),\n",
    "])\n",
    "\n",
    "\n",
    "# Note that this dataset will output PIL images and not numpy arrays nor PyTorch tensors\n",
    "albumentations_pil_dataset = AlbumentationsPilDataset(\n",
    "    file_paths=['./images/image_1.jpg', './images/image_2.jpg', './images/image_3.jpg'],\n",
    "    labels=[1, 2, 3],\n",
    "    transform=albumentations_pil_transform,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Albumentations equivalents for torchvision transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| torchvision transform                                                                                                   | Albumentations transform                                                                                                                                            | Albumentations example                                                                                                                |\n",
    "|-------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------|\n",
    "| [Compose](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.Compose)                             | [Compose](https://albumentations.ai/docs/api_reference/core/composition/#albumentations.core.composition.Compose)                                                   | ```A.Compose([A.Resize(256, 256), A.RandomCrop(224, 224)])```                                                                         |\n",
    "| [CenterCrop](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.CenterCrop)                       | [CenterCrop](https://albumentations.ai/docs/api_reference/augmentations/crops/transforms/#albumentations.augmentations.crops.transforms.CenterCrop)                 | ```A.CenterCrop(256, 256)```                                                                                                          |\n",
    "| [ColorJitter](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.ColorJitter)                     | [HueSaturationValue](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.HueSaturationValue)             | ```A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.5)```                                         |\n",
    "| [Pad](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.Pad)                                     | [PadIfNeeded](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.PadIfNeeded)                           | ```A.PadIfNeeded(min_height=512, min_width=512)```                                                                                    |\n",
    "| [RandomAffine](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomAffine)                   | [Affine](https://albumentations.ai/docs/api_reference/augmentations/geometric/transforms/#albumentations.augmentations.geometric.transforms.Affine)                 | ```A.Affine(scale=(0.9, 1.1), translate_percent=(0.0, 0.2), rotate=(-45, 45), shear=(-15, 15), mode=cv2.BORDER_REFLECT_101, p=0.5)``` |\n",
    "| [RandomCrop](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomCrop)                       | [RandomCrop](https://albumentations.ai/docs/api_reference/augmentations/crops/transforms/#albumentations.augmentations.crops.transforms.RandomCrop)                 | ```A.RandomCrop(256, 256)```                                                                                                          |\n",
    "| [RandomGrayscale](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomGrayscale)             | [ToGray](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.ToGray)                                     | ```A.ToGray(p=0.5)```                                                                                                                 |\n",
    "| [RandomHorizontalFlip](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomHorizontalFlip)   | [HorizontalFlip](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.HorizontalFlip)                     | ```A.HorizontalFlip(p=0.5)```                                                                                                         |\n",
    "| [RandomPerspective](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomPerspective)         | [Perspective](https://albumentations.ai/docs/api_reference/augmentations/geometric/transforms/#albumentations.augmentations.geometric.transforms.Perspective)       | ```A.Perspective(scale=(0.2, 0.4), fit_output=True, p=0.5)```                                                                         |\n",
    "| [RandomRotation](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomRotation)               | [Rotate](https://albumentations.ai/docs/api_reference/augmentations/geometric/rotate/#albumentations.augmentations.geometric.rotate.Rotate)                         | ```A.Rotate(limit=45, p=0.5)```                                                                                                       |\n",
    "| [RandomVerticalFlip](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomVerticalFlip)       | [VerticalFlip](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.VerticalFlip)                         | ```A.VerticalFlip(p=0.5)```                                                                                                           |\n",
    "| [Resize](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.Resize)                               | [Resize](https://albumentations.ai/docs/api_reference/augmentations/geometric/resize/#albumentations.augmentations.geometric.resize.Resize)                         | ```A.Resize(256, 256)```                                                                                                              |\n",
    "| [GaussianBlur](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.GaussianBlur)                   | [GaussianBlur](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.GaussianBlur)                         | ```A.GaussianBlur(blur_limit=(3, 7), p=0.5)```                                                                                        |\n",
    "| [RandomInvert](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomInvert)                   | [InvertImg](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.InvertImg)                               | ```A.InvertImg(p=0.5)```                                                                                                              |\n",
    "| [RandomPosterize](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomPosterize)             | [Posterize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Posterize)                               | ```A.Posterize(num_bits=4, p=0.5)```                                                                                                  |\n",
    "| [RandomSolarize](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomSolarize)               | [Solarize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Solarize)                                 | ```A.Solarize(threshold=127, p=0.5)```                                                                                                |\n",
    "| [RandomAdjustSharpness](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomAdjustSharpness) | [Sharpen](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Sharpen)                                   | ```A.Sharpen(alpha=(0.2, 0.5), lightness=(0.5, 1.0), p=0.5)```                                                                        |\n",
    "| [RandomAutocontrast](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomAutocontrast)       | [RandomBrightnessContrast](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.RandomBrightnessContrast) | ```A.RandomBrightnessContrast(brightness_limit=0, contrast_limit=0.2, p=0.5)```                                                       |\n",
    "| [RandomEqualize](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomEqualize)               | [Equalize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Equalize)                                 | ```A.Equalize(p=0.5)```                                                                                                               |\n",
    "| [RandomErasing](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.RandomErasing)                 | [CoarseDropout](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.CoarseDropout)                       | ```A.CoarseDropout(min_height=8, max_height=32, min_width=8, max_width=32, p=0.5)```                                                  |\n",
    "| [Normalize](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.Normalize)                         | [Normalize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize)                               | ```A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])```                                                              |"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}