{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Exercise 9: Semantic Segmentation with Deep Learning\n", "\n", "In the previous exercise, we implemented a full routine for training and validating DL models for image classification.\n", "Now, we shall take it to the next level and perform pixel-wise classification, also known as semantic segmentation.\n", "This is of particular interest to remote sensing, as it allows us to e.g. obtain spatially well-resolved land cover maps, among other products.\n", "\n", "The basic ingredients are exactly the same as for image classification (optimiser, loss function, training loop, etc.). We do have some changes to make, though:\n", "* Dataset: this time, we don't just want a single output number (class), but one value per spatial location. Essentially, our dataset should also provide a second image where each pixel has the class index as its value.\n", "* Model: likewise, we need a suitable model that provides spatial outputs rather than a single class vector. You have seen some examples in the lecture. In this exercise we shall use a flavour of the [Hypercolumn](https://www.cv-foundation.org/openaccess/content_cvpr_2015/papers/Hariharan_Hypercolumns_for_Object_2015_CVPR_paper.pdf) to do this job." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1. Setup" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 1.1 Install dependencies\n", "\n", "Please re-run this as we now need another package (for downloading the data)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import sys\n", "!{sys.executable} -m pip install torch torchvision\n", "!{sys.executable} -m pip install matplotlib\n", "!{sys.executable} -m pip install tqdm # this gives us a pretty progress bar\n", "\n", "# for downloading files from Google drive\n", "!{sys.executable} -m pip install gdown" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 1.2 Check if GPU available\n", "\n", "This is even more important for semantic segmentation, as our models and data tensors are going to be significantly larger.\n", "\n", "Run the following code block and proceed if the response is `True`. Else see an instructor." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "\n", "print(torch.cuda.is_available())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 1.3 Random seed" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "seed = 323444 # the seed value used to initialise the random number generator of PyTorch\n", "torch.manual_seed(seed)\n", "torch.cuda.manual_seed(seed)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 2. Dataset\n", "\n", "For this exercise we shall be using the [ISPRS Vaihingen](https://www2.isprs.org/commissions/comm2/wg4/benchmark/2d-sem-label-vaihingen/) semantic segmentation dataset.\n", "This is a set of fully-labelled satellite image-segmentation mask pairs, with 9cm resolution and six land cover classes: Impervious, Buildings, Low Vegetation, Tree, Car, Clutter. The images come from a large satellite scene over the town Vaihingen in Germany and were divided into 33 patches, some of which are available with ground truth. These patches are still too large for our model – we would quickly run out of GPU memory if we tried to process an image of e.g. 4000x3000 pixels. Hence, they need to be further divided into even smaller patches. This has already been done for you – all you need to do is to download the image-label pair patches (sized 512x512 pixels) by running the code cell below." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!gdown --id 16OkRr9Ck-XGKy3LjNqGxCTx0pmU1aoUE\n", "\n", "!unzip /content/Vaihingen_512_512.zip\n", "data_root = 'Vaihingen_512_512'" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Like before, we need to wrap these images in a Dataset class.\n", "It's quite involved this time, however, since we don't just have simple RGB images, but need to collect multiple satellite products. The code below does all of this for you." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "import torch\n", "from torch.utils.data import dataset\n", "from torch.utils.data import DataLoader\n", "import torchvision.transforms as T # transformations that can be used e.g. for data conversion or augmentation\n", "import numpy as np\n", "from PIL import Image\n", "\n", "\n", "class VaihingenDataset(dataset.Dataset):\n", " '''\n", " Custom Dataset class that loads images and ground truth segmentation\n", " masks from a directory.\n", " '''\n", "\n", " # image statistics, calculated in advance as averages across the full\n", " # training data set\n", " IMAGE_MEANS = (\n", " (121.03431026287558, 82.52572736507886, 81.92368178210943), # IR-R-G tiles\n", " (285.34753853934154) # DSM\n", " )\n", " IMAGE_STDS = (\n", " (54.21029197978022, 38.434924159900554, 37.040640374137475), # IR-R-G tiles\n", " (6.485453035150256) # DSM\n", " )\n", "\n", "\n", " # label class names\n", " LABEL_CLASSES = (\n", " 'Buildings', 'Tree', 'Low Vegetation', 'Clutter', 'Car', 'Impervious'\n", " )\n", "\n", "\n", " def __init__(self, data_root):\n", " '''\n", " Dataset class constructor. Here we initialize the dataset instance\n", " and retrieve file names (and other metadata, if present) for all the\n", " images and labels (ground truth semantic segmentation maps).\n", " '''\n", " super().__init__()\n", "\n", " self.data_root = data_root\n", " # List all the files in image folder and make a list of samples\n", " imgs_folder = os.path.join(self.data_root,'imgs')\n", " samples_list = os.listdir(imgs_folder)\n", " self.data = []\n", " for sample in samples_list:\n", " if 'tif' in sample:\n", " sample_name = '_'.join(sample.rsplit('_',2)[1:])\n", " self.data.append(sample_name)\n", "\n", "\n", " def __len__(self):\n", " '''\n", " This function tells the Data Loader how many images there are in\n", " this dataset.\n", " '''\n", " return len(self.data)\n", "\n", " \n", " def __getitem__(self, idx):\n", " '''\n", " Here's where we load, prepare, and convert the images and\n", " segmentation mask for the data element at the given \"idx\".\n", " '''\n", " item = self.data[idx]\n", "\n", " # load image\n", " img_filename = 'top_mosaic_09cm_'+item\n", " image = Image.open(os.path.join(self.data_root,'imgs',img_filename))\n", " # load dsm \n", " dsm_filename = 'dsm_09cm_matching_'+item\n", " dsm = Image.open(os.path.join(self.data_root,'dsm',dsm_filename))\n", " # load segmentation mask (groundtruth)\n", " labels = Image.open(os.path.join(self.data_root,'gts',img_filename))\n", " labels = np.array(labels, dtype=np.int64) # convert to NumPy array temporarily\n", "\n", " # NOTE: at this point it would make sense to perform data augmentation.\n", " # However, the default augmentations built-in to PyTorch (resp.\n", " # Torchvision) (i.) only support RGB images; (ii.) only work on the\n", " # images themselves. In our case, however, we have multispectral data\n", " # and need to also transform the segmentation mask.\n", " # This is not difficult to do, but goes beyond the scope of this exercise.\n", " # For the sake of brevity, we'll leave it out accordingly.\n", " # What we will have to do, however, is to normalize the image data.\n", " image = np.array(image, dtype=np.float32)\n", " image = (image - self.IMAGE_MEANS[0]) / self.IMAGE_STDS[0]\n", " dsm = np.array(dsm, dtype=np.float32)\n", " dsm = (dsm - self.IMAGE_MEANS[1]) / self.IMAGE_STDS[1]\n", "\n", " # finally, we need to convert our data into the torch.Tensor format. For\n", " # the images, we already have a \"ToTensor\" transform available, but we\n", " # need to concatenate the images together.\n", " image = T.ToTensor()(image)\n", " dsm = T.ToTensor()(dsm)\n", " #tensors = [T.ToTensor()(i) for i in images]\n", " inputs = torch.cat([image,dsm], dim=0).float() # concatenate along spectral dimension and make sure it's in 32-bit floating point\n", "\n", " # For the labels, we need to convert the PIL image to a torch.Tensor.\n", " labels = torch.from_numpy(labels).long() # labels need to be in 64-bit integer format\n", "\n", " return inputs, labels\n", "\n", "\n", "\n", "\n", "\n", "# we also create a function for the data loader here (see Section 2.6 in Exercise 6)\n", "def load_dataloader(batch_size, split='train'):\n", " return DataLoader(\n", " VaihingenDataset(os.path.join(data_root, split)),\n", " batch_size=batch_size,\n", " shuffle=(split=='train'), # we shuffle the image order for the training dataset\n", " num_workers=2 # perform data loading with two CPU threads\n", " )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let's visualise some images." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "%matplotlib inline\n", "import matplotlib.pyplot as plt\n", "from matplotlib.colors import ListedColormap\n", "\n", "#discrete color scheme\n", "cMap = ListedColormap(['grey', 'darkgreen', 'lawngreen', 'red', 'orange', 'black']) # #'Buildings', 'Tree', 'Low Vegetation', 'Clutter', 'Car', 'Impervious'\n", "\n", "dataset_train = VaihingenDataset(os.path.join(data_root, 'train'))\n", "\n", "# draw a random sample\n", "idx = torch.randint(0, len(dataset_train), (1,))\n", "data, target = dataset_train.__getitem__(idx)\n", "print(f'Image tensor size: {data.size()}')\n", "print(f'Label tensor size: {target.size()}')\n", "\n", "# visualise\n", "plt.figure()\n", "plt.imshow(data[:3,...].permute(1,2,0).numpy()) # first three bands: NIR-R-G\n", "plt.title('Input: NIR-R-G satellite imagery')\n", "plt.show()\n", "plt.figure()\n", "plt.imshow(data[3,...].squeeze().numpy()) # band 4: DSM\n", "plt.title('Input: DSM')\n", "plt.show()\n", "fig = plt.figure()\n", "cax = plt.imshow(target.squeeze().numpy(), cmap=cMap) # target: segmentation mask\n", "cbar = fig.colorbar(cax, ticks=list(range(len(dataset_train.LABEL_CLASSES))))\n", "cbar.ax.set_yticklabels(list(dataset_train.LABEL_CLASSES))\n", "plt.title('Target: segmentation mask')\n", "plt.show()\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3. Model\n", "\n", "Now let's define our semantic segmentation model! Technically, the model must produce an output that is exactly the same size as its input in space. As stated above we shall use a [Hypercolumn](https://www.cv-foundation.org/openaccess/content_cvpr_2015/papers/Hariharan_Hypercolumns_for_Object_2015_CVPR_paper.pdf) for this task. This is a _Fully Convolutional Network_ (FCN), which means that it does not use a single fully-connected layer, but always preserves some notion of space (so it can use convolutions, pooling, ReLU, etc.). The nice property of FCNs is that they can accept any spatial input of sufficient size and can scale up to the limits of the GPU memory.\n", "\n", "A Hypercolumn basically performs downsampling via convolutions, poolings, etc., like you have been doing in Exercise 6 for image classification. However, unlike a classifier, it keeps every intermediate output, upsamples (interpolates) them to the original image's size, stacks them together to a large tensor (a hypercolumn) and uses this to perform pixel-wise classification:\n", "\n", "![Hypercolumn](https://www.researchgate.net/profile/Devis-Tuia/publication/323273293/figure/fig3/AS:614258178027521@1523461970498/Hypercolumn-based-architecture-used-in-all-our-experiments-Note-that-all-the-layers-are.png)\n", "Image source: Marcos, D., Volpi, M., Kellenberger, B. and Tuia, D., 2018. Land cover mapping at very high resolution with rotation equivariant CNNs: Towards small yet accurate models. ISPRS journal of photogrammetry and remote sensing, 145, pp.96-107.\n", "\n", "\n", "Let's implement a Hypercolumn with the following architecture:\n", "1. BLOCK 1:\n", " 1. 2D convolution, 32 kernels of size 5x5, stride 4, zero-padding 0\n", " 2. 2D max pool, kernel size 2x2, stride 1\n", " 2. Batch Normalisation\n", " 3. ReLU\n", "2. BLOCK 2:\n", " 1. 2D convolution, 64 kernels of size 5x5, stride 4, zero-padding 0\n", " 2. 2D max pool, kernel size 2x2, stride 1\n", " 2. Batch Normalisation\n", " 3. ReLU\n", "3. BLOCK 3:\n", " 1. 2D convolution, 128 kernels of size 5x5, stride 2, zero-padding 0\n", " 2. 2D max pool, kernel size 2x2, stride 1\n", " 2. Batch Normalisation\n", " 3. ReLU\n", "4. BLOCK 4:\n", " 1. 2D convolution, 256 kernels of size 3x3, stride 1, zero-padding 0\n", " 2. 2D max pool, kernel size 2x2, stride 1\n", " 2. Batch Normalisation\n", " 3. ReLU\n", "5. HYPERCOLUMN: here you do the following:\n", " 1. Take all outputs of the input and BLOCKs 1, 2, 3 and 4 (after the ReLU)\n", " 2. Interpolate them to the original input's spatial size (tip: use an instance of [torch.nn.Upsample](https://pytorch.org/docs/stable/generated/torch.nn.Upsample.html))\n", " 3. Concatenate them together (tip: `torch.cat((tensor1, tensor2, ...), dim=1))\n", "6. FINAL BLOCK: this works on the output of 5. HYPERCOLUMN:\n", " 1. 2D convolution, 256 kernels of size 1x1, stride 1, zero-padding 0\n", " 2. Batch Normalisation\n", " 3. ReLU\n", " 4. 2D convolution, 6 kernels of size 1x1, stride 1, zero-padding 0 (output of model)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch.nn as nn\n", "\n", "\n", "class Hypercolumn(nn.Module):\n", "\n", " def __init__(self):\n", " super(Hypercolumn, self).__init__()\n", "\n", " #TODO: define your architecture and forward pass here\n", " # ...\n", " \n", "\n", " def forward(self, x):\n", " #TODO\n", " # ..." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let's test it!" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "dataloader_train = load_dataloader(batch_size=2, split='train')\n", "model = Hypercolumn()\n", "\n", "data, _ = iter(dataloader_train).__next__()\n", "pred = model(data)\n", "\n", "assert pred.size(1) == len(dataset_train.LABEL_CLASSES), f'ERROR: invalid number of model output channels (should be # classes {len(dataset_train.LABEL_CLASSES)}, got {pred.size(1)})'\n", "assert pred.size(2) == data.size(2), f'ERROR: invalid spatial height of model output (should be {data.size(2)}, got {pred.size(2)})'\n", "assert pred.size(3) == data.size(3), f'ERROR: invalid spatial width of model output (should be {data.size(3)}, got {pred.size(3)})'" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 4. Model training\n", "\n", "All the rest is exactly the same principle as for the image classification part!\n", "Hence, what you can do here is to simply copy-paste all your code cells from Section 4 (\"Implement training routine\") of the previous exercise." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#TODO: copy-paste criterion block here\n", "criterion = ..." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#TODO: copy-paste optimiser block here" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#TODO: copy-paste training block here" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#TODO: copy-paste validation block here" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The rest (Section 5 and later from Exercise 7) is also the same, but we change the model name and parameters a bit here, which is why these code blocks are given to you below." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import glob\n", "\n", "os.makedirs('cnn_states/Hypercolumn', exist_ok=True)\n", "\n", "def load_model(epoch='latest'):\n", " model = Hypercolumn()\n", " modelStates = glob.glob('cnn_states/Hypercolumn/*.pth')\n", " if len(modelStates) and (epoch == 'latest' or epoch > 0):\n", " modelStates = [int(m.replace('cnn_states/Hypercolumn/','').replace('.pth', '')) for m in modelStates]\n", " if epoch == 'latest':\n", " epoch = max(modelStates)\n", " stateDict = torch.load(open(f'cnn_states/Hypercolumn/{epoch}.pth', 'rb'), map_location='cpu')\n", " model.load_state_dict(stateDict)\n", " else:\n", " # fresh model\n", " epoch = 0\n", " return model, epoch\n", "\n", "\n", "def save_model(model, epoch):\n", " torch.save(model.state_dict(), open(f'cnn_states/Hypercolumn/{epoch}.pth', 'wb'))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# define hyperparameters\n", "device = 'cuda'\n", "start_epoch = 0 # set to 0 to start from scratch again or to 'latest' to continue training from saved checkpoint\n", "batch_size = 2\n", "learning_rate = 0.1\n", "weight_decay = 0.001\n", "num_epochs = 10\n", "\n", "\n", "\n", "# initialise data loaders\n", "dl_train = load_dataloader(batch_size, 'train')\n", "dl_val = load_dataloader(batch_size, 'val')\n", "\n", "# load model\n", "model, epoch = load_model(epoch=start_epoch)\n", "optim = setup_optimiser(model, learning_rate, weight_decay)\n", "\n", "# do epochs\n", "while epoch < num_epochs:\n", "\n", " # training\n", " model, loss_train, oa_train = train_epoch(dl_train, model, optim, device)\n", "\n", " # validation\n", " loss_val, oa_val = validate_epoch(dl_val, model, device)\n", "\n", " # print stats\n", " print('[Ep. {}/{}] Loss train: {:.2f}, val: {:.2f}; OA train: {:.2f}, val: {:.2f}'.format(\n", " epoch+1, num_epochs,\n", " loss_train, loss_val,\n", " 100*oa_train, 100*oa_val\n", " ))\n", "\n", " # save model\n", " epoch += 1\n", " save_model(model, epoch)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 5. Model validation\n", "\n", "Like in Exercise 7 we could do a final accuracy evaluation now. We don't have access to the Vaihingen dataset's test image labels, since these are hidden on an official Web evaluation server (the Vaihingen dataset was at some point a contest where people could submit their scores and compete against each other!).\n", "\n", "But, we can do something else that we could not do in Exercise 7: visualise our results! Our model provides segmentation masks after all… So let's do this!" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%matplotlib inline\n", "import matplotlib.pyplot as plt\n", "\n", "def visualize(dataLoader, epochs, numImages=5):\n", " models = [load_model(e)[0] for e in epochs]\n", " numModels = len(models)\n", " for idx, (data, labels) in enumerate(dataLoader):\n", " if idx == numImages:\n", " break\n", "\n", " _, ax = plt.subplots(nrows=1, ncols=numModels+1, figsize = (20, 15))\n", "\n", " # plot ground truth\n", " ax[0].imshow(labels[0,...].cpu().numpy())\n", " ax[0].axis('off')\n", " if idx == 0:\n", " ax[0].set_title('Ground Truth')\n", "\n", " for mIdx, model in enumerate(models):\n", " model = model.to(device)\n", " with torch.no_grad():\n", " pred = model(data.to(device))\n", "\n", " # get the label (i.e., the maximum position for each pixel along the class dimension)\n", " yhat = torch.argmax(pred, dim=1)\n", "\n", " # plot model predictions\n", " ax[mIdx+1].imshow(yhat[0,...].cpu().numpy())\n", " ax[mIdx+1].axis('off')\n", " if idx == 0:\n", " ax[mIdx+1].set_title(f'Epoch {epochs[mIdx]}')\n", "\n", "\n", "# visualize predictions for a number of epochs\n", "dl_val_single = load_dataloader(1, 'val')\n", "\n", "# load model states at different epochs\n", "epochs = [0, 1, 5, 'latest'] #TODO: modify this vector according to your wishes, resp. for how many model states you have trained\n", "\n", "visualize(dl_val_single, epochs, numImages=5)" ] } ], "metadata": { "interpreter": { "hash": "d87036bccc21f2c1ac8988b25c74afd5927b201ca402dbbb169419cef9fe6115" }, "kernelspec": { "display_name": "Python 3.7.11 64-bit ('env540': conda)", "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.7.11" }, "orig_nbformat": 4 }, "nbformat": 4, "nbformat_minor": 2 }