"""
@brief  Collection of data loaders for different datasets.
@author Luis Carlos Garcia Peraza Herrera (luiscarlos.gph@gmail.com).
@date   2 Jun 2021.
"""

import os
import torch
import torchvision
import PIL
import numpy as np
import random
import colour
import scipy
import matplotlib.pyplot as plt

# My imports
import torchseg.base
import torchseg.utils
import torchseg.data_loader as dl
import monai.data.dataset


class MnistDataLoader(torchseg.base.BaseDataLoader):
    """
    @class MnistDataLoader loads MNIST data.
    """
    def __init__(self, data_dir, batch_size, shuffle=True, 
                 validation_split=0.0, num_workers=1, training=True):
        # Store parameters
        self.data_dir = data_dir

        # Create normalisation transform
        trsfm = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize((0.1307,), (0.3081,))
        ])

        # Create dataset handler
        self.dataset = torchvision.datasets.MNIST(self.data_dir, 
                                                  train=training, 
                                                  download=True, 
                                                  transform=trsfm)

        # Call the constructor of torchseg.base.BaseDataLoader
        super().__init__(self.dataset, batch_size, shuffle, validation_split, 
            num_workers)


class OdsiDbDataLoader(torchseg.base.BaseDataLoader):
    # Input image modes
    modes = [
        'rgbpixel', 
        'rgbpixel_test', 
        'rgbimage',
        'spixel_51',
        'simage_51', 
        'spixel_170', 
        'spixel_170_test', 
        'simage_170',
        'spixel_204',
        'simage_204',
        'boiko',
    ]
    
    # The array of wavelengths depends on the input mode
    mode2wl = {
        'spixel_51' : np.linspace(450, 950, 51),
        'simage_51' : np.linspace(450, 950, 51),
        'spixel_170': np.linspace(450, 950, 170),
        'spixel_170_test': np.linspace(450, 950, 170),
        'simage_170': np.linspace(450, 950, 170),
        'spixel_204': np.linspace(400, 1000, 204),
        'simage_204': np.linspace(400, 1000, 204),
        'boiko'     : np.array([450., 500., 600.]),
    }

    # RGB normalisation stats
    rgb_mean = np.array([0.403129887605, 0.288007343691, 0.299380141092], 
                        dtype=np.float32)
    rgb_std  = np.array([0.159463123785, 0.130188316525, 0.12346950031],
                        dtype=np.float32) 
    
    # Nuance EX normalisation stats
    nuance_ex_mean = np.array([
        0.0776138608947, 0.0783226294712, 0.0915784099028, 
        0.0851693853982, 0.0831637362668, 0.0847127943303, 
        0.085340519524,  0.0845109800243, 0.0790671446523, 
        0.0751274680049, 0.0761701337076, 0.0781188917647, 
        0.0782783973927, 0.0825329193882, 0.0988390369608, 
        0.120948266959,  0.140439376036,  0.154276410607, 
        0.164172320211,  0.171390068447,  0.178038853241, 
        0.184233131521,  0.189796964942,  0.195031949356, 
        0.198907463087,  0.198520651813,  0.200358926766, 
        0.203468102035,  0.201678233612,  0.198429438372, 
        0.193854994342,  0.191312533015,  0.189530537391, 
        0.18868570976,   0.18854169948,   0.189165124177, 
        0.189642954468,  0.189296877222,  0.180562774436, 
        0.177428725348,  0.174671366762,  0.172125912418, 
        0.16933901838,   0.167608797022,  0.165615862068, 
        0.163236092449,  0.16056347767,   0.156600176469, 
        0.143126733362,  0.138918268583,  0.135092694289,
    ], dtype=np.float32)
    nuance_ex_std = np.array([
        0.0038290728339,  0.00386745890456, 0.00461893938556, 
        0.00427318759442, 0.00421193344028, 0.00433005583596, 
        0.00441629367651, 0.00442865858953, 0.00416620537279, 
        0.00409641547929, 0.00419634961461, 0.0043631548856, 
        0.00440318450805, 0.00471044520012, 0.00585655629628, 
        0.00776347677282, 0.00967594428827, 0.0112555765706, 
        0.0125962817541,  0.0135586755183,  0.0144026812212, 
        0.015031395142,   0.0156616627971,  0.0162450380476, 
        0.016777610301,   0.0171646905704,  0.0174596376809, 
        0.0175789718117,  0.0174627609924,  0.0171899509512, 
        0.0168640925817,  0.0167477550425,  0.0166329796535, 
        0.016384616472,   0.0162204907189,  0.0160845513226, 
        0.015832708941,   0.0154426560894,  0.0146519935725, 
        0.0140359361608,  0.0134760286473,  0.0129535803584, 
        0.0124506272206,  0.0117099903208,  0.0112119812513, 
        0.0107739222483,  0.0103392189441,  0.0100008831543, 
        0.00827122500705, 0.00774512242113, 0.00726368979869,
    ], dtype=np.float32)

    # Specim IQ normalisation stats
    specim_iq_mean = np.array([
        0.112085002548,  0.102140085151,  0.0943113060797, 0.0885467882906, 
        0.0846466315166, 0.0820774967744, 0.0802850078561, 0.0791292167016, 
        0.0784498196042, 0.0780782467903, 0.0780644673204, 0.0784523791739, 
        0.0792390083204, 0.0804284923453, 0.0818753314213, 0.0834831640881, 
        0.0851511772687, 0.0867686173785, 0.0883751009176, 0.0898204818059, 
        0.0912809000747, 0.0926440870137, 0.0939971489817, 0.0952519009781, 
        0.0965328084729, 0.097690027486,  0.0987274062461, 0.0996016817623, 
        0.100458407892,  0.101202582001,  0.10195598233,   0.102806150098, 
        0.103694965839,  0.104664556992,  0.105653205393,  0.10659280337, 
        0.10746400651,   0.108253455173,  0.108935761923,  0.109321392809, 
        0.109375497886,  0.108865325572,  0.107869014925,  0.106254806345, 
        0.104402478702,  0.102317535445,  0.100404782722,  0.0988052551902, 
        0.0977005111836, 0.0969098524837, 0.0966750172086, 0.0968155206467, 
        0.0974678382789, 0.0984965585346, 0.0997345260659, 0.100758822558, 
        0.101404000532,  0.101469627798,  0.101135541391,  0.100144589499, 
        0.0986686515618, 0.0974109084362, 0.0973515855222, 0.0992310783409, 
        0.103559105299,  0.110339683424,  0.119200190181,  0.129200049368, 
        0.139490848038,  0.149170142667,  0.158144175815,  0.166132853615, 
        0.173451190929,  0.179981349582,  0.185961931089,  0.191155738785, 
        0.195689772047,  0.199596934812,  0.203163066338,  0.206382320236, 
        0.209400294235,  0.212104318164,  0.214609173942,  0.216990035595, 
        0.219227218176,  0.221412241784,  0.223472401843,  0.225275373864, 
        0.227072332864,  0.228744232706,  0.230303433887,  0.231838015095, 
        0.233366165061,  0.234887439646,  0.236506400265,  0.237957146659, 
        0.239588794172,  0.241068114,     0.242681232832,  0.244230795064, 
        0.245745089833,  0.24692634115,   0.248064798697,  0.248622792136, 
        0.249324235054,  0.249977211781,  0.250498342546,  0.250987490527, 
        0.251156369557,  0.251248702127,  0.25096838481,   0.250822536009, 
        0.250102103495,  0.249356839651,  0.248629243344,  0.247524143344, 
        0.246625750813,  0.245947435307,  0.245181069521,  0.243986731324, 
        0.242938204373,  0.242353213664,  0.241708338683,  0.241195778181, 
        0.241245919936,  0.241264032642,  0.241610304121,  0.242049243169, 
        0.2421130799,    0.242264614189,  0.242367022354,  0.2409313546, 
        0.241159969908,  0.242161924135,  0.242228185646,  0.242473373156, 
        0.242266678267,  0.24223316472,   0.241815605688,  0.241536003903, 
        0.240810299738,  0.239586301619,  0.238636256726,  0.238204898501, 
        0.237551727056,  0.236373644054,  0.235011936757,  0.233764670484, 
        0.232366132656,  0.231146941133,  0.23080009416,   0.230533240929, 
        0.230076653688,  0.229484626509,  0.228599976517,  0.227914845773, 
        0.227556968449,  0.22725171334,   0.226303809768,  0.225347473023, 
        0.224614910223,  0.224014662255,  0.22339760286,   0.222414620401, 
        0.221382271403,  0.220377157411,  0.219338854373,  0.218130296434, 
        0.217091878005,  0.215903149325,  0.214644602196,  0.213494551926, 
        0.212338519749,  0.211226395102,  0.210369281188,  0.209304051737, 
        0.207966701355,  0.206544497363,  0.204824802113,  0.203001666349, 
        0.201089720809,  0.199036227562,  0.196907502338,  0.194647061506, 
        0.191861582388,  0.188844590578,  0.18527816507,   0.181248328695, 
        0.1777737335,    0.175078837825,  0.1733158301,    0.172441615107, 
        0.171608116792,  0.171192572651,  0.171049377519,  0.170862948778, 
        0.171274492165,  0.171573746293,  0.172177039069,  0.172957085572, 
        0.173301498106,  0.174158032792,  0.174963397296,  0.161150677433
    ], dtype=np.float32)
    specim_iq_std = np.array([
        0.0051204327844,  0.0049010894168,  0.00490408951459, 0.00501242131618, 
        0.00515869582467, 0.00531481044734, 0.00547126338076, 0.00563346798113, 
        0.00579447352576, 0.00596555732144, 0.00615566882416, 0.00635973070577, 
        0.00657340461617, 0.0067967285048,  0.0070286671205,  0.00727708514704, 
        0.00752424001689, 0.00775700912999, 0.00798220085672, 0.00818689563432, 
        0.00838465923951, 0.00856988323327, 0.0087392056371,  0.00889675897692, 
        0.00905206000264, 0.0091959914293,  0.009312072202,   0.00940073464096, 
        0.00947625651694, 0.00953227890046, 0.00957837186835, 0.00964017941934, 
        0.00969605544023, 0.00977571641123, 0.00985408362882, 0.00993917538033, 
        0.0100143769771,  0.0100859121458,  0.0101409295895,  0.0101669455923, 
        0.0101443092583,  0.0100463008317,  0.00989182845609, 0.00966438387347, 
        0.00942136906458, 0.00917258006915, 0.008953890893,   0.00878593637129, 
        0.00867668098681, 0.00860270653642, 0.00859433373635, 0.0086204846356, 
        0.00869749103483, 0.00882064339864, 0.00896339789551, 0.00909145584191, 
        0.00918067503764, 0.00920361403469, 0.00919296874221, 0.00911463416519, 
        0.00897490062522, 0.00886443804711, 0.00888037267064, 0.00910847772243, 
        0.00962114957633, 0.0104625277379,  0.0116591370095,  0.0131564475032, 
        0.014852059166,   0.0165929757477,  0.0183182801335,  0.019935233596, 
        0.0214820271417,  0.0229137210796,  0.0242616624892,  0.0254825335416, 
        0.0265475514249,  0.0274784619776,  0.0283352254374,  0.0291056673323, 
        0.0298353614964,  0.0304790896885,  0.0310573625231,  0.0316169604888, 
        0.0321468398482,  0.0326668677501,  0.0331730392939,  0.0335956874147, 
        0.0340282795814,  0.0344635042052,  0.0348561658458,  0.0352599262886, 
        0.0356911594505,  0.0360880062706,  0.0365148940448,  0.0369118063317, 
        0.0373465100767,  0.0377660923652,  0.0382188255657,  0.0386452354932, 
        0.0390436577487,  0.0393508233,     0.0396008200152,  0.0396906081602, 
        0.0398572928164,  0.0399856573532,  0.0400535027338,  0.0401183618839, 
        0.0400770409866,  0.0400355447514,  0.0398820349639,  0.0397914604514, 
        0.0395229589731,  0.0392744169267,  0.0390286206185,  0.0386696315616, 
        0.0383528940588,  0.0381384245943,  0.0378284388325,  0.0374353366368, 
        0.0370882825179,  0.0368806895126,  0.0366293825145,  0.0364590994265, 
        0.036409711205,   0.0363757080885,  0.0364176906163,  0.0364910524371, 
        0.0364356384961,  0.0364550964887,  0.0363572534936,  0.035713651417, 
        0.0357296075661,  0.0360257484172,  0.036304159429,   0.0367355386762, 
        0.0366726984998,  0.0366845383865,  0.0365487921643,  0.0364630193957, 
        0.0362154434343,  0.0358276964578,  0.0355990390691,  0.0355736606419, 
        0.0354255001813,  0.0351446246617,  0.0347914317484,  0.0344802785938, 
        0.0341160734367,  0.0337945665366,  0.0336832205541,  0.0336308778566, 
        0.0334985680856,  0.0333494362052,  0.0331237910006,  0.0329646205593, 
        0.0328705891889,  0.0328222972035,  0.0326090837709,  0.0324180583771, 
        0.0322387706372,  0.0320901724742,  0.0319185152311,  0.0316350319596, 
        0.0313471149088,  0.031120219116,   0.0308533330801,  0.0305213851429, 
        0.0302389431622,  0.029960570957,   0.0296949319738,  0.0294617564081, 
        0.0291807764587,  0.0288812632713,  0.0286547736342,  0.0283600323959, 
        0.0279933843565,  0.0276050068161,  0.027095059129,   0.0265664812623, 
        0.026083041933,   0.0256118408366,  0.025073052516,   0.0244892219447, 
        0.0238096923167,  0.0230908555588,  0.0222759083261,  0.0213622231765, 
        0.0205653866989,  0.0199707523122,  0.0194896643608,  0.0192172825747, 
        0.0189468026241,  0.0187808803217,  0.0186485068087,  0.0184880364592, 
        0.0184822248112,  0.0184771573814,  0.0184898120673,  0.0185603547897, 
        0.018496862837,   0.0185234983434,  0.0185140609897,  0.0130384043098, 
    ], dtype=np.float32)

    # Normalisation stats for interpolated bands 450-950nm (51 bands)
    hyper_51_mean = np.array([
        0.0810067996349, 0.0829831647364, 0.093483960472,  0.0899734763065, 
        0.0894423085798, 0.0915193721587, 0.092661139279,  0.0915338960809, 
        0.0857899568086, 0.0817884475053, 0.0828441900691, 0.0852258059802, 
        0.0848364019569, 0.0874677529617, 0.105330525865,  0.130889929167, 
        0.152258716436,  0.167234953479,  0.177657939938,  0.185337380804, 
        0.192140167946,  0.198199267684,  0.203636422124,  0.208865833329, 
        0.213144227103,  0.213884648171,  0.215777113836,  0.218022347793, 
        0.216236726977,  0.213079549878,  0.208936060479,  0.206582950682, 
        0.205533508108,  0.205109621906,  0.204955697994,  0.205414174429, 
        0.205367417871,  0.204263258639,  0.197107520177,  0.193805749636, 
        0.191465798416,  0.189074327189,  0.186375384941,  0.184456663586, 
        0.182047172551,  0.179237997628,  0.176214905006,  0.172386557792, 
        0.161391944603,  0.156415102283,  0.150869385165
    ], dtype=np.float32)
    hyper_51_std = np.array([
        0.00514014020925, 0.00539040833008, 0.00603190822433, 0.0059271066889, 
        0.00597948945687, 0.00616100200841, 0.00629645035981, 0.00619439587162, 
        0.00575553630156, 0.00557603115215, 0.00569343612698, 0.00594936225045, 
        0.00592971065021, 0.00609018586403, 0.00776464603105, 0.0109345712314, 
        0.0139807234261,  0.0163833898851,  0.0181979949318,  0.0195246053556, 
        0.0206441248362,  0.0215073191503,  0.0223579382643,  0.0231928340066, 
        0.0240259381558,  0.0246012694501,  0.0249232918579,  0.0248939753191, 
        0.0246073320615,  0.0241213238954,  0.0235886267345,  0.0233070117303, 
        0.023279264433,   0.0231155168478,  0.022900319036,   0.0229853245983, 
        0.022653123846,   0.0221113475025,  0.0214060499578,  0.0206702027159, 
        0.0202024043012,  0.0196976652234,  0.0191932768321,  0.0184724650055, 
        0.017832398352,   0.0172010450473,  0.0166074947096,  0.0160897372069, 
        0.0146111361932,  0.0136754420132,  0.0125149234869,
    ], dtype=np.float32)

    # Normalisation stats for interpolated bands 450-950nm (170 bands)
    hyper_170_mean = np.array([
        0.0810067996349, 0.0816043995643, 0.0821970031859, 0.0827673146001, 
        0.0849229440499, 0.0880378175387, 0.0911460114559, 0.0932456460856, 
        0.0922341744158, 0.09118836566,   0.0901252695013, 0.0898196994614, 
        0.0896553436568, 0.089513093682,  0.0897392633201, 0.0903613267805, 
        0.090979011488,  0.0915592812428, 0.0919451792569, 0.0923012143522, 
        0.0925915664477, 0.0925672795887, 0.0923141765019, 0.0918983693193, 
        0.0909842773978, 0.0893221808685, 0.0875870100037, 0.0858583116221, 
        0.084507488729,  0.0833029689943, 0.0822162716762, 0.0818705060695, 
        0.0820900963864, 0.0824551648085, 0.0829875784395, 0.0837494090013, 
        0.0844853113882, 0.0851180245572, 0.0852873398232, 0.0852510529314, 
        0.085032962395,  0.0850213767293, 0.0854903869798, 0.0862520932472, 
        0.0877033510983, 0.0921965406489, 0.0974439886213, 0.103359354549, 
        0.110504447195,  0.1181816836,    0.125735154289,  0.132916819701, 
        0.139431940908,  0.14572431233,   0.151778762252,  0.156579367882, 
        0.161062623956,  0.16533550223,   0.16897751279,   0.172123531667, 
        0.175163861539,  0.178042253332,  0.180372587267,  0.182632682035, 
        0.184851135499,  0.186934765385,  0.188969781225,  0.19097210744, 
        0.192868746172,  0.194688788255,  0.196478620273,  0.198231629835, 
        0.199843773024,  0.201451354722,  0.203056473628,  0.204639118831, 
        0.206169527005,  0.207723786193,  0.209189575556,  0.210462771206, 
        0.211735450466,  0.212993707719,  0.21341177993,   0.213682033696, 
        0.213833884439,  0.214179921524,  0.214760809264,  0.215310997878, 
        0.21587149132,   0.216600461745,  0.217276433301,  0.217878384209, 
        0.217710349168,  0.217203954087,  0.216615029805,  0.215917558555, 
        0.214973406228,  0.214003477609,  0.213096996496,  0.21194623461, 
        0.210718301228,  0.209438320294,  0.208491630456,  0.207786470717, 
        0.207085517874,  0.206505790068,  0.206151632208,  0.205824085203, 
        0.205570909186,  0.205465349108,  0.205318419385,  0.205186642554, 
        0.204999346857,  0.204648895535,  0.204734935304,  0.204967520688, 
        0.205123997043,  0.205306585329,  0.205378787542,  0.205464869726, 
        0.205440862705,  0.205448678606,  0.205270700196,  0.204834264548, 
        0.204479281473,  0.204276377615,  0.202394212481,  0.20025954336, 
        0.1980644597,    0.196538453003,  0.195484550546,  0.194468031395, 
        0.193685258656,  0.1930344551,    0.192345946229,  0.191614290102, 
        0.190845453654,  0.190099431202,  0.189436042822,  0.188795821318, 
        0.188020086576,  0.187164688065,  0.186352484731,  0.18579908009, 
        0.185262986887,  0.184676669537,  0.18398799608,   0.183275393733, 
        0.182564743891,  0.181809907514,  0.180977080903,  0.180169855655, 
        0.179324102047,  0.178414858592,  0.177525557928,  0.176636254971, 
        0.175626738953,  0.174556558391,  0.17342668134,   0.171943895798, 
        0.168758784331,  0.165489093699,  0.162187915121,  0.160305287515, 
        0.158842194262,  0.157354797792,  0.15586327933,   0.154299476375, 
        0.152647749572,  0.150869385165,
    ], dtype=np.float32)  
    hyper_170_std = np.array([
        0.00514014020925, 0.00518764780204, 0.00525823383173, 
        0.00535012794449, 0.00547291214664, 0.00563897203488, 
        0.00584719693034, 0.00601840401274, 0.00597312393174, 
        0.00594230808148, 0.00592796769322, 0.00592149651265, 
        0.00593073969924, 0.00595863235588, 0.00599366047689, 
        0.00603598420629, 0.00609498985838, 0.00616283944647, 
        0.00619021429061, 0.00623129889715, 0.00628079540453, 
        0.00628164484446, 0.0062554232367,  0.00622091327839, 
        0.00614385296943, 0.00600149880839, 0.00586976669371, 
        0.00575943407638, 0.00566313714922, 0.00560383617056, 
        0.00557784065547, 0.00557070429369, 0.00558360532025, 
        0.00563066298374, 0.00570444278682, 0.00577308272315, 
        0.00585341427136, 0.00593415400159, 0.0059523230271, 
        0.00595304568405, 0.00594175737341, 0.0059186263843, 
        0.00591387748421, 0.00596412250667, 0.00610436154186, 
        0.00643399196802, 0.00690699378334, 0.00753532691265, 
        0.0083193032524,  0.0092326761507,  0.0102199661071, 
        0.0112077588638,  0.0121033252134,  0.0130072069277, 
        0.0139072203158,  0.0146611876077,  0.0153828289821, 
        0.0160744121867,  0.0166766603911,  0.0172164795582, 
        0.0177488468134,  0.0182632136293,  0.0186566346647, 
        0.0190425397251,  0.0194357015034,  0.0197617010843, 
        0.0200872805015,  0.0204344186764,  0.0207311245281, 
        0.0209694099278,  0.0212335403014,  0.0215115949242, 
        0.0217373517193,  0.0219897650935,  0.0222568826317, 
        0.0224981370852,  0.0227274337641,  0.0229864699088, 
        0.0232474138528,  0.0234793138343,  0.0237303123809, 
        0.0239932552179,  0.0241723466339,  0.0243493688858, 
        0.0245137766694,  0.024635517016,   0.0247200187774, 
        0.0248183733464,  0.0249209704549,  0.0248972215491, 
        0.0248901565673,  0.0248923538104,  0.0248399712793, 
        0.024752651806,   0.0246569935811,  0.0245488673966, 
        0.0243823754977,  0.0242277833625,  0.024122808071, 
        0.023967193481,   0.0237996737411,  0.0236418743447, 
        0.0235121271635,  0.0234115922515,  0.0233334145963, 
        0.0232919621708,  0.0232425509873,  0.0232313866122, 
        0.0232695414247,  0.0232177091961,  0.0231520231536, 
        0.0231285687793,  0.0230275623595,  0.0228080246666, 
        0.0228104850908,  0.0229009641254,  0.0229397383441, 
        0.0230192843686,  0.0229851797824,  0.0229261270692, 
        0.0228131414284,  0.0227337707128,  0.0225904626606, 
        0.0223573308616,  0.0221935579431,  0.0221160037899, 
        0.0219219949626,  0.0217030911434,  0.0214877019102, 
        0.021263340512,   0.021007989729,   0.0207843640743, 
        0.0206411914581,  0.020503339569,   0.0203641347982, 
        0.0202278523375,  0.0200458075605,  0.0198785898809, 
        0.0197536648597,  0.0196417853907,  0.0194847954329, 
        0.0193236536251,  0.0191825470482,  0.0189384395597, 
        0.0187318605473,  0.0185380521754,  0.0183197827216, 
        0.0181161554718,  0.0179453503841,  0.0177620880403, 
        0.0175395335466,  0.0173615623233,  0.0172140745547, 
        0.0170036939038,  0.0168258604467,  0.0166715328373, 
        0.0165087393902,  0.0163600570587,  0.0162155713658, 
        0.0160108820457,  0.0154940480883,  0.0150403349807, 
        0.0146808652842,  0.0143833065811,  0.0141067654461, 
        0.0138384200987,  0.0135452274736,  0.0131943819911, 
        0.0128539926089,  0.0125149234869,
    ], dtype=np.float32)

    # Normalisation stats for interpolated bands 400-1000nm (204 bands)
    hyper_204_mean = np.array([
        0.0853443115378, 0.082848078348,  0.0810082485533, 0.0797739291089, 
        0.0789765010856, 0.0784253976231, 0.0780743386258, 0.0778704430139, 
        0.0777675948717, 0.0777807791937, 0.0779223368812, 0.0781914794129, 
        0.0785791016654, 0.0790429332303, 0.0795490719257, 0.0800657360188, 
        0.0805705773184, 0.0810564855142, 0.0816537555639, 0.0822434301913, 
        0.0828129241917, 0.0851697518554, 0.0882823012806, 0.0913843392534, 
        0.0931698296219, 0.0921555102572, 0.0911103278881, 0.0900456833142, 
        0.089806860604,  0.0896448332975, 0.0895035055374, 0.0897823078737, 
        0.0904041013803, 0.0910202057368, 0.0915858176773, 0.0919696644263, 
        0.0923231893207, 0.0926073570099, 0.0925576733003, 0.0922939783025, 
        0.0918693242863, 0.090887923841,  0.0892229863784, 0.0874854547235, 
        0.0857666488026, 0.0844349513903, 0.0832400678965, 0.0821597277192, 
        0.0818779649125, 0.0821034052056, 0.0824762599571, 0.0830232482897, 
        0.0837875118399, 0.0845195548363, 0.0851460471168, 0.0852897622035, 
        0.0852479044979, 0.0850210136633, 0.0850393938706, 0.0855107406393, 
        0.086287057134,  0.0878600273907, 0.0923785274538, 0.0976499201371, 
        0.103583522281,  0.110778056089,  0.118450928734,  0.12599042822, 
        0.133139807533,  0.139638323898,  0.145918215417,  0.151959042544, 
        0.156715107509,  0.161186999087,  0.165450053057,  0.169061771918, 
        0.172201978179,  0.175236667448,  0.178097874942,  0.180423778624, 
        0.18268025368,   0.184895774742,  0.186974093397,  0.18900679109, 
        0.191006481496,  0.192898356772,  0.194716583326,  0.196504052264, 
        0.198253142208,  0.199863596941,  0.201469554295,  0.203073066396, 
        0.204653778571,  0.206182218116,  0.207735383038,  0.209197545352, 
        0.210469742978,  0.211741056628,  0.212998041639,  0.213412478345, 
        0.213682428528,  0.213833922756,  0.214179626296,  0.214759964092, 
        0.215309688396,  0.21586877854,   0.216597401741,  0.217272827733, 
        0.217874829296,  0.217713402978,  0.217208848111,  0.216620579664, 
        0.215926719795,  0.214984752083,  0.214015036535,  0.213108572771, 
        0.211962954426,  0.210738131305,  0.209458724713,  0.208503667422, 
        0.20779950603,   0.20709852571,   0.206512866728,  0.206159266439, 
        0.205829899335,  0.2055763301,    0.205469072134,  0.2053216321, 
        0.205190373083,  0.205011808434,  0.204647795705,  0.204727073534, 
        0.204963079918,  0.205117748587,  0.205304463145,  0.205374891607, 
        0.205465791763,  0.205440357818,  0.205452973033,  0.205285706192, 
        0.20485040917,   0.204493097663,  0.204284354293,  0.202475507375, 
        0.200349367799,  0.198158744982,  0.196583004897,  0.195532072024, 
        0.194514200173,  0.193716582944,  0.193065301919,  0.192380201088, 
        0.191651233395,  0.190885340272,  0.190137014281,  0.189469187682, 
        0.188831440083,  0.18806680562,   0.187212345013,  0.186388536669, 
        0.185829852653,  0.185294547243,  0.18471523209,   0.184031349326, 
        0.18331896063,   0.182609573035,  0.181863755887,  0.181028442121, 
        0.180224994041,  0.179381357649,  0.178474862566,  0.177586557723, 
        0.176697226046,  0.175701928521,  0.174635658965,  0.17350939096, 
        0.172175970531,  0.168995748617,  0.165735941283,  0.162440184803, 
        0.160415869749,  0.158957310826,  0.157473121414,  0.155979883857, 
        0.154430096451,  0.152785336609,  0.151022345404,  0.149900274821, 
        0.148793261038,  0.147873480083,  0.147190194871,  0.146765078222, 
        0.146512095475,  0.146301063977,  0.146197658023,  0.146153203837, 
        0.14612417137,   0.146241604267,  0.146334423564,  0.146512443726, 
        0.146740819307,  0.146845406329,  0.147084946832,  0.147322538854,
    ])
    hyper_204_std = np.array([
        0.00429063019463, 0.00421556938322, 0.00421326782134, 0.00424388475654, 
        0.00428636463332, 0.00433271456692, 0.00438247193098, 0.00443299962565, 
        0.00448713852853, 0.00454708491045, 0.00461157550422, 0.00467946718575, 
        0.00475101854674, 0.00482716376698, 0.00490837051352, 0.00498946941246, 
        0.00506774484486, 0.00514309477465, 0.0051925135872,  0.00526480835903, 
        0.00535829764719, 0.00548445997478, 0.00565378259084, 0.00586489236965, 
        0.00601440273555, 0.005970148491,   0.00594068114053, 0.00592745749255, 
        0.00592144445538, 0.00593224420545, 0.00596114517579, 0.00599609616221, 
        0.00603945896088, 0.00609961389108, 0.00616417685938, 0.00619258440357, 
        0.00623438905956, 0.00628420787443, 0.00628055551672, 0.00625363016914, 
        0.00621884501468, 0.0061352346863,  0.00599351875816, 0.00586262239631, 
        0.00575361016223, 0.00565861485312, 0.00560151576211, 0.00557717177346, 
        0.00557049875569, 0.00558494534593, 0.00563381330001, 0.00570727512085, 
        0.0057769231621,  0.00585751955194, 0.00593806224132, 0.00595258611539, 
        0.00595319388812, 0.00594124665851, 0.00591775837127, 0.0059141902236, 
        0.00596699872254, 0.00611386079954, 0.00644855789032, 0.00692696261297, 
        0.00756087973826, 0.00835009031748, 0.00926676960146, 0.0102550364563, 
        0.0112383325076,  0.0121326939975,  0.0130359966992,  0.0139347779985, 
        0.0146831125972,  0.0154032868349,  0.0160932585878,  0.0166910923292, 
        0.0172302476572,  0.0177618612655,  0.0182727020506,  0.0186654340517, 
        0.0190508131667,  0.0194438153118,  0.0197677705176,  0.0200934680362, 
        0.0204406290093,  0.0207347602012,  0.0209732759897,  0.0212375094924, 
        0.0215144401639,  0.0217402755271,  0.0219927754858,  0.0222597296813, 
        0.0225002660856,  0.0227294251226,  0.022988508771,   0.0232487835734, 
        0.0234806482899,  0.0237314659303,  0.0239941872785,  0.0241727708999, 
        0.0243496449606,  0.0245138515016,  0.0246354792911,  0.0247198722799, 
        0.024818114158,   0.02492103543,    0.0248973159425,  0.0248901019265, 
        0.0248923472587,  0.0248404151896,  0.0247535516486,  0.0246577808842, 
        0.0245505008988,  0.0243843574091,  0.0242294347995,  0.0241238050661, 
        0.0239695315257,  0.0238023597521,  0.023644131753,   0.023513948889, 
        0.0234134082593,  0.0233344132812,  0.0232932802508,  0.0232432945068, 
        0.0232309333753,  0.0232681973267,  0.0232199005178,  0.0231528377813, 
        0.0231294812682,  0.0230352229922,  0.0228088477442,  0.0228078705025, 
        0.0229006847232,  0.0229372223784,  0.023020913354,   0.0229853138712, 
        0.0229304161225,  0.0228162353109,  0.0227383430924,  0.0225990666575, 
        0.0223655419096,  0.022199526403,   0.0221186160642,  0.0219295642769, 
        0.0217117120569,  0.021496332392,   0.0212742212559,  0.0210190235232, 
        0.0207939541841,  0.0206486668801,  0.0205093053681,  0.0203708455962, 
        0.020234317687,   0.0200555277685,  0.0198863727191,  0.0197592893651, 
        0.0196486847768,  0.0194943242135,  0.0193319664984,  0.0191953265613, 
        0.0189512298341,  0.0187432517636,  0.0185500554975,  0.0183334317457, 
        0.0181273213302,  0.0179558892652,  0.0177777718174,  0.0175522171077, 
        0.0173725481567,  0.0172229668295,  0.0170170874327,  0.0168379915155, 
        0.016681685115,   0.0165205523857,  0.0163709393398,  0.0162251579454, 
        0.0160519388455,  0.0155285059359,  0.0150706924364,  0.0147043076919, 
        0.0144054270517,  0.0141271505124,  0.0138587851936,  0.0135721412487, 
        0.0132223164765,  0.0128809380303,  0.0125428954494,  0.0122095514875, 
        0.0118863380612,  0.011623718658,   0.0114240502952,  0.0112818544874, 
        0.0111888996149,  0.0111072365138,  0.0110557638921,  0.0110133352409, 
        0.0109718697148,  0.0109764311879,  0.0109805840783,  0.010994444414, 
        0.0110266496781,  0.0110136046495,  0.0110310144291,  0.0110412626426,
    ])

    # The means and stds also depend on the input mode
    mode2mean = {
        'spixel_51'  : hyper_51_mean,
        'simage_51'  : hyper_51_mean,
        'spixel_170' : hyper_170_mean, 
        'spixel_170_test' : hyper_170_mean, 
        'simage_170' : hyper_170_mean,
        'spixel_204' : hyper_204_mean,
        'simage_204' : hyper_204_mean,
        'boiko'      : None,
    }
    mode2std = {
        'spixel_51'  : hyper_51_std,
        'simage_51'  : hyper_51_std,
        'spixel_170' : hyper_170_std, 
        'spixel_170_test' : hyper_170_std, 
        'simage_170' : hyper_170_std,
        'spixel_204' : hyper_204_std,
        'simage_204' : hyper_204_std,
        'boiko'      : None,
    }
    
    class LoadImage(torch.nn.Module):
        """
        @class LoadImage is a transform that loads a TIF image and label. 
                         This is not supported by MONAI.
        """
        #def __init__(self, keys, mode='rgb', training=True):
        def __init__(self, keys, mode='rgbimage'):
            """
            @param[in]  mode  'rgbpixel': Load an RGB pixel at a time.  
                              'spixel':   Load a pixel-wide (1x1x51) 
                                          multispectral vector.
                              'rgbimage': Load the RGB version of the input 
                                          image.
                              'simage':   Load the multispectral version of 
                                          the input image.
            """
            assert(mode in OdsiDbDataLoader.modes)
            self.keys = keys
            self.mode = mode
            #self._training = training
            super().__init__()
        
        @staticmethod
        def get_cmfs(cmf_name: str):
            """
            @brief Get colour matching funtions for different standards.
            @param[in]  cmf_name  Choose one from [cie1931_2, cie1964_10, 
                                  cie2012_2, cie2012_10].
                                  From Wikipia: the CIE's color matching 
                                  functions are the numerical description of 
                                  the chromatic response of the observer. 
                                  They can be thought of as the spectral 
                                  sensitivity curves of three linear light 
                                  detectors yielding the CIE tristimulus 
                                  values X, Y and Z. 
                                  Collectively, these three functions describe 
                                  the CIE standard observer.
            """
            cmf_full_name = {
                'cie_2_1931':  'CIE 1931 2 Degree Standard Observer',
                'cie_10_1964': 'CIE 1964 10 Degree Standard Observer',
                'cie_2_2012':  'CIE 2012 2 Degree Standard Observer',
                'cie_10_2012': 'CIE 2012 10 Degree Standard Observer',
            }
            if cmf_name not in cmf_full_name: 
                err_msg = '[ERROR] Wrong CMF name. The available options are: '
                err_msg += ', '.join(cmf_full_name.keys())
                raise AttributeError(err_msg)
            standard_obs = colour.colorimetry.MSDS_CMFS_STANDARD_OBSERVER
            return standard_obs[cmf_full_name[cmf_name]]
        
        @staticmethod
        def gamma_correction(srgb_chan):
            """
            @brief Perform gamma correction on an sRGB image (i.e. the range of 
                   the data in the channel is [0., 1.]).
            @param[in]  u  One of the three sRGB channels.
            """
            u = srgb_chan[srgb_chan <= 0.0031308]      
            u_r = srgb_chan[srgb_chan > 0.0031308]      

            corrected = np.empty_like(srgb_chan)
            corrected[srgb_chan <= 0.0031308] = 12.92 * u
            corrected[srgb_chan > 0.0031308] = 1.055 * np.power(u_r, 0.416) \
                                               - 0.055

            return corrected
        
        @staticmethod
        def get_single_wl_im(raw_im_hyper, raw_wl, wl):
            """
            @brief Given an image and the corresponding wavelengths this function 
                   returns the image corresponding to a specific channel.
            @param[in]  raw_im_hyper  Hyperspectral time, 
                                      shape (h, w, nbands).
            @param[in]  raw_wl        Wavelengths, shape (bands,).
            @param[in]  wl            Specific wavelegth of the image you want to 
                                      extract. 
            @returns an array of dimenstion (h, w).
            """
            idx = (np.abs(raw_wl - wl)).argmin()
            return raw_im_hyper[:, :, idx].copy()
        
        @staticmethod
        def filter_bands(raw_im_hyper, raw_wl, min_wl, max_wl):
            """
            @brief Given a hyperspectral image this function removes those
                   bands outside the [min_wl, max_wl] range.
            @param[in]  raw_im_hyper  Hyperspectral time, 
                                      shape (h, w, nbands).
            @param[in]  raw_wl        Wavelengths, shape (bands,).
            @param[in]  min_wl        Minimum wavelength.
            @param[in]  max_wl        Maximum wavelength.
            @returns a tuple (new_im_hyper, new_wl) where new_im_hyper is the new 
                     hyperspectral image and new_wl is the new vector of wavelengths 
                     of the channels contained in the new_im_hyper.
            """
            # Create new array of wavelengths
            new_wl = list(filter(lambda x: (x >= min_wl and x <= max_wl), raw_wl))

            # Create new hyperspectral image 
            new_im_hyper = np.empty((raw_im_hyper.shape[0], raw_im_hyper.shape[1], 
                                     len(new_wl)),
                                    dtype=raw_im_hyper.dtype)
            for chan_idx, wl in enumerate(new_wl):
                new_im_hyper[:, :, chan_idx] = \
                    OdsiDbDataLoader.LoadImage.get_single_wl_im(raw_im_hyper, raw_wl, wl)

            return new_im_hyper, np.array(new_wl)

        @staticmethod
        def get_additive_correction(cmf_name, wl):
            """
            @brief TODO: although not necessary for ODSI-DB because the data goes until 900nm
                         at least, we should compensate for the data missing the right side of
                         the spectrum too
            """

            # Get sampled colour matching functions
            cmfs = OdsiDbDataLoader.LoadImage.get_cmfs(cmf_name)
            xbar_y, ybar_y, zbar_y = colour.utilities.tsplit(cmfs.values)
            
            # Get x-coordinates of the missing wavelengths
            x_corr = [x for x in cmfs.wavelengths if x < min(wl)]
            epsilon = min(wl) - x_corr[-1]
            offset = max(x_corr) - min(x_corr) + epsilon
            x_corr = [x + offset for x in x_corr]
            
            # Get additive correction for xbar
            xbar_y_corr = [y for x, y in zip(cmfs.wavelengths, xbar_y) if x < min(wl)]
            xbar_y_corr.reverse()

            # Get additive correction for ybar
            ybar_y_corr = [y for x, y in zip(cmfs.wavelengths, ybar_y) if x < min(wl)]
            ybar_y_corr.reverse()
            
            # Get additive correction for zbar
            zbar_y_corr = [y for x, y in zip(cmfs.wavelengths, zbar_y) if x < min(wl)]
            zbar_y_corr.reverse()
            
            # Get interpolated functions for the additive correction
            f_xbar_corr = scipy.interpolate.PchipInterpolator(x_corr, xbar_y_corr, 
                                                              extrapolate=False)
            f_ybar_corr = scipy.interpolate.PchipInterpolator(x_corr, ybar_y_corr,
                                                              extrapolate=False)
            f_zbar_corr = scipy.interpolate.PchipInterpolator(x_corr, zbar_y_corr,
                                                              extrapolate=False)
            return f_xbar_corr, f_ybar_corr, f_zbar_corr

        @staticmethod
        def get_corrected_cmf(cmf_name, wl):
            """
            @brief Get colour matching function modified to colour-compensate 
                   for missing wavelengths in the input images.
            @param[in]  cmf_name  TODO
            @param[in]  wl        TODO
            """
            # Get colour matching function (wavelength -> xyz_bar)
            cmfs = OdsiDbDataLoader.LoadImage.get_cmfs(cmf_name)
            xbar_y, ybar_y, zbar_y = colour.utilities.tsplit(cmfs.values)
            f_xbar = scipy.interpolate.PchipInterpolator(cmfs.wavelengths, 
                                                         xbar_y, 
                                                         extrapolate=False)
            f_ybar = scipy.interpolate.PchipInterpolator(cmfs.wavelengths, 
                                                         ybar_y, 
                                                         extrapolate=False)
            f_zbar = scipy.interpolate.PchipInterpolator(cmfs.wavelengths, 
                                                         zbar_y, 
                                                         extrapolate=False)

            # Get additive correction for missing wavelengths (e.g. Nuance EX)
            f_xbar_corr, f_ybar_corr, f_zbar_corr = \
                OdsiDbDataLoader.LoadImage.get_additive_correction(cmf_name, wl)
            
            # Get corrected colour matching function 
            xbar_y = np.nan_to_num(f_xbar(wl)) + np.nan_to_num(f_xbar_corr(wl))
            ybar_y = np.nan_to_num(f_ybar(wl)) + np.nan_to_num(f_ybar_corr(wl))
            zbar_y = np.nan_to_num(f_zbar(wl)) + np.nan_to_num(f_zbar_corr(wl))
            f_xbar = scipy.interpolate.PchipInterpolator(wl, xbar_y, extrapolate=False)
            f_ybar = scipy.interpolate.PchipInterpolator(wl, ybar_y, extrapolate=False)
            f_zbar = scipy.interpolate.PchipInterpolator(wl, zbar_y, extrapolate=False)
            
            return f_xbar, f_ybar, f_zbar

        @staticmethod
        def hyper2rgb(raw_im_hyper: np.ndarray, raw_wl: np.ndarray, 
                cmf_name: str = 'cie_2_1931', illuminant_name: str = 'D65', 
                min_wl: float = 300., max_wl: float = 830.):
            """
            @brief Convert hyperspectral image into RGB for debugging purposes.
            @details The code in this function follows the paper:
                     "Creating RGB images from hyperspectral images using a color
                     matching function" Magnusson et al. 2020.
                     More info here: https://help.commonvisionblox.com/Spectral/
                                             html_english_color-conversion.html

            @param[in]  raw_im_hyper     Hyperspectral band image read from 
                                         ODSI-DB. The intensities are in the 
                                         range [0.0, 1.0].
                                         Shape (h, w, nbands).

            @param[in]  raw_wl           Wavelengths, shape (bands,). 
                                         In the ODSI-DB, the smallest 
                                         wavelength is 397.32nm, and the 
                                         largest is 1003.58nm.

            @param[in]  cmf_name         One of [cie_2_1931, cie_10_1964, 
                                         cie_2_2012, cie_10_2021].
            
            @param[in]  illuminant_name  See colour.SDS_ILLUMINANTS for more 
                                         information.

            @return RGB image in range [0, 255], shape (H, W, 3), dype=np.uint8.
            """
            # Remove bands outside the 300 nm to 830 nm range, they do not
            # provide any information and this is the range where the D65 illuminant 
            # is defined 
            im_hyper, wl = OdsiDbDataLoader.LoadImage.filter_bands(raw_im_hyper, 
                                                                   raw_wl, min_wl, 
                                                                   max_wl)
            h, w, nbands = im_hyper.shape

            # Get colour matching functions
            f_xbar, f_ybar, f_zbar = \
                OdsiDbDataLoader.LoadImage.get_corrected_cmf(cmf_name, wl)

            # Get illuminant function (wavelength -> relative power) 
            il = colour.SDS_ILLUMINANTS[illuminant_name]
            f_illum = scipy.interpolate.PchipInterpolator(il.wavelengths, 
                                                          il.values, 
                                                          extrapolate=True)

            # Convert hyperspectral image to XYZ (normalisation scales the 
            # values so that y = 1 is the white colour)
            one_over_n = 1. / np.trapz(f_ybar(wl) * f_illum(wl), wl)
            x = one_over_n * np.trapz(f_xbar(wl) * im_hyper * f_illum(wl), wl, 
                                      axis=2)
            y = one_over_n * np.trapz(f_ybar(wl) * im_hyper * f_illum(wl), wl, 
                                      axis=2)
            z = one_over_n * np.trapz(f_zbar(wl) * im_hyper * f_illum(wl), wl, 
                                      axis=2)

            # Convert XYZ to RGB (in the range [0, 1])
            red = 3.2406255 * x - 1.5372080 * y - 0.4986286 * z;
            green = -0.9689307 * x + 1.8757561 * y + 0.0415175 * z; 
            blue = 0.0557101 * x - 0.2040211 * y + 1.0569959 * z;

            # Gamma correction (without gamma correction, images look overly 
            # dark)
            red = OdsiDbDataLoader.LoadImage.gamma_correction(red)
            green = OdsiDbDataLoader.LoadImage.gamma_correction(green)
            blue = OdsiDbDataLoader.LoadImage.gamma_correction(blue)

            # Reshape images back to 2D
            im_rgb = np.dstack((red, green, blue))

            # Convert sRGB [0., 1.] to RGB in uint8 [0, 255]
            im_rgb = np.clip((255. * im_rgb).round(), 0, 255).astype(np.uint8)

            return im_rgb

        @staticmethod
        def interp_spectra(im_hyper, wl, new_wl, eps=1., mode='linear'):
            """
            @brief Given a hyperspectral image sampled at certain wavelengths,
                   produce a new hyperspectral image sampled at different wavelengths.
            @param[in]  im_hyper  Hyperspectral image of shape (h, w, c).
            @param[in]  wl        Wavelengths of each of the channels of 'im_hyper', 
                                  shape (c,).
            @param[in]  new_wl    Wavelengths of each of the channels of the wanted
                                  interpolated image to be returned, shape (c,).
            @param[in]  mode      Interpolation mode, 'nearest' or 'linear' available.
            @returns a new hyperspectral image of shape (h, w, c).
            """
            # Make sure that the new wavelengths are not outside the provided ones
            assert(min(wl) - eps < min(new_wl))
            assert(max(wl) + eps > max(new_wl))

            # Compute the distance between wavelengths
            dist = scipy.spatial.distance.cdist(wl.reshape((wl.shape[0], 1)), 
                                                new_wl.reshape((new_wl.shape[0], 1)), 
                                                metric='euclidean')

            # Compute the nearest wavelength for each band
            nearest = np.argmin(dist, axis=0)

            # Synthesize new hyperspectral image
            h, w, _ = im_hyper.shape
            im_hyper_new = np.empty((h, w, new_wl.shape[0]), dtype=im_hyper.dtype)
            if mode == 'nearest':
                # Loop over the new bands
                for i in range(new_wl.shape[0]):
                    im_hyper_new[:, :, i] = im_hyper[:, :, nearest[i]].copy()
            elif mode == 'linear':
                # Loop over the new bands
                for i in range(new_wl.shape[0]):
                    # Find the bands before and after
                    if wl[nearest[i]] > new_wl[i]:  
                        j = nearest[i] - 1
                        k = nearest[i]
                    else:
                        j = nearest[i]
                        k = nearest[i] + 1

                    # Clip values at the extremes
                    if j < 0:
                        j = 0
                    if k == wl.shape[0]: 
                        k = wl.shape[0] - 1

                    # Compute the weights of the bands before and after
                    if j == k:
                        w1 = 0.5
                        w2 = 0.5
                    else:
                        d1 = new_wl[i] - wl[j]
                        d2 = wl[k] - new_wl[i]
                        norm = d1 + d2
                        w1 = d2 / norm
                        w2 = d1 / norm

                    # Add interpolated channel to the image
                    im_hyper_new[:, :, i] = w1 * im_hyper[:, :, j] + w2 * im_hyper[:, :, k]
            
            return im_hyper_new

        @staticmethod
        def read_image(path, mode, new_wl=np.linspace(400, 1000, 204)):
            """
            @returns an image with shape (c, h, w) containing values [0, 1].
            """
            if mode not in OdsiDbDataLoader.modes:
                raise ValueError('[ERROR] The input mode ' + mode \
                    + ' is unknown.')

            # Read raw TIFF file
            im_hyper, wl, _, metadata = torchseg.data_loader.read_stiff(path,
                silent=True, rgb_only=False)
            
            # Select the type of input wanted by the user
            if mode in ['rgbpixel', 'rgbpixel_test', 'rgbimage']:
                # Generate RGB reconstruction
                im_rgb = OdsiDbDataLoader.LoadImage.hyper2rgb(im_hyper, wl)

                # Convert RGB to float and normalise it to the range [0, 1]
                im = im_rgb.astype(np.float32) / 255.

                # Center and scale the image
                im -= OdsiDbDataLoader.rgb_mean
                im /= OdsiDbDataLoader.rgb_std
                
                # Convert the image to (C, H, W)
                im = im.transpose((2, 0, 1)).astype(np.float32)

            elif mode in ['spixel_51', 'simage_51', 
                          'spixel_170', 'spixel_170_test', 'simage_170', 'boiko']:
                # In the ODSI-DB dataset there are images with 51 bands
                # (450-950nm, both inclusive) and images with 204 bands
                # (400-1000nm)
                #
                # The boiko mode comes from "Deep Learning for Dental 
                # Image Analysis" by Boiko et al. 2019

                # Get the array of interpolated wavelengths according to the
                # input mode selected by the user 
                new_wl = OdsiDbDataLoader.mode2wl[mode]

                # Interpolate hyperspectral image to the requested range
                interp_func = dl.OdsiDbDataLoader.LoadImage.interp_spectra
                im = interp_func(im_hyper, wl, new_wl)

                # Center and scale the image
                im -= OdsiDbDataLoader.mode2mean[mode]
                im /= OdsiDbDataLoader.mode2std[mode]

                # Convert the image to (C, H, W)
                im = im.transpose((2, 0, 1)).astype(np.float32)

            elif mode in ['spixel_204', 'simage_204']:
                # Get the array of interpolated wavelengths according to the
                # input mode selected by the user 
                new_wl = OdsiDbDataLoader.mode2wl[mode]

                # Get the function that will interpolate the hyperspectral
                # images into a fixed set of wavelengths
                interp_func = dl.OdsiDbDataLoader.LoadImage.interp_spectra
                
                # If the image comes from the Nuance EX, let's amend it
                if im_hyper.shape[2] == 51:
                    # Get an array of the wavelengths in [450, 950]
                    inside_indices = np.array([idx for idx, x in enumerate(new_wl.tolist()) \
                        if x > 450. and x < 950.])
                    inside_wl = new_wl[inside_indices]

                    # Interpolate hyperspectral image to the requested range
                    inside_im = interp_func(im_hyper, wl, inside_wl)

                    # Extrapolate to 204 bands 
                    nbands = 204
                    im = np.empty((im_hyper.shape[0], im_hyper.shape[1], nbands), 
                                  dtype=inside_im.dtype)
                    j = 0
                    for i in range(nbands):
                        if i < np.min(inside_indices):
                            im[:, :, i] = inside_im[:, :, 0]
                        elif i > np.max(inside_indices):
                            im[:, :, i] = inside_im[:, :, -1]
                        else:
                            im[:, :, i] = inside_im[:, :, j] 
                            j += 1

                    # Center and scale the image
                    im -= OdsiDbDataLoader.mode2mean[mode]
                    im /= OdsiDbDataLoader.mode2std[mode]

                    # Convert the image to (C, H, W)
                    im = im.transpose((2, 0, 1)).astype(np.float32)
                else:
                    # Interpolate hyperspectral image to the requested range
                    im = interp_func(im_hyper, wl, new_wl)
                
                    # Center and scale the image
                    im -= OdsiDbDataLoader.mode2mean[mode]
                    im /= OdsiDbDataLoader.mode2std[mode]

                    # Convert the image to (C, H, W)
                    im = im.transpose((2, 0, 1)).astype(np.float32)
            else:
                raise ValueError('[ERROR] ODSI-DB mode unknown.')

            return im
        
        @staticmethod
        def read_label(path):
            raw_tiff = torchseg.data_loader.read_mtiff(path)

            # Get the shape of the labels (which should be identical to the
            # height and width of the image)
            shape = raw_tiff[list(raw_tiff.keys())[0]].shape

            # Create tensor of labels
            n_classes = len(OdsiDbDataLoader.OdsiDbDataset.classnames)
            label = np.zeros((n_classes, *shape), dtype=np.float32)
            
            # Build dictionaries to convert quickly from index to class name
            # and vice versa
            idx2class = OdsiDbDataLoader.OdsiDbDataset.classnames
            class2idx = {y: x for x, y in idx2class.items()}
            
            # Populate the binary array for each class
            for k, gt in raw_tiff.items():
                if k not in class2idx:
                    raise ValueError('[ERROR] Unknown <' + k + '> label.')
                
                # Find the index of the class with name 'k'
                idx = class2idx[k]

                # Set the label for the class with name 'k' 
                label[idx] = gt.astype(np.float32)

            # If all the classes are zero for a pixel, we activate the 
            # background
            #label[class2idx['Background'], label.sum(axis=0) == 0] = 1

            return label

        def forward(self, data):
            """
            @param[in]  data  File name to read.
            """
            # Open TIF image
            im = OdsiDbDataLoader.LoadImage.read_image(data['image'], self.mode)

            # Open ground truth segmentation
            label = OdsiDbDataLoader.LoadImage.read_label(data['label'])
            
            # Make default output
            dic = {
                'image': im, 
                'label': label, 
                'mode' : self.mode,
                'path' : data['image'],
            }
            
            '''
            # If the user wants pixel-wise multispectral predictions, i.e. what
            # we call 'spixel' mode, we change the output to contain just a
            # random pixel from the image
            #if self.mode == 'spixel' and self.training:
            if self.mode == 'spixel':
                # Find all the pixels properly labelled
                ann_idx = np.array((np.sum(label, axis=0) == 1).nonzero())
                
                # Take only labelled pixels, shape: (npixels, bands)
                pixels = im[:, ann_idx[0], ann_idx[1]].transpose((1, 0))
                pixel_labels = label[:, ann_idx[0], ann_idx[1]].transpose((1, 0))

                # Reshape (npixels, bands, 1, 1)
                pixels = np.expand_dims(pixels, axis=(2, 3))
                pixel_labels = np.expand_dims(pixel_labels, axis=(2, 3))
                
                dic = {'image': pixels, 'label': pixel_labels, 'mode': self.mode}

                # Choose a random pixel
                #pix = np.random.randint(0, ann_idx.shape[1])
                #rand_row = ann_idx[0, pix]
                #rand_col = ann_idx[1, pix]
                    
                # Get annotated hyperspectral pixel
                #pixel = im[:, rand_row, rand_col].reshape(
                #    (im.shape[0], 1, 1))
                #pixel_label = label[:, rand_row, rand_col].reshape(
                #    (label.shape[0], 1, 1))
                #dic = {'image': pixel, 'label': pixel_label, 'mode': self.mode}

                # TODO: Make many images of 1x1 pixel instead of picking a
                #        pixel randomly
                #dic = {'image': [], 'label': []}
                #for i in range(im.shape[1]):
                #    for j in range(im.shape[2]):
                #        pixel = im[:, i, j].reshape((im.shape[0], 1, 1))
                #        pixel_label = label[:, i, j].reshape((label.shape[0], 1, 1))
                #        dic['image'].append(pixel)
                #        dic['label'].append(pixel_label)
            '''

            return dic
        
        #@property
        #def training(self):
        #    return self._training

        #@training.setter
        #def training(self, training):
        #    self._training = training


    class OdsiDbDataset(monai.data.dataset.PersistentDataset):
        # Class indices and names for the ODSI-DB dataset
        classnames = {
            0 : u'Attached gingiva',
            1 : u'Attrition/Erosion',
            2 : u'Blood vessel',
            3 : u'Calculus',
            4 : u'Dentine caries',
            5 : u'Enamel',
            6 : u'Fibroma',
            7 : u'Fluorosis',
            8 : u'Gingivitis',
            9 : u'Hair',
            10: u'Hard palate',
            11: u'Inflammation',
            12: u'Initial caries',
            13: u'Leukoplakia',
            14: u'Lip',
            15: u'Makeup',
            16: u'Malignant lesion',
            17: u'Marginal gingiva',
            18: u'Metal',
            19: u'Microfracture',
            20: u'Mole',
            21: u'Oral mucosa',
            22: u'Out of focus area',
            23: u'Pigmentation',
            24: u'Plaque',
            25: u'Plastic',
            26: u'Prosthetics',
            27: u'Root',
            28: u'Shadow/Noise',
            29: u'Skin',
            30: u'Soft palate',
            31: u'Specular reflection',
            32: u'Stain',
            33: u'Tongue',
            34: u'Ulcer',
        }
        classnames_reverse = {v:k for k, v in classnames.items()}

        #def __init__(self, data_dir, mode='rgb', training='True'):
        def __init__(self, data_dir, mode='rgbimage'):
            """
            @param[in]  mode  Input type, see LoadImage constructor for more
                              details.
            """
            # Store parameters
            assert(mode in OdsiDbDataLoader.modes)
            self.data_dir = data_dir 
            self.mode = mode

            # Get list of TIFF files
            tiffs = [f for f in torchseg.utils.listdir(self.data_dir) \
                if '.tif' in f]

            # Get list of segmentation files
            segs = [f for f in tiffs if f.endswith('_masks.tif')]

            # Get list of image files (we ignore images without segmentation)
            imgs = [f.replace('_masks.tif', '.tif') for f in segs]
            for im in imgs:
                assert(os.path.isfile(os.path.join(data_dir, im)))

            # Build list of dictionaries, where each dictionary 
            # corresponds to an input image
            data = [{
                'image': os.path.join(data_dir, im), 
                'label': os.path.join(data_dir, seg)} \
                    for im, seg in zip(imgs, segs)]

            # Create data augmentation transform
            self.image_loader = OdsiDbDataLoader.LoadImage(
                keys=['image', 'label'], mode=self.mode)
                #keys=['image', 'label'], mode=self.mode, training=training)
            trsfm = torchvision.transforms.Compose([
                self.image_loader,
                monai.transforms.ToTensord(keys=['image', 'label']),
            ])
            
            # Call PersistentDataset constructor
            super().__init__(data, trsfm)

        #@property
        #def training(self):
        #    return self.image_loader.training

        #@training.setter
        #def training(self, training):
        #    self.image_loader.training = training


    """
    @class OdsiDbLoader loads data from the ODSI-DB dataset.
    @details Link to the dataset: https://cs.uef.fi/pub/color/spectra/ODSI-DB 
    """
    def __init__(self, data_dir, batch_size, mode='rgbimage', shuffle=True, 
            validation_split=0.0, num_workers=1, ignore_labels=[]):
        #    validation_split=0.0, num_workers=1, training=True):
        # Store parameters
        self.data_dir = data_dir
        self.mode = mode
        self.ignore_labels = ignore_labels

        # Create dataset handler
        self.dataset = OdsiDbDataLoader.OdsiDbDataset(self.data_dir, 
            mode=self.mode)
        #    mode=self.mode, training=training)

        # Call the constructor of torchseg.base.BaseDataLoader
        super().__init__(self.dataset, batch_size, shuffle, validation_split, 
            num_workers, collate_fn=self.collate_fn)
        #    num_workers, collate_fn=OdsiDbDataLoader.collate_fn) 

    def _split_sampler(self, split):
        """
        @brief Split the list of inputs between training and validation.
        @param[in]  split  Ratio between 0.0 and 1.0 that will correspond to 
                           validation.
        @returns TODO
        """
        if split == 0.0:
            return None, None

        if isinstance(split, int):
            assert(split > 0 and split < self.n_samples) 
            split = float(split) / self.n_samples

        # Fix seeds for reproducibility
        np.random.seed(0)
        random.seed(0)

        # Make a list of the indices of all the images
        remaining = np.arange(self.n_samples).tolist()
        
        # Get list of labels to check the classes that are present in each image
        labels = [os.path.join(self.data_dir, im_dic['label']) \
            for im_dic in self.dataset.data]

        # Make a dictionary from each class to the indices of the images that
        # contain pixels of such class
        n_classes = len(OdsiDbDataLoader.OdsiDbDataset.classnames)
        class2img = {c: [] for c in range(n_classes)}
        for im_idx in remaining:
            label_path = labels[im_idx] 
            label = OdsiDbDataLoader.LoadImage.read_label(label_path).sum(axis=(1,
                2)).astype(np.int64)
            for c in label.nonzero()[0].tolist():
                class2img[c].append(im_idx)
        
        # Put an image of each class in the training split, unless we are in 
        # testing mode, in which case self.validation split should be equal to 1.0
        train_idx = []
        if self.validation_split < 1.0:
            for c in class2img:
                random.shuffle(class2img[c])
                im_idx = class2img[c][0]
                if im_idx in remaining: 
                    train_idx.append(im_idx)
                    remaining.remove(im_idx)
        
        # Split the rest of the images according to the given probability
        valid_idx = []
        while remaining:
            im_idx = remaining.pop()
            if np.random.binomial(1, 1 - split):  
                train_idx.append(im_idx)
            else:
                valid_idx.append(im_idx)
        train_idx = np.asarray(train_idx)
        valid_idx = np.asarray(valid_idx)

        # Create data samplers for each split
        train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_idx)
        valid_sampler = torch.utils.data.sampler.SubsetRandomSampler(valid_idx)
        
        # Turn off shuffle option which is mutually exclusive with sampler
        self.shuffle = False
        self.n_samples = len(train_idx)

        return train_sampler, valid_sampler
        
    #@staticmethod
    def collate_fn(self, batch, max_h=256, max_w=256):
        """
        @brief This function hijacks the batch and crops all the elements of the
               batch randomly so that they are the same size. The minimum width
               and height of the images in the batch is used as crop size.
        @param[in]  batch  List of dictionaries. Each dictionary represents an 
                           input image.
        @param[in]  max_h  Maximum height for each image in the batch. If an image
                           has a larger height it will be cropped. This only 
                           applies to the 'spixel' mode.
        @param[in]  max_w  Maximum width for each image in the batch. If an image
                           is larger it will be cropped. This only applies to the
                           'spixel' mode.
        """
        # If we are in pixel mode we crop a region of 10 * bs x 10 * bs pixels, 
        # where bs is the batch size
        #if 'mode' in batch[0] and batch[0]['mode'] == 'spixel':
            # Reshape the batch, one pixel per batch element
            #bands = batch[0]['image'].shape[0]
            #classes = batch[0]['label'].shape[0]
            #pixels_per_im = min_w * min_h
            #new_batch = []
            #for item in batch:
            #    im = item['image'].reshape((bands, -1)).permute(1, 
            #        0).reshape((pixels_per_im, bands, 1, 1)).unbind()
            #    label = item['label'].reshape((classes, -1)).permute(1, 
            #        0).reshape((pixels_per_im, classes, 1, 1)).unbind()
            #    new_batch += [{'image': pix, 'label': pix_label} \
            #        for pix, pix_label in zip(im, label)]
            #batch = new_batch
        '''

        if 'mode' in batch[0] and batch[0]['mode'] == 'spixel':
            # Find minimum number of pixels
            min_pixels = np.inf
            for item in batch:
                if item['image'].shape[0] < min_pixels:
                    min_pixels = item['image'].shape[0]

            # Reduce the number of pixels for performance reasons
            #max_npixels = 128
            #min_pixels = min(min_pixels, max_npixels)
            
            # Crop number of pixels so that all the elements in the batch have
            # the same
            new_batch = []
            for item in batch:
                npixels = item['image'].shape[0]
                diff = npixels - min_pixels 
                if diff > 0:
                    offset = np.random.randint(0, diff) 
                    item['image'] = item['image'][offset:offset + min_pixels, :]
                    item['label'] = item['label'][offset:offset + min_pixels, :]

                # Reshape the batch, we want a batch of pixels
                pixels = item['image'].unbind()
                labels = item['label'].unbind()
                new_batch += [{'image': pix, 'label': pixl} \
                    for pix, pixl in zip(pixels, labels)]

            # Use batch of pixels
            batch = new_batch
        else:
        '''

        # Find minimum dimensions
        min_h = np.inf
        min_w = np.inf
        for item in batch: 
            if item['image'].shape[1] < min_h:
                min_h = item['image'].shape[1]
            if item['image'].shape[2] < min_w:
                min_w = item['image'].shape[2]

        # Crop image for the 'rgbpixel' and 'spixel' cases due to memory and 
        # speed restrictions
        #if self.training:
        mode = batch[0].get('mode')
        if mode in ['rgbpixel', 'spixel_170']:
            min_h = min(min_h, max_h)
            min_w = min(min_w, max_w)

        # Crop image and labels
        for item in batch:
            good_crop = False
            im = None 
            label = None
            while not good_crop:
                # Try a crop
                im, label = OdsiDbDataLoader.random_crop(item['image'], item['label'], min_h, min_w)
                #item['image'], item['label'] = OdsiDbDataLoader.random_crop(item['image'], item['label'], min_h, min_w)

                # Does it have annotated pixels?
                ann_idx = torch.sum(label, dim=0) == 1
                num_valid_pixels = ann_idx.sum()
                if num_valid_pixels > 0:
                    good_crop = True
            item['image'] = im
            item['label'] = label

        return torch.utils.data.dataloader.default_collate(batch)

    @staticmethod
    def random_crop(image, label, crop_h, crop_w):
        """
        @param[in]  image   Torch tensor, shape (C, H, W).
        @param[in]  label   Torch tensor, shape (C, H, W).
        @param[in]  crop_h  Crop height.
        @param[in]  crop_w  Crop width.
        """
        # Get image dimensions
        h = image.shape[1]
        w = image.shape[2]
        assert(h >= crop_h)
        assert(w >= crop_w)

        # Compute random crop location
        if crop_h < h:
            off_row = np.random.randint(0, h - crop_h, size=1)[0]
        else:
            off_row = 0
        if crop_w < w:
            off_col = np.random.randint(0, w - crop_w, size=1)[0]
        else:
            off_col = 0

        # Crop image and label
        new_image = image[:, off_row:off_row + crop_h, off_col:off_col + crop_w]
        new_label = label[:, off_row:off_row + crop_h, off_col:off_col + crop_w]

        return new_image, new_label
        

