Satellite_Imagery/Satellite_Images_Hometask.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h5oP7zLQACZR"
   },
   "source": [
    "# Intro to Digital Agriculture\n",
    "## Week 6\n",
    "Instructors: Maria  Pukalchik, Dmitry Shadrin\n",
    "\n",
    "TAs for this week: Svetlana Illarionova, Ivan Matvienko\n",
    "\n",
    "On this practical we will work with the satellite images, look for agricultural fields and visualize them"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QD9-AvLKCEQJ"
   },
   "source": [
    "# Part 0\n",
    "### Data download and preparation\n",
    "Lets first download the data (can take several minutes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6_yz9jPtAB6q"
   },
   "outputs": [],
   "source": [
    "!wget --no-check-certificate \"https://onedrive.live.com/download?cid=1E2DE865E90D4259&resid=1E2DE865E90D4259%21258622&authkey=AL59NvN10qXgesk\" -O Farmpins.zip\n",
    "!wget --no-check-certificate \"https://onedrive.live.com/download?cid=1E2DE865E90D4259&resid=1E2DE865E90D4259%21195470&authkey=AFPw5W-8uzm5vpM\" -O unet_parts.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ywXB6zwlSlaB"
   },
   "outputs": [],
   "source": [
    "!unzip ./Farmpins.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gX0DUzrPFPQc"
   },
   "source": [
    "And import necessary libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "X-xXliNb-QhN"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import imageio as io\n",
    "import pandas as pd\n",
    "import skimage.transform as transforms\n",
    "import skimage.util as utils\n",
    "import os\n",
    "from matplotlib import pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torch.optim import Adam, lr_scheduler\n",
    "import gdal\n",
    "from unet_parts import *\n",
    "import torchvision\n",
    "from tqdm import tnrange, tqdm\n",
    "from IPython.display import clear_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "xpihfoyP-QhW"
   },
   "outputs": [],
   "source": [
    "from sklearn.metrics import accuracy_score, confusion_matrix, recall_score, f1_score, precision_score, classification_report, SCORERS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CA5dTlCoGZlp"
   },
   "source": [
    "# Part 1. EDA\n",
    "\n",
    "Lets load the images and look what we have inside"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "2oqFprb7GlFq"
   },
   "outputs": [],
   "source": [
    "data_f = gdal.Open('./20170101_mosaic_cropped.tif')\n",
    "data_mask_train = gdal.Open('./train_crops.tif')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hrUGLAEnUUjB"
   },
   "source": [
    "`fields` - data tensor (satellite image) with surface reflectance values\n",
    "\n",
    "`train_mask` - crop labels for field (used for training)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "QOxtI5WzGuuF"
   },
   "outputs": [],
   "source": [
    "fields = np.array(data_f.ReadAsArray())\n",
    "train_mask = np.array(data_mask_train.ReadAsArray())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "aezLtcIqHEK3"
   },
   "outputs": [],
   "source": [
    "train_id = gdal.Open('./train_field_id.tif')\n",
    "train_field_id_map = np.array(train_id.ReadAsArray()).astype(np.uint16)\n",
    "train_field_id_map = np.hstack([train_field_id_map, np.zeros((train_field_id_map.shape[0], 1))])\n",
    "train_field_id_map = train_field_id_map*(train_mask > 0)\n",
    "# train_field_id = np.unique(train_field_id_map)\n",
    "# train_field_id = train_field_id[train_field_id > 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hKcQKTwDUrw_"
   },
   "source": [
    "Lets plot `train_field_id_map`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1fVbzbaOUIde"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 15))\n",
    "plt.imshow(train_field_id_map)\n",
    "# plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZV313q5TU434"
   },
   "source": [
    "As we can see, `train_field_id_map` contains masks with field ids (pixels corresponding to some particular field #$n$ will have all value $n$)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "juCo_lK3VXpB",
    "outputId": "34db42e3-700f-4d70-d611-9f4ab0871643"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=uint8)"
      ]
     },
     "execution_count": 10,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_mask[train_mask > 15] = 0\n",
    "np.unique(train_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "b1dZ7EYdVqpB"
   },
   "source": [
    "We can also see, that it has only 10 unique values, thus it is most probably (and really is) crop types (out labels for field). Lets visualize them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "p9TaQQkzVPqW"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 15))\n",
    "plt.imshow(train_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K3wOkKPEWKYo"
   },
   "source": [
    "Now it is time to look at satellite image itself. Look at image shape. It should have 13 channels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eKRuw4kaWax4",
    "outputId": "4b430561-42f2-4ee8-c7e8-806b064cc19a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(13, 7467, 8292)"
      ]
     },
     "execution_count": 12,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fields.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KOv9_rpMWeWV"
   },
   "source": [
    "Since we are working with the Sentinel - it captures multispectral images: in our case it has 13 channels. In order to plot the RGB or any other image of the planet surface we should Google tech sheet for Sentinel 2B satellite and search for channels description :[Sentinel 2 Wiki Page](https://en.wikipedia.org/wiki/Sentinel-2#Instruments)\n",
    "\n",
    "Lets plot the NIR channel (7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3X1mlv07ahCO"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 15))\n",
    "plt.imshow(fields[7, :2000, :3000], cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "37SvP2PLcUcg"
   },
   "source": [
    "One can play and plot different channels and convince yourself that different channels have different spatial resolution. Sentinel have 3 types of resolutions (60m, 20m and 10m per pix). They all have been upsampled up to 10m per pixel, but visually one can notice the difference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ADUo4E4lc-PY"
   },
   "source": [
    "Lets visualize agro fields on the image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UY9zeQZZc-vh"
   },
   "outputs": [],
   "source": [
    "from matplotlib import cm\n",
    "from matplotlib.colors import ListedColormap, LinearSegmentedColormap\n",
    "\n",
    "#here we define the color scheme for our filed, such that every crop should have\n",
    "#its own constant color\n",
    "viridis = cm.get_cmap('viridis', 256)\n",
    "train_colors = viridis(np.linspace(0, 1, 256))\n",
    "transp = np.array([0, 0, 0, 0])\n",
    "train_colors[:25, :] = transp\n",
    "train_cmp = ListedColormap(train_colors)\n",
    "\n",
    "plt.figure(figsize=(15, 15))\n",
    "plt.imshow(fields[12, :2000, :3000], cmap='gray')\n",
    "plt.imshow(train_mask[:2000, :3000], alpha=0.5, cmap=train_cmp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Nz_FidMOXWxg"
   },
   "source": [
    "**Task 1 (1 point for each index)**. Plot the NDVI, EVI, NDRE indices for patch of original image for coordinates(pixels) in range for x in [0, 3000] and y in [500:1500]. Just Google the indices, if you do not know exact formulae.\n",
    "\n",
    " *Hint: mind the proper order when slicing the tensor*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TaP27XSaYxw_"
   },
   "source": [
    "Since our dataset is imbalanced, training of even simplest segmentation model could be quite difficult. One of the simplest ways to fight with this issue is to assign the weight for each class in a loss function.\n",
    "\n",
    "**Task 2 (2 points)**. Calculate the number of *pixels* (1 pt) and *fields* (1 pt) for each crop and print them out. This information will help us further in training loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TyDqhsp-dxQG"
   },
   "source": [
    "# Part 2. Download data for training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Jlh7ZQoOT6hg"
   },
   "source": [
    "In order to train the network, we need to split our big satellite image into patches in a such way, that on the one hand, we could process it using neural networks, but on the other hand we could split it on training/validation set with equal contribution of all classes.\n",
    "\n",
    "One way to do this is to crop each field from the original image and resize them (for classification) or not (for segmentation). Another way is to just use patches with constant size. In this work we will use second approach, even though it is much harder to properly split it with equal class proportions. One can try to do that for additional points, but I suggest you to download prepared already split dataset with the link below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "efW4k3VGT535"
   },
   "outputs": [],
   "source": [
    "!wget --no-check-certificate \"https://onedrive.live.com/download?cid=1E2DE865E90D4259&resid=1E2DE865E90D4259%21258624&authkey=AHWYbEgRMryIkKQ\" -O Patches.zip\n",
    "!unzip ./Patches.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ExS53xOS1MlV"
   },
   "source": [
    "The main task is to train networks for classification using fully convolutional network like UNet with Squeeze and Excitation blocks.\n",
    "\n",
    "Downloaded data have the following structure: it has the `train` and `val` folder with the patches for *train* and *validation* purpouses. Each folder contains 2 folders: `images` with the multispectral image patches and `labels` of corresponding crops. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yyEYzwvh2Oo3"
   },
   "source": [
    "Bellow the structure of a network is defined. Ones who are interested can look through it and investigate it. It is basically the UNet with channel-wise attention mechanism called Squeeze and Excitation (SE_block here)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "zCxqgk2P-QhX"
   },
   "outputs": [],
   "source": [
    "class SE_block(nn.Module):\n",
    "    def __init__(self, channels, squeese_rate=1):\n",
    "        super(SE_block, self).__init__()\n",
    "        self.fc1 = nn.Linear(channels, channels//squeese_rate)\n",
    "        self.fc2 = nn.Linear(channels//squeese_rate, channels)\n",
    "    \n",
    "    def forward(self, input):\n",
    "        g_avg = torch.mean(input, [-1, -2])\n",
    "        x = self.fc1(g_avg).relu()\n",
    "        x = self.fc2(x).sigmoid()\n",
    "        return torch.unsqueeze(torch.unsqueeze(x, 2), 3)*input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "VR02P7Yj-QhX"
   },
   "outputs": [],
   "source": [
    "class UNet(nn.Module):\n",
    "    def __init__(self, n_channels, n_classes):\n",
    "        super(UNet, self).__init__()\n",
    "        self.inc = inconv(n_channels, 64)\n",
    "        self.down1 = down(64, 128)\n",
    "        self.se1 = SE_block(128)\n",
    "        self.down2 = down(128, 256)\n",
    "        self.se2 = SE_block(256)\n",
    "        self.down3 = down(256, 256) #was 512\n",
    "        self.se3 = SE_block(256) #was 512\n",
    "        self.down4 = down(512, 512)\n",
    "        self.up1 = up(1024, 256, bilinear=False)\n",
    "        self.up2 = up(512, 128, bilinear=False)\n",
    "        self.up3 = up(256, 64, bilinear=False)\n",
    "        self.up4 = up(128, 64, bilinear=False)\n",
    "        self.outc = outconv(64, n_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x1 = self.inc(x)\n",
    "        x2 = self.down1(x1)\n",
    "        x3 = self.down2(self.se1(x2))\n",
    "        x4 = self.down3(self.se2(x3))\n",
    "        # x5 = self.down4(self.se3(x4))\n",
    "        # x = self.up1(x5, x4)\n",
    "        x = self.up2(x4, x3)\n",
    "        x = self.up3(x, x2)\n",
    "        x = self.up4(x, x1)\n",
    "        x = self.outc(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YgiShdCJ2y7k"
   },
   "source": [
    "Since we have relatively low amount of data, we should use augmentation in order to prevent overfitting and increase generalization a little bit.\n",
    "\n",
    "Below are the parameters for augmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "_RFyphL1-QhY"
   },
   "outputs": [],
   "source": [
    "img_size = 100\n",
    "pad_size = img_size//2\n",
    "shift = img_size//5\n",
    "rot_angle = 60"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "epyH1Eps7tue"
   },
   "source": [
    "It is also important to normalize the data, so do it.\n",
    "\n",
    "\n",
    "**Task 3 (1 point)** Calculate the *mean* and *standart deviation* for each channel in the input image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "iAnc6Nbh-QhY"
   },
   "outputs": [],
   "source": [
    "'''\n",
    "    mean, std: torch.tensor or nd.array of shape (channels)\n",
    "'''\n",
    "# Your code here\n",
    "# mean = ...\n",
    "# std = ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WfOvktGQ88rD"
   },
   "source": [
    "**Task 4 (2 points)** Below you can see the bodies for the functions to additional features represented by different vegetation indices: we suggest you to fill in these functions with the code for generation features, but you can also find some additional features and implement them here in a similar fashion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "Jr4-Tjo--QhY"
   },
   "outputs": [],
   "source": [
    "def generate_NDVI(features):\n",
    "    '''\n",
    "    Arguments:\n",
    "        features: torch.tensor of shape (Chanels, H, W)\n",
    "    Return:\n",
    "        ndvi: torch.tensor of shape (H, W)\n",
    "    '''\n",
    "    # Your code here\n",
    "    # ndvi = ...\n",
    "    return ndvi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "id": "S82GTXfK-QhZ"
   },
   "outputs": [],
   "source": [
    "def generate_EVI(features):\n",
    "    '''\n",
    "    Arguments:\n",
    "        features: torch.tensor of shape (Chanels, H, W)\n",
    "    Return:\n",
    "        evi: torch.tensor of shape (H, W)\n",
    "    '''\n",
    "    # Your code here\n",
    "    # evi = ...\n",
    "    return evi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "id": "g6cSJjth-QhZ"
   },
   "outputs": [],
   "source": [
    "def generate_NDRE(features):\n",
    "    '''\n",
    "    Arguments:\n",
    "        features: torch.tensor of shape (Chanels, H, W)\n",
    "    Return:\n",
    "        ndre: torch.tensor of shape (H, W)\n",
    "    '''\n",
    "    # Your code here\n",
    "    # ndre = ...\n",
    "    return ndre"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "IlPAMTm9-QhZ"
   },
   "outputs": [],
   "source": [
    "def generate_MSAVI(features):\n",
    "    '''\n",
    "    Arguments:\n",
    "        features: torch.tensor of shape (Chanels, H, W)\n",
    "    Return:\n",
    "        msavi: torch.tensor of shape (H, W)\n",
    "    '''\n",
    "    # Your code here\n",
    "    # msavi = ...\n",
    "    return msavi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BsFqEam40vom"
   },
   "source": [
    "Here we gather all generated features in one tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "id": "ctC1nCI0-QhZ"
   },
   "outputs": [],
   "source": [
    "def generate_all_indices(features, u2b4evi=True):\n",
    "    ndvi = generate_NDVI(features)\n",
    "    evi = generate_EVI(features, u2b4evi)\n",
    "    ndre = generate_NDRE(features)\n",
    "    msavi = generate_MSAVI(features)\n",
    "    return torch.cat([features, ndvi.unsqueeze(0), evi.unsqueeze(0), ndre.unsqueeze(0), msavi.unsqueeze(0)], dim=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vmyM9OlL0_Z7"
   },
   "source": [
    "This is the class for `torch` `Dataset` which loads the data from folders and do the preprocessing after loading"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FPUvCEwT1SYF"
   },
   "source": [
    "We have computed `mean` and `std` in previous task, thus we should apply it. Keep in mind that one should apply normalization after all the features are generated and concatenated to the original features.\n",
    "\n",
    "**Task 4 (1 point)** Using function `generate_all_features()` you are suggested to calculate all indices and normalize the satellite image. \n",
    "\n",
    "*Hint: It is recommended to normalize only satellite image part and leave generated indices as it is*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "id": "uvyCD_hh-Qha"
   },
   "outputs": [],
   "source": [
    "class CropFieldsDataset(Dataset):\n",
    "\n",
    "    def __init__(self, images_dir, labels_dir, transform=None):\n",
    "        self.images_dir = images_dir\n",
    "        self.labels_dir = labels_dir\n",
    "        self.transform = transform\n",
    "        i_list = os.listdir(self.images_dir)\n",
    "        self.im_list = []\n",
    "        for image in i_list:\n",
    "            if (image[-3:] == 'npy'):\n",
    "                self.im_list.append(image)\n",
    "                \n",
    "        self.mean = mean[:, None, None]\n",
    "        self.std = std[:, None, None]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.im_list)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_name = os.path.join(self.images_dir, self.im_list[idx])\n",
    "        image = np.load(img_name)#*255\n",
    "        labels_name = os.path.join(self.labels_dir, self.im_list[idx])\n",
    "        labels = np.load(labels_name)\n",
    "        if(self.transform != None):\n",
    "\n",
    "            randangle = np.random.randint(-rot_angle, rot_angle)\n",
    "            sc = tuple(np.random.uniform(0.75, 1, 2))\n",
    "            tf = transforms.AffineTransform(scale = sc)\n",
    "            image = transforms.rotate(image, randangle, mode='reflect')\n",
    "            image = utils.pad(image, ((pad_size, pad_size), (pad_size, pad_size), (0, 0)), mode='reflect')\n",
    "            image = transforms.warp(image, tf, mode='reflect')\n",
    "            labels = transforms.rotate(labels, randangle, mode='reflect', order=0, preserve_range=True)\n",
    "            labels = utils.pad(labels, ((pad_size, pad_size), (pad_size, pad_size)), mode='reflect')\n",
    "            labels = transforms.warp(labels, tf, mode='reflect', order=0)\n",
    "                        \n",
    "            M, N, D = image.shape\n",
    "            randshiftx = np.random.randint(-shift, shift)\n",
    "            randshifty = np.random.randint(-shift, shift)\n",
    "            image = utils.crop(image, (((M + randshiftx - img_size)//2, (M - randshiftx + 1 - img_size)//2), \n",
    "                                       ((N + randshifty - img_size)//2, (N - randshifty + 1 - img_size)//2), (0, 0)))\n",
    "            labels= utils.crop(labels,(((M + randshiftx - img_size)//2, (M - randshiftx + 1 - img_size)//2), \n",
    "                                       ((N + randshifty - img_size)//2, (N - randshifty + 1 - img_size)//2)))           \n",
    "            \n",
    "            \n",
    "        image = torchvision.transforms.ToTensor()(image)\n",
    "        labels = torch.tensor(labels).to(torch.long)\n",
    "\n",
    "        ## Your code here \n",
    "        # image = ...\n",
    "        # norm_image = ...\n",
    "        ## ------------------\n",
    "        return norm_image, labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o-cjbbsKB8dU"
   },
   "source": [
    "Lets make the instance for `Dataset` and make the dataloade (one could try different batch sizes which is appropriate for GPU you use)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "id": "BtwHK6_e-Qha"
   },
   "outputs": [],
   "source": [
    "train_set = CropFieldsDataset(\"./train/images\", \"./train/labels\", transform=True)\n",
    "train_loader = DataLoader(train_set, batch_size=8, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2uLqImLUCTrR"
   },
   "source": [
    "And lets visualize the how augmentation works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4T7gHkAm-Qhb"
   },
   "outputs": [],
   "source": [
    "image, labels = train_set[10]\n",
    "\n",
    "_, ax = plt.subplots(ncols=2)\n",
    "ax[0].imshow(image[12, :, :], cmap='gray')\n",
    "ax[1].imshow(labels[:, :])\n",
    "image.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z6c9Y8yVCcxD"
   },
   "source": [
    "Here comes the validation dataset also"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "id": "rWFmm-fY-Qhd"
   },
   "outputs": [],
   "source": [
    "val_set = CropFieldsDataset(\"./val/images\", \"./val/labels\")\n",
    "val_loader = DataLoader(val_set, batch_size=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xn0rRbjGCjJk"
   },
   "source": [
    "One could look for the expected result, which is generated from validation dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "vQWWd9Ty-Qhe"
   },
   "outputs": [],
   "source": [
    "from matplotlib import cm\n",
    "from matplotlib.colors import ListedColormap, LinearSegmentedColormap\n",
    "\n",
    "viridis = cm.get_cmap('viridis', 256)\n",
    "train_colors = viridis(np.linspace(0, 1, 256))\n",
    "test_colors = viridis(np.linspace(0, 1, 256))\n",
    "transp = np.array([0, 0, 0, 0])\n",
    "train_colors[:25, :] = transp\n",
    "train_cmp = ListedColormap(train_colors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6_kHRt0w-Qhe"
   },
   "outputs": [],
   "source": [
    "# len(train_set)\n",
    "image, labels = val_set[18]\n",
    "\n",
    "print(image.shape)\n",
    "_, ax = plt.subplots(ncols=2, figsize=(10, 10), dpi=100)\n",
    "ax[0].imshow(image[2, :, :], cmap='gray')\n",
    "ax[1].imshow(image[2, :, :], cmap='gray')\n",
    "ax[1].imshow(labels[:, :], cmap=train_cmp)\n",
    "ax[0].set_axis_off()\n",
    "ax[1].set_axis_off()\n",
    "ax[0].set_title('Band 3')\n",
    "ax[1].set_title('Crop map')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WMfnDODUCrQb"
   },
   "source": [
    "Now we are going to train the network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "id": "EHV_Y6ey-Qhe"
   },
   "outputs": [],
   "source": [
    "if(torch.cuda.is_available()):\n",
    "    device = torch.device('cuda:0')\n",
    "else:\n",
    "    device = torch.device('cpu')\n",
    "\n",
    "# Here we initialize the network: UNet(input_features, n_classes)\n",
    "# If you generate additional features, you should put here number of channels\n",
    "# of your final input (13 + number_of_generated_indices)\n",
    "model = UNet(17, 10).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "id": "munZdXfM-Qhf"
   },
   "outputs": [],
   "source": [
    "def get_lr(optimizer):\n",
    "    for param_group in optimizer.param_groups:\n",
    "        return param_group['lr']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SaAmFnXtDUzh"
   },
   "source": [
    "**Task 5 (1 point)** Complete training loop with missing parts (look at the comments for guidance) and \n",
    "\n",
    "**Task 6 (2 points)** Train your network. You will get full score (**2 points**) if the $accuracy$ on validation $\\geq 63\\%$, **1 point** if it will be in range $50\\% \\leq accuracy < 63\\%$, and **no points** if $accuracy < 50\\%$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Af57Az-8-Qhf"
   },
   "outputs": [],
   "source": [
    "train_weights = torch.tensor([0, 0.14852941, 0, 0.28836312, 0.04097108, 0.07013557, 0.10067904, 0.07735483, 0.02633347, 0]).to(device)\n",
    "val_weights = torch.tensor([0, 0.20184716, 0, 0.23679423, 0.03584126, 0.06965259, 0.09625471, 0.07470775, 0.02500757, 0]).to(device)\n",
    "optimizer = Adam(model.parameters(), lr = 0.01, weight_decay=0.0001)\n",
    "\n",
    "train_criterion = nn.CrossEntropyLoss(ignore_index=0, weight=train_weights, reduction='mean').to(device)\n",
    "val_criterion = nn.CrossEntropyLoss(ignore_index=0, weight=val_weights, reduction='mean').to(device)\n",
    "\n",
    "lr_sch = lr_scheduler.StepLR(optimizer, 1, gamma=0.977)\n",
    "epochs_num = 350\n",
    "\n",
    "loss_list = []\n",
    "acc_list = []\n",
    "\n",
    "best_acc = 0.0 #0.628\n",
    "\n",
    "for epoch in range(epochs_num):  # loop over the dataset multiple times\n",
    "\n",
    "    running_train_loss = 0.0\n",
    "    train_accuracy_pix = 0\n",
    "    train_sum_pix = 0\n",
    "    model.train()\n",
    "    for data in tqdm(train_loader):\n",
    "        # get the inputs; data is a list of [inputs, labels]\n",
    "        inputs, labels = data\n",
    "        # transfer data to the device\n",
    "        # inputs = ...\n",
    "        # labels = ...\n",
    "        # zero the parameter gradients\n",
    "        # ...\n",
    "        # make forward pass, compute the loss, make backward pass and optimization step\n",
    "        # outputs = ...\n",
    "        # loss = ...\n",
    "        # ...\n",
    "        # ...\n",
    "\n",
    "        # print statistics\n",
    "        running_train_loss += loss.item()\n",
    "        \n",
    "        result = F.softmax(outputs, dim=1).detach().cpu().numpy()\n",
    "        pred = np.argmax(result, axis=1)\n",
    "        labels = labels.cpu().numpy()\n",
    "        fl_labels = labels.flatten()\n",
    "        train_accuracy_pix += np.sum(pred.flatten()[fl_labels > 0] == fl_labels[fl_labels > 0])\n",
    "        train_sum_pix += np.sum(labels > 0)\n",
    "#     train_writer.add_scalar('Loss', running_train_loss, global_step=epoch)\n",
    "\n",
    "    running_train_loss /= len(train_set)\n",
    "    loss_list.append(running_train_loss)\n",
    "    acc_list.append(train_accuracy_pix/train_sum_pix)\n",
    "    \n",
    "    model.eval()\n",
    "    running_val_loss = 0.0\n",
    "    val_accuracy_pix = 0\n",
    "    val_sum_pix = 0\n",
    "    \n",
    "    for data in tqdm(val_loader):\n",
    "        inputs, labels = data\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "        with torch.no_grad():\n",
    "            outputs = model(inputs)\n",
    "            loss = val_criterion(outputs, labels)\n",
    "        \n",
    "        running_val_loss += loss.item()\n",
    "        \n",
    "        result = F.softmax(outputs, dim=1).detach().cpu().numpy()\n",
    "        pred = np.argmax(result, axis=1)\n",
    "        labels = labels.cpu().numpy()\n",
    "        fl_labels = labels.flatten()\n",
    "        val_accuracy_pix += np.sum(pred.flatten()[fl_labels > 0] == fl_labels[fl_labels > 0])\n",
    "        val_sum_pix += np.sum(labels > 0)\n",
    "        \n",
    "    running_val_loss /= len(val_set)\n",
    "\n",
    "    # train_writer.add_scalar('Accuracy', train_accuracy_pix/train_sum_pix, global_step=epoch)\n",
    "    # train_writer.add_scalar(\"Loss\",  running_train_loss, global_step=epoch)\n",
    "    \n",
    "    # val_writer.add_scalar('Accuracy', val_accuracy_pix/val_sum_pix, global_step=epoch)\n",
    "    # val_writer.add_scalar(\"Loss\", running_val_loss, global_step=epoch)\n",
    "    print(\"Epoch: {0:3d} | LR: {1}\\n    Train Loss:  {2:4f} \\n    Train Acc: \\t {3:4.1f}%\\n    Val Loss:    {4:4f} \\n    Val Acc: \\t {5:4.1f}%\".format(\n",
    "                epoch+1, get_lr(optimizer), \n",
    "                running_train_loss, \n",
    "                100*train_accuracy_pix/train_sum_pix,\n",
    "                running_val_loss,\n",
    "                100*val_accuracy_pix/val_sum_pix))\n",
    "    \n",
    "    if (val_accuracy_pix/val_sum_pix > best_acc):\n",
    "        best_acc = val_accuracy_pix/val_sum_pix\n",
    "        torch.save(model.state_dict(), 'best_model.pth')\n",
    "    \n",
    "    lr_sch.step()\n",
    "        \n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_86qYlsypUem"
   },
   "source": [
    "Next cells are just for validation\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4Srzo8TJ-Qhh"
   },
   "outputs": [],
   "source": [
    "model.eval()\n",
    "val_accuracy_pix = 0\n",
    "val_sum_pix = 0\n",
    "\n",
    "correct_class_pix = np.zeros((10), dtype=np.int)\n",
    "class_pix = np.zeros((10), dtype=np.int)\n",
    "\n",
    "preds_list = []\n",
    "labels_list = []\n",
    "\n",
    "for data in tqdm(val_loader):\n",
    "    inputs, labels = data\n",
    "    inputs = inputs.to(device)\n",
    "    labels = labels.to(device)\n",
    "    with torch.no_grad():\n",
    "        outputs = model(inputs)\n",
    "        \n",
    "    result = F.softmax(outputs, dim=1).detach().cpu().numpy()\n",
    "    pred = np.argmax(result, axis=1)\n",
    "    preds_list.append(pred.flatten())\n",
    "    labels = labels.cpu().numpy()\n",
    "    fl_labels = labels.flatten()\n",
    "    labels_list.append(fl_labels)\n",
    "    val_accuracy_pix += np.sum(pred.flatten()[fl_labels > 0] == fl_labels[fl_labels > 0])\n",
    "    val_sum_pix += np.sum(labels > 0)\n",
    "    for i in range(10):\n",
    "        class_pix[i] += np.sum(labels == i)\n",
    "        correct_class_pix[i] += np.sum((pred == i)*(labels == i))\n",
    "\n",
    "class_acc = correct_class_pix/class_pix\n",
    "\n",
    "pix_df = pd.DataFrame(data=np.vstack([correct_class_pix, class_pix]), \n",
    "                      index=['correct_pix', 'pix'], \n",
    "                      columns=np.arange(10))\n",
    "acc_df = pd.DataFrame(data=np.round(class_acc[None, :], 3), index=['accuracy'], columns=np.arange(10))\n",
    "\n",
    "print('Accuracy', val_accuracy_pix/val_sum_pix)\n",
    "\n",
    "display(pix_df)\n",
    "display(acc_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "id": "sNp8Aivp-Qhh"
   },
   "outputs": [],
   "source": [
    "all_preds = np.hstack(preds_list)\n",
    "all_labels = np.hstack(labels_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AmF942STqzDx"
   },
   "source": [
    "Cell below are for visualization of correctness of the models prediction (green means correct, red - incorrect). Just launch all remaining cells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rs86Fpj3-Qhi"
   },
   "outputs": [],
   "source": [
    "images = []\n",
    "results = []\n",
    "llabels = []\n",
    "with torch.no_grad():\n",
    "    for i, data in tqdm(enumerate(val_loader)):\n",
    "        # get the inputs; data is a list of [inputs, labels]\n",
    "        inputs, labels = data\n",
    "        images.append(inputs.numpy())\n",
    "        llabels.append(labels.numpy())\n",
    "        inputs = inputs.to(device)\n",
    "\n",
    "        # forward\n",
    "        outputs = model(inputs)\n",
    "        result = F.softmax(outputs, dim=1).detach().cpu().numpy()\n",
    "        pred = np.argmax(result, axis=1)\n",
    "        results.append(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "id": "vxTa72ke-Qhj"
   },
   "outputs": [],
   "source": [
    "from matplotlib import cm\n",
    "from matplotlib.colors import ListedColormap, LinearSegmentedColormap\n",
    "\n",
    "viridis = cm.get_cmap('viridis', 256)\n",
    "train_colors = viridis(np.linspace(0, 1, 256))\n",
    "test_colors = viridis(np.linspace(0, 1, 256))\n",
    "transp = np.array([0, 0, 0, 0])\n",
    "train_colors[:24, :] = transp\n",
    "train_cmp = ListedColormap(train_colors)\n",
    "pink = np.array([248/256, 24/256, 148/256, 1])\n",
    "test_colors[:128, :] = transp\n",
    "test_colors[128:, :] = pink\n",
    "test_cmp = ListedColormap(test_colors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "id": "Os5aSTDA-Qhj"
   },
   "outputs": [],
   "source": [
    "test_scale = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-nUlm167-Qhj"
   },
   "outputs": [],
   "source": [
    "for i in range(len(val_set)):\n",
    "    right = (results[i][0, ...] == llabels[i][0, ...]).astype(int) - (results[i][0, ...] != llabels[i][0, ...]).astype(int)\n",
    "    pic1 = results[i][0, ...]*(llabels[i][0, ...] > 0)\n",
    "    pic2 = llabels[i][0, ...]#.numpy()\n",
    "    pic1[0, :10] = test_scale\n",
    "    pic2[0, :10] = test_scale\n",
    "\n",
    "    _, ax = plt.subplots(ncols=3, figsize=(12, 4))\n",
    "    ax[0].imshow(images[i][0, 3, ...], cmap='gray')\n",
    "    fig1 = ax[0].imshow(pic1, alpha=1, cmap=train_cmp)\n",
    "    ax[0].set_title('Model Prediction')\n",
    "    fig2 = ax[1].imshow(pic2, cmap=viridis)\n",
    "    ax[1].set_title('Ground Truth')\n",
    "    ax[2].imshow(right*(llabels[i][0, ...] > 0), cmap='RdYlGn')\n",
    "    ax[2].set_title('Correctness')\n",
    "    # plt.colorbar(fig1, ax=ax[0])\n",
    "    # plt.colorbar(fig2, ax=ax[1])\n",
    "    ax[0].xaxis.set_visible(False)\n",
    "    ax[0].yaxis.set_visible(False)\n",
    "    ax[1].xaxis.set_visible(False)\n",
    "    ax[1].yaxis.set_visible(False)\n",
    "    ax[2].xaxis.set_visible(False)\n",
    "    ax[2].yaxis.set_visible(False)\n",
    "    \n",
    "#     plt.savefig('Figures/fig_{}.jpg'.format(i), dpi=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "g6ETVnG-pKVu"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Satellite_Images_Hometask.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}