2.Datasets/Dataset_class_and_utils.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Dataset class and utils.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "j525m489D5q-"
      },
      "source": [
        "import numpy as np\r\n",
        "import math\r\n",
        "import scipy\r\n",
        "import pandas as pd\r\n",
        "import PIL\r\n",
        "import gdal\r\n",
        "import matplotlib.pyplot as plt\r\n",
        "plt.style.use('ggplot')\r\n",
        "import sys, os\r\n",
        "from pathlib import Path\r\n",
        "import time\r\n",
        "import random\r\n",
        "import collections, functools, operator\r\n",
        "import csv\r\n",
        "import subprocess\r\n",
        "import datetime\r\n",
        "\r\n",
        "from osgeo import gdal,osr\r\n",
        "from gdalconst import *\r\n",
        "import subprocess\r\n",
        "from osgeo.gdalconst import GA_Update\r\n",
        "\r\n",
        "import torch\r\n",
        "import torch.nn as nn\r\n",
        "from torch.utils.data import Dataset\r\n",
        "from torch.autograd import Variable\r\n",
        "from torch.nn import Linear, ReLU, CrossEntropyLoss, MSELoss, Sequential, Conv2d, MaxPool2d, Module, Softmax, BatchNorm2d, Dropout, Sigmoid\r\n",
        "from torch.optim import Adam, SGD\r\n",
        "from torchvision import transforms, utils\r\n",
        "\r\n",
        "import skimage\r\n",
        "from skimage import io, transform\r\n",
        "import sklearn\r\n",
        "import sklearn.metrics\r\n",
        "from sklearn.feature_extraction import image\r\n",
        "from sklearn import svm"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cLJFffF4MSyF"
      },
      "source": [
        "# Satellite dataset class"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "frX8AE5VEJAW"
      },
      "source": [
        "class SatelliteDataset(Dataset):\r\n",
        "    \"\"\"Satellite dataset.\"\"\"\r\n",
        "\r\n",
        "    def __init__(self, dirs, normalization_options = None, patches_options = None, readWhileRunning=False, readFromPatches=False, cropClassification=False):\r\n",
        "        \"\"\"\r\n",
        "        Args:\r\n",
        "            dirs (dict): Dictionary of lists of directories of scene \"input\" and \"target\" bands.\r\n",
        "            normalization_options (obj): Object normalization_options with information about normalization technique to be applied.\r\n",
        "            patches_options (obj): Object patches_options with information about the patching options (size, step, etc)\r\n",
        "            readWhileRunning (bool): Whether you want to read online files or local files\r\n",
        "            readFromPatches (bool): Whether you want to read from patches in a numpy array data structure.\r\n",
        "            cropClassification (bool): Whether you want to perform crop classification. It will vary the reading workflow.\r\n",
        "        \"\"\"\r\n",
        "        self.dirs = dirs\r\n",
        "        self.normalization_options = normalization_options\r\n",
        "        self.patches_options = patches_options\r\n",
        "        self.readWhileRunning = readWhileRunning\r\n",
        "        self.readFromPatches = readFromPatches\r\n",
        "        self.data_augm = False\r\n",
        "        self.cropClassification = cropClassification\r\n",
        "\r\n",
        "        self.bandsToGet = ['band1', 'band2', 'band3', 'band4', 'band5', 'band7',\r\n",
        "                           '_B02_', '_B03_', '_B04_', '_B05_', '_B06_', '_B07_', '_B08_', '_B09_', '_B10_',\r\n",
        "                           \r\n",
        "                           'B008', 'B009', 'B010', 'B011', 'B012', 'B013', 'B014', 'B015', 'B016', 'B017',\r\n",
        "                           'B018', 'B019', 'B020', 'B021', 'B022', 'B023', 'B024', 'B025', 'B026', 'B027',\r\n",
        "                           'B028', 'B029', 'B030', 'B031', 'B032', 'B033', 'B034', 'B035', 'B036', 'B037',\r\n",
        "                           'B038', 'B039', 'B040', 'B041', 'B042', 'B043', 'B044', 'B045', 'B046', 'B047',\r\n",
        "                           'B048', 'B049', 'B050', 'B051', 'B052', 'B053', 'B054', 'B055', 'B056', 'B057',\r\n",
        "                           \r\n",
        "                           'B079', 'B080', 'B081', 'B082', 'B083', 'B084', 'B085', 'B086', 'B087', \r\n",
        "                           'B088', 'B089', 'B090', 'B091', 'B092', 'B093', 'B094', 'B095', 'B096', 'B097', \r\n",
        "                           'B098', 'B099', 'B100', 'B101', 'B102', 'B103', 'B104', 'B105', 'B106', 'B107', \r\n",
        "                           'B108', 'B109', 'B110', 'B111', 'B112', 'B113', 'B114', 'B115', 'B116', 'B117', \r\n",
        "                           'B118', 'B119', 'B120',\r\n",
        "                           \r\n",
        "                           'B131', 'B132', 'B133', 'B134', 'B135', 'B136', 'B137', 'B138', 'B139', 'B140', \r\n",
        "                           'B141', 'B142', 'B143', 'B144', 'B145', 'B146', 'B147', 'B148', 'B149', 'B150', \r\n",
        "                           'B151', 'B152', 'B153', 'B154', 'B155', 'B156', 'B157', 'B158', 'B159', 'B160', \r\n",
        "                           'B161', 'B162', 'B163', 'B164', 'B165',\r\n",
        "                           \r\n",
        "                           'B181', 'B182', 'B183', 'B184', 'B185', 'B186', 'B187', 'B188', 'B189', 'B190', \r\n",
        "                           'B191', 'B192', 'B193', 'B194', 'B195', 'B196', 'B197', 'B198', 'B199', 'B200', \r\n",
        "                           'B201', 'B202', 'B203', 'B204', 'B205', 'B206', 'B207', 'B208', 'B209', 'B210', \r\n",
        "                           'B211', 'B212', 'B213', 'B214', 'B215', 'B216', 'B217', 'B218', 'B219', 'B220', \r\n",
        "                           'B221', 'B222', 'B223']\r\n",
        "\r\n",
        "        if self.readWhileRunning == False:\r\n",
        "            self.dataset = self.read()\r\n",
        "        \r\n",
        "    def readAllPatchesFromScene(self):\r\n",
        "        scene = {}\r\n",
        "        scene['input'] = []\r\n",
        "        for fname in sorted(os.listdir(self.dirs['input'][0])):\r\n",
        "            if any(b in fname for b in self.bandsToGet) and 'modified3' in fname:\r\n",
        "                band = gdal.Open(os.path.join(self.dirs['input'][0], fname)).ReadAsArray()\r\n",
        "                band = band.astype('float32')\r\n",
        "                band /= 65535.0\r\n",
        "                band = skimage.util.shape.view_as_windows(band, (64,64), step=64)\r\n",
        "                scene['input'].append(band)\r\n",
        "        scene['input'] = np.array(scene['input']).astype('float32')\r\n",
        "        print(scene['input'].shape)\r\n",
        "        scene['input'] = scene['input'].reshape(9,55*18,64,64).transpose(1,0,2,3)\r\n",
        "        print(scene['input'].shape)\r\n",
        "        \r\n",
        "        scene['target'] = []\r\n",
        "        for fname in sorted(os.listdir(self.dirs['target'][0])):\r\n",
        "            if any(b in fname for b in self.bandsToGet):\r\n",
        "                band = gdal.Open(os.path.join(self.dirs['target'][0], fname)).ReadAsArray()\r\n",
        "                band = band.astype('float32')\r\n",
        "                band /= 65535.0\r\n",
        "                band = skimage.util.shape.view_as_windows(band, (64,64), step=64)\r\n",
        "                scene['target'].append(band)\r\n",
        "        scene['target'] = np.array(scene['target']).astype('float32')\r\n",
        "        print(scene['target'].shape)\r\n",
        "        scene['target'] = scene['target'].reshape(170,55*18,64,64).transpose(1,0,2,3)\r\n",
        "        print(scene['target'].shape)\r\n",
        "\r\n",
        "        return scene\r\n",
        "\r\n",
        "\r\n",
        "    def __len__(self):\r\n",
        "        if self.readWhileRunning == False and self.readFromPatches == False:\r\n",
        "            return self.dataset['input'].shape[0]\r\n",
        "        else:\r\n",
        "            return len(os.listdir(self.dirs['input']))\r\n",
        "\r\n",
        "    def __getitem__(self, idx): # Now it returns whole batch\r\n",
        "        if self.readWhileRunning:\r\n",
        "            dataset = self.getitem_online(idx)\r\n",
        "        else:\r\n",
        "            dataset = self.getitem_local(idx)\r\n",
        "\r\n",
        "        dataset['input'] = torch.from_numpy(dataset['input'])\r\n",
        "        dataset['target'] = torch.from_numpy(dataset['target'])\r\n",
        "        \r\n",
        "        if self.data_augm:\r\n",
        "            if bool(random.getrandbits(1)):\r\n",
        "                dataset['input'] = transforms.RandomHorizontalFlip(1).forward(dataset['input'])\r\n",
        "                dataset['target'] = transforms.RandomHorizontalFlip(1).forward(dataset['target'])\r\n",
        "            if bool(random.getrandbits(1)):\r\n",
        "                dataset['input'] = transforms.RandomVerticalFlip(1).forward(dataset['input'])\r\n",
        "                dataset['target'] = transforms.RandomVerticalFlip(1).forward(dataset['target'])\r\n",
        "        return dataset\r\n",
        "    \r\n",
        "    def getitem_online(self, idx):\r\n",
        "        fnameInput = sorted(os.listdir(self.dirs['input']))[idx]\r\n",
        "        fnameTarget = sorted(os.listdir(self.dirs['target']))[idx]\r\n",
        "        dataset = {'input': np.load(os.path.join(self.dirs['input'], fnameInput)),\r\n",
        "                  'target': np.load(os.path.join(self.dirs['target'], fnameTarget))}\r\n",
        "        if 'crop' in self.dirs.keys():\r\n",
        "            dataset['crop'] = 'Not loaded'\r\n",
        "        return dataset\r\n",
        "    def getitem_local(self, idx):\r\n",
        "        dataset = {'input': self.dataset['input'][idx],\r\n",
        "                'target': self.dataset['target'][idx]}\r\n",
        "        if 'crop' in self.dirs.keys():\r\n",
        "            dataset['crop'] = self.dataset['crop'][idx]\r\n",
        "        return dataset\r\n",
        "    \r\n",
        "\r\n",
        "\r\n",
        "    def read(self):\r\n",
        "        if self.cropClassification == True:\r\n",
        "            dataset = self.readAllPatchesFromScene()\r\n",
        "        elif self.readFromPatches == True:\r\n",
        "            dataset = self.read_patches()\r\n",
        "        else:\r\n",
        "            dataset = self.read_scenes()\r\n",
        "        return dataset\r\n",
        "\r\n",
        "    def read_patches(self):\r\n",
        "        input=[]\r\n",
        "        target=[]\r\n",
        "        for fnameInput, fnameTarget in zip(sorted(os.listdir(self.dirs['input'])), sorted(os.listdir(self.dirs['target']))):\r\n",
        "            input.append(np.load(os.path.join(self.dirs['input'],fnameInput)))\r\n",
        "            target.append(np.load(os.path.join(self.dirs['target'],fnameTarget)))\r\n",
        "        dataset = {'input': input, 'target': target}\r\n",
        "        return dataset\r\n",
        "\r\n",
        "    def read_scenes(self):\r\n",
        "        dataset = {}\r\n",
        "        dataset['input'] = np.array([])\r\n",
        "        dataset['target'] = np.array([])\r\n",
        "        dataset['crop'] = np.array([])\r\n",
        "\r\n",
        "        for dirNameInput, dirNameTarget in zip(dirs['input'], dirs['target']):\r\n",
        "            scene = {}\r\n",
        "            scene['input'] = []\r\n",
        "            scene['target'] = []\r\n",
        "            scene['crop'] = []\r\n",
        "            for fname in sorted(os.listdir(dirNameInput)):\r\n",
        "                if any(b in fname for b in self.bandsToGet) and 'modified3' in fname:\r\n",
        "                    band = gdal.Open(os.path.join(dirNameInput, fname)).ReadAsArray()\r\n",
        "                    scene['input'].append(band)\r\n",
        "            for fname in sorted(os.listdir(dirNameTarget)):\r\n",
        "                if any(b in fname for b in self.bandsToGet):\r\n",
        "                    band = gdal.Open(os.path.join(dirNameTarget, fname)).ReadAsArray()\r\n",
        "                    scene['target'].append(band)\r\n",
        "            sceneCropClass = None\r\n",
        "            if 'crop' in self.dirs.keys():\r\n",
        "                sceneCropClass = gdal.Open(self.dirs['crop']).ReadAsArray()\r\n",
        "\r\n",
        "            scene['input'] = np.array(scene['input']).astype('float32')\r\n",
        "            scene['target'] = np.array(scene['target']).astype('float32')\r\n",
        "\r\n",
        "            cloudMask = None\r\n",
        "            if self.patches_options and self.patches_options.evadeClouds:\r\n",
        "                cloudMask = self.get_cloudMask(scene)\r\n",
        "            if self.normalization_options:\r\n",
        "                scene = self.normalize(scene)\r\n",
        "            if self.patches_options:\r\n",
        "                scene = self.extract_patches(scene, cloudMask, sceneCropClass)\r\n",
        "                cloudMask = None\r\n",
        "            \r\n",
        "            if len(scene['input'].shape) >= 3 and len(scene['target'].shape) >= 3:\r\n",
        "                dataset['input'] = np.concatenate((dataset['input'], scene['input'])) if dataset['input'].size else scene['input']\r\n",
        "                dataset['target'] = np.concatenate((dataset['target'], scene['target'])) if dataset['target'].size else scene['target']\r\n",
        "                if 'crop' in self.dirs.keys():\r\n",
        "                    dataset['crop'] = np.concatenate((dataset['crop'], scene['crop'])) if dataset['crop'].size else scene['crop']\r\n",
        "\r\n",
        "\r\n",
        "            print(dataset['input'].shape)\r\n",
        "            print(dataset['target'].shape)\r\n",
        "            if 'crop' in self.dirs.keys():\r\n",
        "                print(dataset['crop'].shape)\r\n",
        "            print('')\r\n",
        "\r\n",
        "        return dataset\r\n",
        "\r\n",
        "\r\n",
        "\r\n",
        "\r\n",
        "\r\n",
        "    def visualize(self, datasetName, rBand=None, gBand=None, bBand=None, patched=False):\r\n",
        "        #if len(self.sample[datasetName].shape) == 4:\r\n",
        "        #    raise TypeError('Dataset is in patches. Visualization not available.')\r\n",
        "\r\n",
        "        if datasetName == 'all':\r\n",
        "            fig, axs = plt.subplots(1,2)\r\n",
        "            fig.set_figwidth(24)\r\n",
        "            fig.set_figheight(15)\r\n",
        "            for i, (nameset, scene) in enumerate(self.dataset.items()):\r\n",
        "                if nameset == 'input':\r\n",
        "                    rBand = 3-1\r\n",
        "                    gBand = 2-1\r\n",
        "                    bBand = 1-1\r\n",
        "                elif nameset == 'target':\r\n",
        "                    rBand = 29-8\r\n",
        "                    gBand = 20-8\r\n",
        "                    bBand = 12-8\r\n",
        "\r\n",
        "                R = self.dataset[nameset][rBand]\r\n",
        "                G = self.dataset[nameset][gBand]\r\n",
        "                B = self.dataset[nameset][bBand]\r\n",
        "                img = np.dstack((R, G, B))\r\n",
        "                axs[i].imshow(img)\r\n",
        "            plt.show()\r\n",
        "        else:\r\n",
        "\r\n",
        "            if datasetName == 'input':\r\n",
        "                rBand -= 1\r\n",
        "                gBand -= 1\r\n",
        "                bBand -= 1\r\n",
        "            elif datasetName == 'target':\r\n",
        "                rBand -= 7+1\r\n",
        "                gBand -= 7+1\r\n",
        "                bBand -= 7+1\r\n",
        "\r\n",
        "            if not patched:\r\n",
        "                R = self.dataset[datasetName][rBand]\r\n",
        "                G = self.dataset[datasetName][gBand]\r\n",
        "                B = self.dataset[datasetName][bBand]\r\n",
        "                img = np.dstack((R,G,B))\r\n",
        "                plt.figure(num=None, figsize=(15, 12), dpi=80)\r\n",
        "                plt.imshow(img)\r\n",
        "                plt.show()\r\n",
        "\r\n",
        "            elif patched:\r\n",
        "                fig, axs = plt.subplots(1,10)\r\n",
        "                fig.suptitle(datasetName)\r\n",
        "                fig.set_figwidth(15)\r\n",
        "                for i in range(10):\r\n",
        "                    n = random.randint(0, self.dataset[datasetName].shape[0]-1)\r\n",
        "                    R = self.dataset[datasetName][n][rBand]\r\n",
        "                    G = self.dataset[datasetName][n][gBand]\r\n",
        "                    B = self.dataset[datasetName][n][bBand]\r\n",
        "                    \r\n",
        "                    img = np.dstack((R, G, B))\r\n",
        "                    axs[i].imshow(img)\r\n",
        "                plt.show()\r\n",
        "\r\n",
        "    def get_cloudMask(self, scene):\r\n",
        "        cloudMask = {}\r\n",
        "        cloudMask['input'] = scene['input'][3-1]\r\n",
        "        cloudMask['target'] = scene['target'][29-8]\r\n",
        "        cloudMask = self.bandAdaptive_norm(cloudMask)\r\n",
        "        return cloudMask\r\n",
        "\r\n",
        "    #===================================================#\r\n",
        "    #============== EXTRACT PATCHES UTILS ==============#\r\n",
        "    #===================================================#\r\n",
        "\r\n",
        "    def extract_patches(self, scene, cloudMask=None, sceneCropClass=None):\r\n",
        "        patchesInput = []\r\n",
        "        patchesTarget = []\r\n",
        "        patchesCrop = []\r\n",
        "\r\n",
        "        window = self.patches_options.window\r\n",
        "        step = self.patches_options.step\r\n",
        "\r\n",
        "        stepNextRow = False\r\n",
        "        i = 0\r\n",
        "        while i < (scene['input'][4-2].shape[0] - window): # 3-1 for landsat7\r\n",
        "            j=0\r\n",
        "            while j < (scene['input'][4-2].shape[1] - window): # Rband - 2 because of indexing properly (indexing starts at 0 but first band is the 2nd as n1 is PAN)\r\n",
        "                patchInput = scene['input'][:, i:i+window, j:j+window] \r\n",
        "                patchTarget = scene['target'][29-8][i:i+window, j:j+window] # Rband - 8 because of indexing properly (same as above and 8 or 7 initial bands from hyperion are removed)\r\n",
        "                if self.meaningful_patch(patchInput, patchTarget, cloudMask):\r\n",
        "                    if not self.patches_options.miniPatches:\r\n",
        "                        patchesInput.append(patchInput)\r\n",
        "                        patchesTarget.append(scene['target'][:, i:i+window, j:j+window])\r\n",
        "                        if 'crop' in self.dirs.keys():\r\n",
        "                            patchesCrop.append(sceneCropClass[i:i+window, j:j+window])\r\n",
        "                    elif self.patches_options.miniPatches:\r\n",
        "                        patchesInput.append(self.sliding_window(patchInput).transpose(1,2,0,3,4))\r\n",
        "                        patchesTarget.append(scene['target'][:, i+2:i+window-2, j+2:j+window-2].transpose(1,2,0))\r\n",
        "                        if 'crop' in self.dirs.keys():\r\n",
        "                            patchesCrop.append(sceneCropClass[i+2:i+window-2, j+2:j+window-2])\r\n",
        "                    j += step-1\r\n",
        "                    stepNextRow = True\r\n",
        "                j += 1\r\n",
        "            if stepNextRow:\r\n",
        "                i += step-1\r\n",
        "                stepNextRow = False\r\n",
        "            i += 1\r\n",
        "        \r\n",
        "        scene['input'] = np.array(patchesInput).astype('float32')\r\n",
        "        scene['target'] = np.array(patchesTarget).astype('float32')\r\n",
        "        if 'crop' in self.dirs.keys():\r\n",
        "            scene['crop'] = np.array(patchesCrop)\r\n",
        "        return scene\r\n",
        "    \r\n",
        "    def sliding_window(self, patch):\r\n",
        "        miniPatch = []\r\n",
        "        for bandPatch in patch:\r\n",
        "            miniPatch.append(skimage.util.shape.view_as_windows(bandPatch, (5,5), step=1))\r\n",
        "        return np.array(miniPatch)\r\n",
        "\r\n",
        "    def meaningful_patch(self, patchInput, patchTarget, cloudMask=None):\r\n",
        "        if self.patches_options.evadeClouds == False:\r\n",
        "            if np.any(patchTarget == 0.0) or np.any(patchTarget == 1.0) or np.any(patchInput == 0.0) or np.any(patchInput == 1.0):\r\n",
        "                return False\r\n",
        "        elif self.patches_options.evadeClouds == True:\r\n",
        "            if 'nan' in patchTarget or 0. in patchTarget or patchTarget.mean() < 0.3 or patchTarget.mean() > 0.7 or patchTarget.std() > 0.25:\r\n",
        "                #print('target', patchTarget.mean(), patchTarget.std())\r\n",
        "                return False\r\n",
        "            if 'nan' in patchInput or patchInput.mean() > 0.8 or patchInput.mean() < 0.3 or patchInput.std() > 0.25:\r\n",
        "                #print('input', patchInput.mean(), patchInput.std())\r\n",
        "                return False\r\n",
        "        \r\n",
        "        return True\r\n",
        "\r\n",
        "\r\n",
        "    #===================================================#\r\n",
        "    #=============== NORMALIZATION UTILS ===============#\r\n",
        "    #===================================================#\r\n",
        "    def normalize(self, dataset):\r\n",
        "        norm_type = self.normalization_options.norm_type\r\n",
        "        if norm_type == 'bandAdaptive_norm':\r\n",
        "            return self.bandAdaptive_norm(dataset)\r\n",
        "        elif norm_type == 'sceneAdaptive_norm':\r\n",
        "            return self.sceneAdaptive_norm(dataset)\r\n",
        "        elif norm_type == 'sensorAdaptive_norm':\r\n",
        "            return self.sensorAdaptive_norm(dataset)\r\n",
        "        elif norm_type == 'fullRange_norm':\r\n",
        "            return self.fullRange_norm(dataset)\r\n",
        "        elif norm_type == 'none':\r\n",
        "            return dataset\r\n",
        "        else:\r\n",
        "            return Exception('No norm_type set')\r\n",
        "\r\n",
        "    def sensorAdaptive_norm(self, dataset):\r\n",
        "        dataset['input'] = self.band_max_values(dataset['input'], 8)\r\n",
        "        dataset['input'] /= 3000.0\r\n",
        "\r\n",
        "        dataset['target'] = self.band_max_values(dataset['target'], 8)\r\n",
        "        #dataset['target'] -= 100.0\r\n",
        "        #dataset['target'] /= 2000.0 - 100.0\r\n",
        "        dataset['target'] /= 2000.0\r\n",
        "        dataset['target'][dataset['target'] < 0] = 0.0\r\n",
        "        dataset['target'][dataset['target'] > 1] = 1.0\r\n",
        "        \r\n",
        "        return dataset\r\n",
        "\r\n",
        "\r\n",
        "    def sceneAdaptive_norm(self, dataset):\r\n",
        "        pass\r\n",
        "\r\n",
        "\r\n",
        "    def fullRange_norm(self, dataset):\r\n",
        "        dataset['input'] = self.band_max_values(dataset['input'], 8)\r\n",
        "        dataset['input'] /= 65535.0\r\n",
        "        dataset['target'] = self.band_max_values(dataset['target'], 8)\r\n",
        "        dataset['target'] /= 65535.0\r\n",
        "        '''\r\n",
        "        for nameset, scenes in dataset.items():\r\n",
        "            if isinstance(dataset[nameset], list): # Dataset with a list of scenes\r\n",
        "                for scene_i, scene in enumerate(scenes):\r\n",
        "                    dataset[nameset][scene_i] = self.band_max_values(scene, 8)\r\n",
        "                    dataset[nameset][scene_i] = dataset[nameset][scene_i].astype('float') / 65535.0\r\n",
        "\r\n",
        "            elif len(dataset[nameset].shape) == 3: # Dataset with only 1 scene (with some .tif bands)\r\n",
        "                dataset[nameset] = self.band_max_values(scenes, 8)\r\n",
        "                dataset[nameset] = dataset[nameset].astype('float') / 65535.0\r\n",
        "\r\n",
        "            elif len(dataset[nameset].shape) == 2: # Dataset with only 1 band of 1 scene (one .tif)\r\n",
        "                dataset[nameset] = self.stretch_hist(scenes)\r\n",
        "                dataset[nameset] = dataset[nameset].astype('float') / 65535.0\r\n",
        "        '''\r\n",
        "        \r\n",
        "        return dataset\r\n",
        "\r\n",
        "\r\n",
        "    def bandAdaptive_norm(self, dataset):\r\n",
        "        for nameset, scenes in dataset.items():\r\n",
        "            if isinstance(scenes, list): # Dataset with a list of scenes\r\n",
        "                for scene_i, scene in enumerate(scenes):\r\n",
        "                    scene = self.band_max_values(scene, 8)\r\n",
        "                    for band_i, band in enumerate(scene):\r\n",
        "                        scene[band_i] = self.stretch_hist(band) # this should be the same for different scenes\r\n",
        "                    dataset[nameset][scene_i] = scene\r\n",
        "\r\n",
        "            elif len(dataset[nameset].shape) == 3: # Dataset with only 1 scene (with some .tif bands)\r\n",
        "                scenes = self.band_max_values(scenes, 8)\r\n",
        "                for i, band in enumerate(scenes):\r\n",
        "                    scenes[i] = self.stretch_hist(band)\r\n",
        "                dataset[nameset] = scenes\r\n",
        "\r\n",
        "            elif len(dataset[nameset].shape) == 2: # Dataset with only 1 band of 1 scene (one .tif)\r\n",
        "                dataset[nameset] = self.stretch_hist(scenes)\r\n",
        "                \r\n",
        "        return dataset\r\n",
        "\r\n",
        "\r\n",
        "    def band_max_values(self, bandList, m = 8., verbose = False):\r\n",
        "        '''\r\n",
        "        Finds outliers value pixels from bands and limits the max possible value of that band to the median of its neighbours\r\n",
        "        '''\r\n",
        "        # Get max values from each band\r\n",
        "        maxs = []\r\n",
        "        for band in bandList:\r\n",
        "            if isinstance(band, gdal.Dataset):\r\n",
        "                band = band.ReadAsArray() # To numpy array\r\n",
        "\r\n",
        "            maxs.append(np.max(band))\r\n",
        "        \r\n",
        "        # Fix outlier pixels\r\n",
        "        if verbose:\r\n",
        "            print('Input:')\r\n",
        "            plt.figure(1)\r\n",
        "            plt.plot(maxs)\r\n",
        "            plt.show()\r\n",
        "        \r\n",
        "        maxs = np.array(maxs)\r\n",
        "        d = np.abs(maxs - np.median(maxs))\r\n",
        "        mdev = np.median(d)\r\n",
        "        s = d/mdev if mdev else 0.\r\n",
        "        outlierBands = np.where(s>m)[0]\r\n",
        "\r\n",
        "        for band in outlierBands:\r\n",
        "            neighbours = np.append(maxs[band-2:band], maxs[band+1:band+3])\r\n",
        "            newMaxValueBand = np.median(neighbours)\r\n",
        "            bandList[band][bandList[band]>newMaxValueBand] = newMaxValueBand\r\n",
        "            maxs[band] = newMaxValueBand\r\n",
        "\r\n",
        "        if verbose:\r\n",
        "            print('Output:')\r\n",
        "            plt.figure(2)\r\n",
        "            plt.plot(maxs)\r\n",
        "            plt.show()\r\n",
        "\r\n",
        "        return bandList\r\n",
        "\r\n",
        "    def stretch_hist(self, band, verbose = False):\r\n",
        "        '''\r\n",
        "        Removes the {percentage} least frequent pixel values from a band in order to decrease noise\r\n",
        "        Admits a single band\r\n",
        "\r\n",
        "        band = band.flatten()\r\n",
        "        bs = math.ceil(band.max() - band.min())\r\n",
        "        hist = plt.hist(band, bins=bs)[0]\r\n",
        "        \r\n",
        "        hist_s = np.array(sorted(hist))\r\n",
        "        hist_sc = hist_s[hist_s>0]\r\n",
        "        plt.figure()\r\n",
        "        plt.hist(hist_sc)\r\n",
        "        print(hist_sc)\r\n",
        "        '''\r\n",
        "\r\n",
        "        if isinstance(band, gdal.Dataset):\r\n",
        "            band.ReadAsArray() # To numpy array\r\n",
        "\r\n",
        "        band_no0 = band[band>0] # Mask will not be taken into account for the histogram stretching\r\n",
        "        min_percent = 5   # Low percentile\r\n",
        "        max_percent = 95  # High percentile\r\n",
        "\r\n",
        "        # Find lower and upper percentile using percentile function, and \"stretch\" pixel range linearly between lower and upper percentiles\r\n",
        "        lo, hi = np.percentile(band_no0, (min_percent, max_percent)) \r\n",
        "\r\n",
        "        # Apply linear \"stretch\" - lo goes to 0, and hi goes to 1\r\n",
        "        res_img = (band.astype(float) - lo) / (hi-lo)\r\n",
        "\r\n",
        "        #Multiply by 1, clamp range to [0, 1] and convert to float32\r\n",
        "        res_img = np.maximum(np.minimum(res_img*1, 1), 0).astype(np.float32)\r\n",
        "\r\n",
        "        if verbose:\r\n",
        "            print('Input:')\r\n",
        "            plt.figure(1)\r\n",
        "            plt.imshow(band, cmap='gray', vmin=0, vmax=np.max(band))\r\n",
        "            plt.show()\r\n",
        "            print('Output (Normalized):')\r\n",
        "            plt.figure(2)\r\n",
        "            plt.imshow(res_img, cmap='gray', vmin=0, vmax=1)\r\n",
        "            plt.show()\r\n",
        "\r\n",
        "        return res_img\r\n",
        "\r\n",
        "\r\n",
        "class NormalizeOptions(object):\r\n",
        "    \"\"\"\r\n",
        "    Normalizes data\r\n",
        "    Fixes pixel values outliers in bands comparing to others\r\n",
        "    Applies normalization between values [0-1]\r\n",
        "    Admits both a list of bands and a single band\r\n",
        "    \"\"\"\r\n",
        "    def __init__(self, norm_type):\r\n",
        "        self.norm_type = norm_type\r\n",
        "\r\n",
        "class ToPatchesOptions(object):\r\n",
        "    \"\"\"Convert dataset to image patches with custom window size and step size\"\"\"\r\n",
        "\r\n",
        "    def __init__(self, window, step, evadeClouds, miniPatches):\r\n",
        "        self.window = window\r\n",
        "        self.step = step\r\n",
        "        #self.setForShape = setForShape # The name of the set to extract the patches from corresponding to its shape. I use Hyperion shape because it is the smallest\r\n",
        "        #self.setsForPatching = setsForPatching # List of the name sets to divide into patches\r\n",
        "        self.evadeClouds = evadeClouds\r\n",
        "        self.miniPatches = miniPatches\r\n",
        "\r\n",
        "\r\n",
        "class ToTensor(object):\r\n",
        "    \"\"\"Convert ndarrays in sample to Tensors.\"\"\"\r\n",
        "    def __call__(self, sample):\r\n",
        "        # numpy image: H x W x C // Not using this format in any moment\r\n",
        "        # torch image: C x H x W // Embraced this for my data\r\n",
        "        \r\n",
        "        sample = torch.from_numpy(sample)\r\n",
        "        return sample"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}