12-pytorch-02-dataloaders.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---\n",
        "title: PyTorch 02 - Dataloaders and Transforms\n",
        "jupyter: python3\n",
        "---\n",
        "\n",
        "## Datasets & DataLoaders\n",
        "\n",
        "[![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/trgardos/ml-549-fa24/blob/main/12-pytorch-02-dataloaders.ipynb)\n",
        "\n",
        "::: {.callout-note}\n",
        "Adapted from [PyTorch Quickstart](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html)\n",
        "to use CIFAR10 dataset instead of FashionMNIST.\n",
        ":::\n",
        "\n",
        "* Code for processing data samples can get messy and hard to maintain\n",
        "* We ideally want our dataset code to be decoupled from our model training\n",
        "code for better readability and modularity. \n",
        "\n",
        "## PyTorch Data Primitives\n",
        "\n",
        "PyTorch provides two data primitives: \n",
        "\n",
        "1. `torch.utils.data.DataLoader` and \n",
        "2. `torch.utils.data.Dataset`\n",
        "\n",
        "that allow you to use pre-loaded datasets as well as your own data.\n",
        "\n",
        "* `Dataset` stores the samples and their corresponding labels, and\n",
        "* `DataLoader` wraps an _iterable_ around the `Dataset` to enable easy\n",
        "   access to the samples.\n",
        "\n",
        "::: {.callout-note}\n",
        "An **iterable** is a Python object capable of returning its members one at a time.\n",
        "It must implement the `__iter__` method or the `__getitem__` method. See\n",
        "[Iterators](https://docs.python.org/3/tutorial/classes.html#iterators)\n",
        "for more details.\n",
        ":::\n",
        "\n",
        "## Pre-loaded Datasets\n",
        "\n",
        "PyTorch provides a number of pre-loaded datasets (such as FashionMNIST or CIFAR10)\n",
        "that subclass `torch.utils.data.Dataset` and implement\n",
        "functions specific to the particular data. \n",
        "\n",
        "They can be used to prototype and benchmark your model. \n",
        "\n",
        "You can find them here:\n",
        "\n",
        "* [Image Datasets](https://pytorch.org/vision/stable/datasets.html)\n",
        "    <details>\n",
        "    <summary>View sublist</summary>\n",
        "    <ul>\n",
        "        <li>image classification</li>\n",
        "        <li>object detection or segmentation</li>\n",
        "        <li>optical flow</li>\n",
        "        <li>stereo matching</li>\n",
        "        <li>image pairs</li>\n",
        "        <li>image captioning</li>\n",
        "        <li>video classification</li>\n",
        "        <li>video prediction</li>\n",
        "    </ul>\n",
        "    </details>\n",
        "* [Text Datasets](https://pytorch.org/text/stable/datasets.html)\n",
        "    <details>\n",
        "    <summary>View sublist</summary>\n",
        "    <ul>\n",
        "        <li>text classification</li>\n",
        "        <li>language modeling</li>\n",
        "        <li>machine translation</li>\n",
        "        <li>sequence tagging</li>\n",
        "        <li>question answering</li>\n",
        "        <li>unsupervised learning</li>\n",
        "    </ul>\n",
        "    </details>\n",
        "* [Audio Datasets](https://pytorch.org/audio/stable/datasets.html)\n",
        "\n",
        "\n",
        "## Loading a Dataset\n",
        "\n",
        "Here is an example of how to load the\n",
        "[CIFAR10](https://pytorch.org/vision/stable/generated/torchvision.datasets.CIFAR10.html#torchvision.datasets.CIFAR10)\n",
        "dataset from TorchVision.\n",
        "\n",
        "We load the\n",
        "[CIFAR10](https://pytorch.org/vision/stable/generated/torchvision.datasets.CIFAR10.html#torchvision.datasets.CIFAR10) \n",
        "with the following parameters:\n",
        "\n",
        "- `root` is the path where the train/test data is stored,\n",
        "- `train` if True, specifies training dataset, if False, specifies test dataset,\n",
        "- `download=True` downloads the data from the internet if it's\n",
        "   not available at `root`.\n",
        "- `transform` and `target_transform` specify the feature and label\n",
        "   transformations. More on that later."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch.utils.data import Dataset\n",
        "from torchvision import datasets\n",
        "from torchvision.transforms import ToTensor\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "training_data = datasets.CIFAR10(\n",
        "    root=\"data\",\n",
        "    train=True,\n",
        "    download=True,\n",
        "    transform=ToTensor()\n",
        ")\n",
        "\n",
        "test_data = datasets.CIFAR10(\n",
        "    root=\"data\",\n",
        "    train=False,\n",
        "    download=True,\n",
        "    transform=ToTensor()\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In this case we simply use the\n",
        "[ToTensor](https://pytorch.org/vision/main/generated/torchvision.transforms.ToTensor.html)\n",
        "transform which converts the image from\n",
        "a (H x W x C) shape to a (C x H x W) shape and converts the pixel values from\n",
        "[0,255] to a torch.FloatTensor in the range [0.0, 1.0]."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "img, label = training_data[0]\n",
        "print(f\"img.shape: {img.shape}\")\n",
        "print(f\"img.dtype: {img.dtype}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's also look at the `data` directory contents.\n",
        "\n",
        "```sh \n",
        "% ls -al data\n",
        "total 333008\n",
        "drwxr-xr-x@  6 tomg  staff        192 Oct  1 22:14 .\n",
        "drwxr-xr-x  62 tomg  staff       1984 Oct 17 11:09 ..\n",
        "drwxr-xr-x@ 10 tomg  staff        320 Jun  4  2009 cifar-10-batches-py\n",
        "-rw-r--r--@  1 tomg  staff  170498071 Sep  4 19:51 cifar-10-python.tar.gz\n",
        "```\n",
        "\n",
        "And the contents of the `cifar-10-batches-py` directory.\n",
        "\n",
        "```sh\n",
        "% ls -al data/cifar-10-batches-py \n",
        "total 363752\n",
        "drwxr-xr-x@ 10 tomg  staff       320 Jun  4  2009 .\n",
        "drwxr-xr-x@  6 tomg  staff       192 Oct  1 22:14 ..\n",
        "-rw-r--r--@  1 tomg  staff       158 Mar 31  2009 batches.meta\n",
        "-rw-r--r--@  1 tomg  staff  31035704 Mar 31  2009 data_batch_1\n",
        "-rw-r--r--@  1 tomg  staff  31035320 Mar 31  2009 data_batch_2\n",
        "-rw-r--r--@  1 tomg  staff  31035999 Mar 31  2009 data_batch_3\n",
        "-rw-r--r--@  1 tomg  staff  31035696 Mar 31  2009 data_batch_4\n",
        "-rw-r--r--@  1 tomg  staff  31035623 Mar 31  2009 data_batch_5\n",
        "-rw-r--r--@  1 tomg  staff        88 Jun  4  2009 readme.html\n",
        "-rw-r--r--@  1 tomg  staff  31035526 Mar 31  2009 test_batch\n",
        "```\n",
        "\n",
        "You see in this case the images aren't stored individually, but rather\n",
        "as combined batches. The CIFAR10 `Dataset` and `DataLoader` classes hide\n",
        "these details from us.\n",
        "\n",
        "### Iterating and Visualizing the Dataset\n",
        "\n",
        "It's very important to understand the data you're working with. Visually inspecting\n",
        "the dataset is a good way to get started.\n",
        "\n",
        "We can index `Datasets` manually like a list: `training_data[index]`. In this\n",
        "case we randomly sample images from the dataset.\n",
        "\n",
        "We use `matplotlib` to visualize some samples in our training data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "labels_map = {\n",
        "    0: \"plane\",\n",
        "    1: \"car\",\n",
        "    2: \"bird\",\n",
        "    3: \"cat\",\n",
        "    4: \"deer\",\n",
        "    5: \"dog\",\n",
        "    6: \"frog\",\n",
        "    7: \"horse\",\n",
        "    8: \"ship\",\n",
        "    9: \"truck\",\n",
        "}\n",
        "\n",
        "figure = plt.figure(figsize=(6, 6))\n",
        "cols, rows = 4, 4\n",
        "\n",
        "for i in range(1, cols * rows + 1):\n",
        "    # Randomly choose indices\n",
        "    sample_idx = torch.randint(len(training_data), size=(1,)).item()\n",
        "\n",
        "    img, label = training_data[sample_idx]\n",
        "    #print(f\"img.shape: {img.shape}\")\n",
        "    #print(f\"label: {label}\")\n",
        "    \n",
        "    figure.add_subplot(rows, cols, i)\n",
        "    plt.title(labels_map[label])\n",
        "    plt.axis(\"off\")\n",
        "    plt.imshow(img.permute(1, 2, 0))\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "::: {.callout-tip}\n",
        "Try re-running the above cell a few times to see different samples from the dataset.\n",
        ":::\n",
        "\n",
        "\n",
        "### Collecting Sample Data to Illustrate Custom Dataset\n",
        "\n",
        "To illustrate creating a custom dataset, we will collect images from the CIFAR10\n",
        "dataset and save them to a local directory. We will also save the labels to a CSV\n",
        "file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import pandas as pd\n",
        "from torchvision import datasets\n",
        "from torchvision.transforms import ToTensor\n",
        "from PIL import Image\n",
        "\n",
        "# Create directories to store images and annotations\n",
        "os.makedirs('cifar10_images', exist_ok=True)\n",
        "annotations_file = 'cifar10_annotations.csv'\n",
        "\n",
        "# Load CIFAR10 dataset\n",
        "cifar10 = datasets.CIFAR10(root='data', train=True, download=True, transform=ToTensor())\n",
        "\n",
        "# Number of images to download\n",
        "n_images = 10\n",
        "\n",
        "# Store images and their labels\n",
        "data = []\n",
        "for i in range(n_images):\n",
        "    img, label = cifar10[i]\n",
        "    img = img.permute(1, 2, 0)  # Convert from (C, H, W) to (H, W, C)\n",
        "    img = (img * 255).byte().numpy()  # Convert to numpy array and scale to [0, 255]\n",
        "    img = Image.fromarray(img)  # Convert to PIL Image\n",
        "\n",
        "    img_filename = f'cifar10_images/img_{i}.png'\n",
        "    img.save(img_filename)  # Save image\n",
        "\n",
        "    data.append([img_filename, label])  # Append image path and label to data list\n",
        "\n",
        "# Write annotations to CSV file\n",
        "df = pd.DataFrame(data, columns=['image_path', 'label'])\n",
        "df.to_csv(annotations_file, index=False)\n",
        "\n",
        "print(f\"Saved {n_images} images and their labels to {annotations_file}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# List the directory `cifar_images\n",
        "import os\n",
        "os.listdir('cifar10_images')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "\n",
        "# Read the annotations file\n",
        "annotations = pd.read_csv('cifar10_annotations.csv')\n",
        "\n",
        "# Display the first 10 lines of the annotations file\n",
        "print(annotations.head(10))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Creating a Custom Dataset Class\n",
        "\n",
        "A custom Dataset class must implement three functions:\n",
        "\n",
        "1. `__init__`, \n",
        "2. `__len__`, and \n",
        "3. `__getitem__`. \n",
        "\n",
        "We'll look at an example implementation.\n",
        "\n",
        "The CIFAR10 images are stored in a directory and their labels are stored\n",
        "separately in a CSV file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import pandas as pd\n",
        "import torch\n",
        "from torch.utils.data import Dataset\n",
        "from torchvision.io import read_image\n",
        "\n",
        "class MyCustomImageDataset(Dataset):\n",
        "    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):\n",
        "        self.img_labels = pd.read_csv(annotations_file)\n",
        "        self.img_dir = img_dir\n",
        "        self.transform = transform\n",
        "        self.target_transform = target_transform\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.img_labels)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        img_path = self.img_labels.iloc[idx, 0]\n",
        "        image = read_image(img_path)\n",
        "        label = self.img_labels.iloc[idx, 1]\n",
        "\n",
        "        if self.transform:\n",
        "            image = self.transform(image)\n",
        "        if self.target_transform:\n",
        "            label = self.target_transform(label)\n",
        "        return image, label"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### `__init__`\n",
        "\n",
        "The `__init__` function is run once when instantiating the Dataset\n",
        "object. We initialize the directory containing the images, the\n",
        "annotations file, and both transforms (covered in more detail in the\n",
        "next section).\n",
        "\n",
        "```{.python}\n",
        "def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):\n",
        "    self.img_labels = pd.read_csv(annotations_file)\n",
        "    self.img_dir = img_dir\n",
        "    self.transform = transform\n",
        "    self.target_transform = target_transform\n",
        "```\n",
        "\n",
        "### `__len__`\n",
        "\n",
        "The `len()` function returns the number of samples in our dataset.\n",
        "\n",
        "Example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def __len__(self):\n",
        "    return len(self.img_labels)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### `__getitem__`\n",
        "\n",
        "The `getitem()` function loads and returns a sample from the dataset\n",
        "at the given index `idx`. \n",
        "\n",
        "Based on the index, it \n",
        "\n",
        "* identifies the image's location on disk, \n",
        "* converts that to a tensor using `read_image`,\n",
        "* retrieves the corresponding label from the csv data in `self.img_labels`, \n",
        "* calls the transform functions on them (if applicable), and \n",
        "* returns the tensor image and corresponding label in a tuple.\n",
        "\n",
        "```{.python}\n",
        "def __getitem__(self, idx):\n",
        "    img_path = self.img_labels.iloc[idx, 0]\n",
        "    image = read_image(img_path)\n",
        "    label = self.img_labels.iloc[idx, 1]\n",
        "\n",
        "    if self.transform:\n",
        "        image = self.transform(image)\n",
        "    if self.target_transform:\n",
        "        label = self.target_transform(label)\n",
        "    return image, label\n",
        "\n",
        "```\n",
        "<br>\n",
        "\n",
        "**Let's try to use it.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from torchvision.transforms import ToTensor\n",
        "\n",
        "mydataset = MyCustomImageDataset(annotations_file='cifar10_annotations.csv', img_dir='cifar10_images') # , transform=ToTensor())\n",
        "\n",
        "img, label = mydataset[0]\n",
        "print(f\"img.shape: {img.shape}\")\n",
        "print(f\"label: {label}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "figure = plt.figure(figsize=(6, 6))\n",
        "cols, rows = 2, 2\n",
        "\n",
        "for i in range(1, cols * rows + 1):\n",
        "    # Randomly choose indices\n",
        "    sample_idx = torch.randint(len(mydataset), size=(1,)).item()\n",
        "\n",
        "    img, label = mydataset[sample_idx]\n",
        "    #print(f\"img.shape: {img.shape}\")\n",
        "    #print(f\"label: {label}\")\n",
        "    \n",
        "    figure.add_subplot(rows, cols, i)\n",
        "    plt.title(labels_map[label])\n",
        "    plt.axis(\"off\")\n",
        "    plt.imshow(img.permute(1, 2, 0))\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "------------------------------------------------------------------------\n",
        "\n",
        "### Preparing your data for training with DataLoaders\n",
        "\n",
        "The `Dataset` retrieves your dataset's features and labels one sample at\n",
        "a time. \n",
        "\n",
        "While training a model, we typically want to \n",
        "\n",
        "* pass samples in \"batches\", \n",
        "* reshuffle the data at every epoch to reduce model overfitting, and \n",
        "* use Python's `multiprocessing` to speed up data retrieval.\n",
        "\n",
        "`DataLoader` is an iterable that abstracts this complexity for us in an\n",
        "easy API.\n",
        "\n",
        "Because `MyCustomImageDataset` is a subclass of `Dataset`, we can use it\n",
        "with a `DataLoader` just as we would use any other `Dataset`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from torch.utils.data import DataLoader\n",
        "\n",
        "train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)\n",
        "test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Iterate through the DataLoader\n",
        "\n",
        "We have loaded that dataset into the `DataLoader` and can iterate\n",
        "through the dataset as needed. \n",
        "\n",
        "Each iteration below returns a batch of\n",
        "`train_features` and `train_labels` (containing `batch_size=64` features\n",
        "and labels respectively). \n",
        "\n",
        "Because we specified `shuffle=True`, after we\n",
        "iterate over all batches the data is shuffled\n",
        "\n",
        "(for finer-grained control over the data loading order, take a look at\n",
        "[Samplers](https://pytorch.org/docs/stable/data.html#data-loading-order-and-sampler))."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Display image and label.\n",
        "train_features, train_labels = next(iter(train_dataloader))\n",
        "print(f\"Feature batch shape: {train_features.size()}\")\n",
        "print(f\"Labels batch shape: {train_labels.size()}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Note from the size that we got a batch of 64 images."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "img = train_features[0].squeeze() # squeeze() removes dimension of size 1, e.g. (1, 3, 32, 32) -> (3, 32, 32)\n",
        "label = train_labels[0]\n",
        "\n",
        "plt.imshow(img.permute(1, 2, 0))\n",
        "plt.show()\n",
        "\n",
        "print(f\"Label: {labels_map[label.item()]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Transforms and Data Augmentation\n",
        "\n",
        "Data does not always come in its final processed form that is required\n",
        "for training machine learning algorithms. \n",
        "\n",
        "We use **transforms** to perform some manipulation of the data and make it suitable for training as well as augmentation.\n",
        "\n",
        "All TorchVision dataset classes have two parameters \n",
        "\n",
        "- `transform` to modify the features and \n",
        "- `target_transform` to modify the labels\n",
        "\n",
        "They accept callables containing the transformation logic. \n",
        "\n",
        "The CIFAR10 images are in PIL Image format, and the labels are\n",
        "integers. \n",
        "\n",
        "For training, we need the features as normalized tensors, and\n",
        "the labels as one-hot encoded tensors. \n",
        "\n",
        "To make these transformations, we use `ToTensor` and `Lambda`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "from torchvision import datasets\n",
        "from torchvision.transforms import ToTensor, Lambda\n",
        "\n",
        "ds = datasets.CIFAR10(\n",
        "    root=\"data\",\n",
        "    train=True,\n",
        "    download=True,\n",
        "    transform=ToTensor(),\n",
        "    target_transform=Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter_(dim=0, index=torch.tensor(y), value=1))\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### ToTensor()\n",
        "\n",
        "[ToTensor](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.ToTensor)\n",
        "converts a PIL image or NumPy `ndarray` into a `FloatTensor` and scales\n",
        "the image's pixel intensity values in the range $0., 1.$\n",
        "\n",
        "### Lambda Transforms\n",
        "\n",
        "Lambda transforms apply any user-defined lambda function. \n",
        "\n",
        "Here, we define a function to turn the integer into a one-hot encoded tensor. \n",
        "\n",
        "It \n",
        "\n",
        "* first creates a zero tensor of size 10 (the number of labels in our dataset) and \n",
        "* calls [`scatter_`](https://pytorch.org/docs/stable/generated/torch.Tensor.scatter_.html)\n",
        "  which assigns a `value=1` on the index as given by the label `y`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "target_transform = Lambda(lambda y: torch.zeros(\n",
        "    10, dtype=torch.float).scatter_(dim=0, index=torch.tensor(y), value=1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Transforms V2 and Data Augmentation\n",
        "\n",
        "The [torchvision.transforms](https://pytorch.org/vision/stable/transforms.html)\n",
        "module offers many commonly-used transforms for image and video data.\n",
        "\n",
        "Adapted from [Getting started with transforms v2](https://pytorch.org/vision/stable/auto_examples/transforms/plot_transforms_getting_started.html#sphx-glr-auto-examples-transforms-plot-transforms-getting-started-py)\n",
        "\n",
        "## Transforms V2 Example\n",
        "\n",
        "We'll illustrate with an example.\n",
        "\n",
        "First, a bit of setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "plt.rcParams[\"savefig.bbox\"] = 'tight'\n",
        "\n",
        "from torchvision.transforms import v2\n",
        "from torchvision.io import read_image\n",
        "\n",
        "import os\n",
        "import urllib.request\n",
        "import zipfile\n",
        "\n",
        "# Check if the directory exists\n",
        "if not os.path.exists('transform_v2_files'):\n",
        "    # Download the zip file\n",
        "    url = 'https://github.com/trgardos/ml-549-fa24/raw/refs/heads/main/transform_v2_files.zip?download='\n",
        "    zip_path = 'transform_v2_files.zip'\n",
        "    urllib.request.urlretrieve(url, zip_path)\n",
        "    \n",
        "    # Unzip the file\n",
        "    with zipfile.ZipFile(zip_path, 'r') as zip_ref:\n",
        "        zip_ref.extractall('.')\n",
        "    \n",
        "    # Remove the zip file\n",
        "    os.remove(zip_path)\n",
        "\n",
        "\n",
        "torch.manual_seed(1)\n",
        "\n",
        "from transform_v2_files.transforms.helpers import plot\n",
        "img = read_image(str(Path('transform_v2_files/assets') / 'astronaut.jpg'))\n",
        "print(f\"{type(img) = }, {img.dtype = }, {img.shape = }\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### The basics\n",
        "\n",
        "The Torchvision transforms behave like a regular [`torch.nn.Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html):\n",
        "\n",
        "* instantiate a transform, \n",
        "* pass an input, \n",
        "* get a transformed output:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "transform = v2.RandomCrop(size=(224, 224))\n",
        "out = transform(img)\n",
        "\n",
        "plot([img, out])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This transform takes a random crop of the image of size 224x224. Try\n",
        "re-running the cell a few times to see different samples.\n",
        "\n",
        "### Transforms for image classification\n",
        "\n",
        "Here's a basic image classification transform pipeline:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "transforms = v2.Compose([\n",
        "    v2.RandomResizedCrop(size=(224, 224), antialias=True),\n",
        "    v2.RandomHorizontalFlip(p=0.5),\n",
        "    v2.ToDtype(torch.float32, scale=True),\n",
        "    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
        "])\n",
        "out = transforms(img)\n",
        "\n",
        "plot([img, out])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Again, try re-running the cell a few times to see different samples.\n",
        "\n",
        "Such transformation pipeline is typically passed as the `transform` argument\n",
        "to the `Datasets` class, e.g. `ImageNet(...,transform=transforms)`.\n",
        "\n",
        "See the [Transforms Gallery](https://pytorch.org/vision/stable/auto_examples/transforms/index.html#transforms-gallery) for examples of how to use augmentation transforms like `CutMix` and `MixUp`.\n",
        "\n",
        "### Detection, Segmentation, Videos\n",
        "\n",
        "Besides images, transforms can also transform bounding boxes, segmentation / detection masks, or videos.\n",
        "\n",
        "Let's look at an example with bounding boxes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from torchvision import tv_tensors  # we'll describe this a bit later\n",
        "\n",
        "boxes = tv_tensors.BoundingBoxes(\n",
        "    [\n",
        "        [15, 10, 370, 510],\n",
        "        [275, 340, 510, 510],\n",
        "        [130, 345, 210, 425]\n",
        "    ],\n",
        "    format=\"XYXY\", canvas_size=img.shape[-2:])\n",
        "\n",
        "transforms = v2.Compose([\n",
        "    v2.RandomResizedCrop(size=(224, 224), antialias=True),\n",
        "    v2.RandomPhotometricDistort(p=1),\n",
        "    v2.RandomHorizontalFlip(p=1),\n",
        "])\n",
        "out_img, out_boxes = transforms(img, boxes)\n",
        "print(type(boxes), type(out_boxes))\n",
        "\n",
        "plot([(img, boxes), (out_img, out_boxes)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This also works for:\n",
        "* masks (`torchvision.tv_tensors.Mask`) for object segmentation or semantic\n",
        "segmentation, or\n",
        "* videos (`torchvision.tv_tensors.Video`)\n",
        "\n",
        "We could have passed them to the transforms in exactly the same way.\n",
        "\n",
        "\n",
        "### What are TVTensors?\n",
        "\n",
        "TVTensors are [`torch.Tensor`](https://pytorch.org/docs/stable/tensors.html#torch.Tensor) subclasses.\n",
        "\n",
        "The available TVTensors are [`Image`](https://pytorch.org/vision/stable/generated/torchvision.tv_tensors.Image.html#torchvision.tv_tensors.Image),\n",
        "[`BoundingBoxes`](https://pytorch.org/vision/stable/generated/torchvision.tv_tensors.BoundingBoxes.html#torchvision.tv_tensors.BoundingBoxes),\n",
        "[`Mask`](https://pytorch.org/vision/stable/generated/torchvision.tv_tensors.Mask.html#torchvision.tv_tensors.Mask), and\n",
        "[`Video`](https://pytorch.org/vision/stable/generated/torchvision.tv_tensors.Video.html#torchvision.tv_tensors.Video).\n",
        "\n",
        "TVTensors look and feel just like regular tensors - they **are** tensors.\n",
        "\n",
        "Everything that is supported on a plain `torch.Tensor` like `.sum()`\n",
        "or any `torch.*` operator will also work on a TVTensor:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "img_dp = tv_tensors.Image(torch.randint(0, 256, (3, 256, 256), dtype=torch.uint8))\n",
        "\n",
        "print(f\"{isinstance(img_dp, torch.Tensor) = }\")\n",
        "print(f\"{img_dp.dtype = }, {img_dp.shape = }, {img_dp.sum() = }\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "These TVTensor classes are at the core of the transforms: \n",
        "\n",
        "in order to transform a given input, the transforms first look at the **class**\n",
        "of the object, and dispatch to the appropriate implementation accordingly.\n",
        "\n",
        "For more info, see the [TVTensors FAQ](https://pytorch.org/vision/stable/auto_examples/transforms/plot_tv_tensors.html#sphx-glr-auto-examples-transforms-plot-tv-tensors-py).\n",
        "\n",
        "### What do I pass as input?\n",
        "\n",
        "Above, we've seen two examples: \n",
        "\n",
        "* one where we passed a single image as input i.e. ``out = transforms(img)``, and \n",
        "* one where we passed both an image and bounding boxes, i.e.\n",
        " ``out_img, out_boxes = transforms(img, boxes)``.\n",
        "\n",
        "Transforms support **arbitrary input structures**. \n",
        "\n",
        "The input can be e.g. a single image, a tuple, an arbitrarily nested dictionary...\n",
        "\n",
        "The same structure will be returned as output. \n",
        "\n",
        "Below, we use the\n",
        "same detection transforms, but pass a tuple (image, target_dict) as input and\n",
        "we're getting the same structure as output:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "target = {\n",
        "    \"boxes\": boxes,\n",
        "    \"labels\": torch.arange(boxes.shape[0]),\n",
        "    \"this_is_ignored\": (\"arbitrary\", {\"structure\": \"!\"})\n",
        "}\n",
        "\n",
        "# Re-using the transforms and definitions from above.\n",
        "out_img, out_target = transforms(img, target)\n",
        "\n",
        "plot([(img, target[\"boxes\"]), (out_img, out_target[\"boxes\"])])\n",
        "print(f\"{out_target['this_is_ignored']}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We passed a tuple so we get a tuple back, and the second element is the\n",
        "tranformed target dict. \n",
        "\n",
        "Transforms don't really care about the structure of the input; as mentioned\n",
        "above, they only care about the **type** of the objects and transforms them\n",
        "accordingly.\n",
        "\n",
        "*Foreign* objects like strings or ints are simply passed-through. This can be\n",
        "useful e.g. if you want to associate a path with every single sample when\n",
        "debugging!\n",
        "\n",
        "\n",
        "## Transforms and Datasets intercompatibility\n",
        "\n",
        "Roughly speaking, the output of the datasets must correspond to the input of\n",
        "the transforms. How to do that depends on whether you're using the torchvision\n",
        "[built-in datatsets](https://pytorch.org/vision/stable/datasets.html#datasets),\n",
        "or your own custom datasets.\n",
        "\n",
        "### Using built-in datasets\n",
        "\n",
        "If you're just doing image classification, you don't need to do anything. Just\n",
        "use `transform` argument of the dataset e.g. `ImageNet(..., transform=transforms)`.\n",
        "\n",
        "#### Compatbility with Older Datasets\n",
        "\n",
        "Torchvision also supports datasets for object detection or segmentation like\n",
        "[`torchvision.datasets.CocoDetection`](https://pytorch.org/vision/stable/generated/torchvision.datasets.CocoDetection.html#torchvision.datasets.CocoDetection). \n",
        "Those datasets predate\n",
        "the existence of the `transforms.v2` module and of the\n",
        "TVTensors, so they don't return TVTensors out of the box.\n",
        "\n",
        "An easy way to force those datasets to return TVTensors and to make them\n",
        "compatible with v2 transforms is to use the\n",
        "`wrap_dataset_for_transforms_v2` function:\n",
        "\n",
        "```python\n",
        "from torchvision.datasets import CocoDetection, wrap_dataset_for_transforms_v2\n",
        "\n",
        "dataset = CocoDetection(..., transforms=my_transforms)\n",
        "dataset = wrap_dataset_for_transforms_v2(dataset)\n",
        "# Now the dataset returns TVTensors!\n",
        "```\n",
        "### Using your own datasets\n",
        "\n",
        "If you have a custom dataset, then you'll need to convert your objects into\n",
        "the appropriate TVTensor classes. \n",
        "\n",
        "To create TVTensor instances refer to\n",
        "[How do I create TVTensors?](https://pytorch.org/vision/stable/auto_examples/transforms/plot_tv_tensors.html#tv-tensor-creation)\n",
        "for more details.\n",
        "\n",
        "There are two main places where you can implement that conversion logic:\n",
        "\n",
        "- At the end of the datasets's ``__getitem__`` method, before returning the\n",
        "  sample (or by sub-classing the dataset).\n",
        "- As the very first step of your transforms pipeline\n",
        "\n",
        "Either way, the logic will depend on your specific dataset.\n",
        "\n",
        "## V2 Transforms\n",
        "\n",
        "See [V2 API Snapshot](https://pytorch.org/vision/stable/transforms.html#v2-api-reference-recommended) for a more complete list.\n",
        "\n",
        "\n",
        "### Further Reading\n",
        "\n",
        "-   [torch.utils.data API](https://pytorch.org/docs/stable/data.html)\n",
        "-   [torchvision.transforms API](https://pytorch.org/vision/stable/transforms.html)\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3",
      "path": "/Users/tomg/Source/courses/ds549/ml-549-fa24/.venv/share/jupyter/kernels/python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}