pytorch/data_loading_tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nData Loading and Processing Tutorial\n====================================\n**Author**: `Sasank Chilamkurthy <https://chsasank.github.io>`_\n\nA lot of effort in solving any machine learning problem goes in to\npreparing the data. PyTorch provides many tools to make data loading\neasy and hopefully, to make your code more readable. In this tutorial,\nwe will see how to load and preprocess/augment data from a non trivial\ndataset.\n\nTo run this tutorial, please make sure the following packages are\ninstalled:\n\n-  ``scikit-image``: For image io and transforms\n-  ``pandas``: For easier csv parsing\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from __future__ import print_function, division\nimport os\nimport torch\nimport pandas as pd\nfrom skimage import io, transform\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, utils\n\n# Ignore warnings\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\nplt.ion()   # interactive mode"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The dataset we are going to deal with is that of facial pose.\nThis means that a face is annotated like this:\n\n.. figure:: /_static/img/landmarked_face2.png\n   :width: 400\n\nOver all, 68 different landmark points are annotated for each face.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>Download the dataset from `here <https://download.pytorch.org/tutorial/faces.zip>`_\n    so that the images are in a directory named 'faces/'.\n    This dataset was actually\n    generated by applying excellent `dlib's pose\n    estimation <http://blog.dlib.net/2014/08/real-time-face-pose-estimation.html>`__\n    on a few images from imagenet tagged as 'face'.</p></div>\n\nDataset comes with a csv file with annotations which looks like this:\n\n::\n\n    image_name,part_0_x,part_0_y,part_1_x,part_1_y,part_2_x, ... ,part_67_x,part_67_y\n    0805personali01.jpg,27,83,27,98, ... 84,134\n    1084239450_e76e00b7e7.jpg,70,236,71,257, ... ,128,312\n\nLet's quickly read the CSV and get the annotations in an (N, 2) array where N\nis the number of landmarks.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "landmarks_frame = pd.read_csv('faces/face_landmarks.csv')\n\nn = 65\nimg_name = landmarks_frame.ix[n, 0]\nlandmarks = landmarks_frame.ix[n, 1:].as_matrix().astype('float')\nlandmarks = landmarks.reshape(-1, 2)\n\nprint('Image name: {}'.format(img_name))\nprint('Landmarks shape: {}'.format(landmarks.shape))\nprint('First 4 Landmarks: {}'.format(landmarks[:4]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's write a simple helper function to show an image and its landmarks\nand use it to show a sample.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def show_landmarks(image, landmarks):\n    \"\"\"Show image with landmarks\"\"\"\n    plt.imshow(image)\n    plt.scatter(landmarks[:, 0], landmarks[:, 1], s=10, marker='.', c='r')\n    plt.pause(0.001)  # pause a bit so that plots are updated\n\nplt.figure()\nshow_landmarks(io.imread(os.path.join('faces/', img_name)),\n               landmarks)\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Dataset class\n-------------\n\n``torch.utils.data.Dataset`` is an abstract class representing a\ndataset.\nYour custom dataset should inherit ``Dataset`` and override the following\nmethods:\n\n-  ``__len__`` so that ``len(dataset)`` returns the size of the dataset.\n-  ``__getitem__`` to support the indexing such that ``dataset[i]`` can\n   be used to get $i$\\ th sample\n\nLet's create a dataset class for our face landmarks dataset. We will\nread the csv in ``__init__`` but leave the reading of images to\n``__getitem__``. This is memory efficient because all the images are not\nstored in the memory at once but read as required.\n\nSample of our dataset will be a dict\n``{'image': image, 'landmarks': landmarks}``. Our datset will take an\noptional argument ``transform`` so that any required processing can be\napplied on the sample. We will see the usefulness of ``transform`` in the\nnext section.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "class FaceLandmarksDataset(Dataset):\n    \"\"\"Face Landmarks dataset.\"\"\"\n\n    def __init__(self, csv_file, root_dir, transform=None):\n        \"\"\"\n        Args:\n            csv_file (string): Path to the csv file with annotations.\n            root_dir (string): Directory with all the images.\n            transform (callable, optional): Optional transform to be applied\n                on a sample.\n        \"\"\"\n        self.landmarks_frame = pd.read_csv(csv_file)\n        self.root_dir = root_dir\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.landmarks_frame)\n\n    def __getitem__(self, idx):\n        img_name = os.path.join(self.root_dir, self.landmarks_frame.ix[idx, 0])\n        image = io.imread(img_name)\n        landmarks = self.landmarks_frame.ix[idx, 1:].as_matrix().astype('float')\n        landmarks = landmarks.reshape(-1, 2)\n        sample = {'image': image, 'landmarks': landmarks}\n\n        if self.transform:\n            sample = self.transform(sample)\n\n        return sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's instantiate this class and iterate through the data samples. We\nwill print the sizes of first 4 samples and show their landmarks.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "face_dataset = FaceLandmarksDataset(csv_file='faces/face_landmarks.csv',\n                                    root_dir='faces/')\n\nfig = plt.figure()\n\nfor i in range(len(face_dataset)):\n    sample = face_dataset[i]\n\n    print(i, sample['image'].shape, sample['landmarks'].shape)\n\n    ax = plt.subplot(1, 4, i + 1)\n    plt.tight_layout()\n    ax.set_title('Sample #{}'.format(i))\n    ax.axis('off')\n    show_landmarks(**sample)\n\n    if i == 3:\n        plt.show()\n        break"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Transforms\n----------\n\nOne issue we can see from the above is that the samples are not of the \nsame size. Most neural networks expect the images of a fixed size.\nTherefore, we will need to write some prepocessing code.\nLet's create three transforms:\n\n-  ``Rescale``: to scale the image\n-  ``RandomCrop``: to crop from image randomly. This is data\n   augmentation.\n-  ``ToTensor``: to convert the numpy images to torch images (we need to\n   swap axes).\n\nWe will write them as callable classes instead of simple functions so\nthat parameters of the transform need not be passed everytime it's\ncalled. For this, we just need to implement ``__call__`` method and \nif required, ``__init__`` method. We can then use a transform like this:\n\n::\n\n    tsfm = Transform(params)\n    transformed_sample = tsfm(sample)\n\nObserve below how these transforms had to be applied both on the image and\nlandmarks.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "class Rescale(object):\n    \"\"\"Rescale the image in a sample to a given size.\n\n    Args:\n        output_size (tuple or tuple): Desired output size. If tuple, output is\n            matched to output_size. If int, smaller of image edges is matched\n            to output_size keeping aspect ratio the same.\n    \"\"\"\n\n    def __init__(self, output_size):\n        assert isinstance(output_size, (int, tuple))\n        self.output_size = output_size\n\n    def __call__(self, sample):\n        image, landmarks = sample['image'], sample['landmarks']\n\n        h, w = image.shape[:2]\n        if isinstance(self.output_size, int):\n            if h > w:\n                new_h, new_w = self.output_size * h / w, self.output_size\n            else:\n                new_h, new_w = self.output_size, self.output_size * w / h\n        else:\n            new_h, new_w = self.output_size\n\n        new_h, new_w = int(new_h), int(new_w)\n\n        img = transform.resize(image, (new_h, new_w))\n\n        # h and w are swapped for landmarks because for images,\n        # x and y axes are axis 1 and 0 respectively\n        landmarks = landmarks * [new_w / w, new_h / h]\n\n        return {'image': img, 'landmarks': landmarks}\n\n\nclass RandomCrop(object):\n    \"\"\"Crop randomly the image in a sample.\n\n    Args:\n        output_size (tuple or int): Desired output size. If int, square crop\n            is made.\n    \"\"\"\n\n    def __init__(self, output_size):\n        assert isinstance(output_size, (int, tuple))\n        if isinstance(output_size, int):\n            self.output_size = (output_size, output_size)\n        else:\n            assert len(output_size) == 2\n            self.output_size = output_size\n\n    def __call__(self, sample):\n        image, landmarks = sample['image'], sample['landmarks']\n\n        h, w = image.shape[:2]\n        new_h, new_w = self.output_size\n\n        top = np.random.randint(0, h - new_h)\n        left = np.random.randint(0, w - new_w)\n\n        image = image[top: top + new_h,\n                      left: left + new_w]\n\n        landmarks = landmarks - [left, top]\n\n        return {'image': image, 'landmarks': landmarks}\n\n\nclass ToTensor(object):\n    \"\"\"Convert ndarrays in sample to Tensors.\"\"\"\n\n    def __call__(self, sample):\n        image, landmarks = sample['image'], sample['landmarks']\n \n        # swap color axis because\n        # numpy image: H x W x C\n        # torch image: C X H X W\n        image = image.transpose((2, 0, 1))\n        return {'image': torch.from_numpy(image),\n                'landmarks': torch.from_numpy(landmarks)}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compose transforms\n~~~~~~~~~~~~~~~~~~\n\nNow, we apply the transforms on an sample.\n\nLet's say we want to rescale the shorter side of the image to 256 and\nthen randomly crop a square of size 224 from it. i.e, we want to compose\n``Rescale`` and ``RandomCrop`` transforms.\n``torchvision.transforms.Compose`` is a simple callable class which allows us\nto do this.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "scale = Rescale(256)\ncrop = RandomCrop(128)\ncomposed = transforms.Compose([Rescale(256),\n                               RandomCrop(224)])\n\n# Apply each of the above transforms on sample.\nfig = plt.figure()\nsample = face_dataset[65]\nfor i, tsfrm in enumerate([scale, crop, composed]):\n    transformed_sample = tsfrm(sample)\n\n    ax = plt.subplot(1, 3, i + 1)\n    plt.tight_layout()\n    ax.set_title(type(tsfrm).__name__)\n    show_landmarks(**transformed_sample)\n\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Iterating through the dataset\n-----------------------------\n\nLet's put this all together to create a dataset with composed\ntransforms. \nTo summarize, every time this dataset is sampled:\n\n-  An image is read from the file on the fly\n-  Transforms are applied on the read image\n-  Since one of the transforms is random, data is augmentated on\n   sampling\n\nWe can iterate over the created dataset with a ``for i in range``\nloop as before.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "transformed_dataset = FaceLandmarksDataset(csv_file='faces/face_landmarks.csv',\n                                           root_dir='faces/',\n                                           transform=transforms.Compose([\n                                               Rescale(256),\n                                               RandomCrop(224),\n                                               ToTensor()\n                                           ]))\n\nfor i in range(len(transformed_dataset)):\n    sample = transformed_dataset[i]\n\n    print(i, sample['image'].size(), sample['landmarks'].size())\n\n    if i == 3:\n        break"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "However, we are losing a lot of features by using a simple ``for`` loop to\niterate over the data. In particular, we are missing out on:\n\n-  Batching the data\n-  Shuffling the data\n-  Load the data in parallel using ``multiprocessing`` workers.\n\n``torch.utils.data.DataLoader`` is an iterator which provides all these\nfeatures. Parameters used below should be clear. One parameter of\ninterest is ``collate_fn``. You can specify how exactly the samples need\nto be batched using ``collate_fn``. However, default collate should work\nfine for most use cases.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "dataloader = DataLoader(transformed_dataset, batch_size=4,\n                        shuffle=True, num_workers=4)\n\n\n# Helper function to show a batch\ndef show_landmarks_batch(sample_batched):\n    \"\"\"Show image with landmarks for a batch of samples.\"\"\"\n    images_batch, landmarks_batch = \\\n            sample_batched['image'], sample_batched['landmarks']\n    batch_size = len(images_batch)\n    im_size = images_batch.size(2)\n\n    grid = utils.make_grid(images_batch)\n    plt.imshow(grid.numpy().transpose((1, 2, 0)))\n\n    for i in range(batch_size):\n        plt.scatter(landmarks_batch[i, :, 0].numpy() + i * im_size,\n                    landmarks_batch[i, :, 1].numpy(),\n                    s=10, marker='.', c='r')\n\n        plt.title('Batch from dataloader')\n\nfor i_batch, sample_batched in enumerate(dataloader):\n    print(i_batch, sample_batched['image'].size(),\n          sample_batched['landmarks'].size())\n\n    # observe 4th batch and stop.\n    if i_batch == 3:\n        plt.figure()\n        show_landmarks_batch(sample_batched)\n        plt.axis('off')\n        plt.ioff()\n        plt.show()\n        break"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Afterword: torchvision\n----------------------\n\nIn this tutorial, we have seen how to write and use datasets, transforms\nand dataloader. ``torchvision`` package provides some common datasets and\ntransforms. You might not even have to write custom classes. One of the\nmore generic datasets available in torchvision is ``ImageFolder``.\nIt assumes that images are organized in the following way: ::\n\n    root/ants/xxx.png\n    root/ants/xxy.jpeg\n    root/ants/xxz.png\n    .\n    .\n    .\n    root/bees/123.jpg\n    root/bees/nsdf3.png\n    root/bees/asd932_.png\n\nwhere 'ants', 'bees' etc. are class labels. Similarly generic transforms\nwhich operate on ``PIL.Image`` like  ``RandomHorizontalFlip``, ``Scale``,\nare also avaiable. You can use these to write a dataloader like this: ::\n\n  import torch\n  from torchvision import transforms, datasets\n\n  data_transform = transforms.Compose([\n          transforms.RandomSizedCrop(224),\n          transforms.RandomHorizontalFlip(),\n          transforms.ToTensor(),\n          transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                               std=[0.229, 0.224, 0.225])\n      ])\n  hymenoptera_dataset = datasets.ImageFolder(root='hymenoptera_data/train',\n                                             transform=data_transform)\n  dataset_loader = torch.utils.data.DataLoader(hymenoptera_dataset,\n                                               batch_size=4, shuffle=True,\n                                               num_workers=4)\n\nFor an example with training code, please see\n:doc:`transfer_learning_tutorial`.\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.6.1"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}