Upload 32 files

.gitattributes

@@ -0,0 +1,16 @@
+GANFilling/data/test_data/context_29TQF_2017-08-16_2018-01-12_5049_5177_441_569_78_158_6_86.npz filter=lfs diff=lfs merge=lfs -text
+GANFilling/data/test_data/context_30UYV_2018-06-19_2018-11-15_1977_2105_2617_2745_30_110_40_120.npz filter=lfs diff=lfs merge=lfs -text
+GANFilling/data/test_data/context_33SXD_2019-05-16_2019-10-12_4025_4153_3129_3257_62_142_48_128.npz filter=lfs diff=lfs merge=lfs -text
+GANFilling/data/test_data/context_34TDP_2018-04-28_2018-09-24_2233_2361_5049_5177_34_114_78_158.npz filter=lfs diff=lfs merge=lfs -text
+GANFilling/data/train_data/33UXQ/33UXQ_2018-03-30_2018-08-26_2873_3001_3513_3641_44_124_54_134.npz filter=lfs diff=lfs merge=lfs -text
+GANFilling/data/train_data/33UXQ/33UXQ_2018-06-18_2018-11-14_2105_2233_3769_3897_32_112_58_138.npz filter=lfs diff=lfs merge=lfs -text
+GANFilling/data/train_data/33UXQ/33UXQ_2018-06-18_2018-11-14_2745_2873_3897_4025_42_122_60_140.npz filter=lfs diff=lfs merge=lfs -text
+GANFilling/data/train_data/33UXQ/33UXQ_2018-07-08_2018-12-04_2617_2745_3513_3641_40_120_54_134.npz filter=lfs diff=lfs merge=lfs -text
+GANFilling/data/train_data/33UXQ/33UXQ_2018-07-08_2018-12-04_697_825_4793_4921_10_90_74_154.npz filter=lfs diff=lfs merge=lfs -text
+GANFilling/data/train_data/33UXQ/33UXQ_2018-07-18_2018-12-14_825_953_4537_4665_12_92_70_150.npz filter=lfs diff=lfs merge=lfs -text
+GANFilling/data/train_data/33UXQ/33UXQ_2018-07-28_2018-12-24_4025_4153_2105_2233_62_142_32_112.npz filter=lfs diff=lfs merge=lfs -text
+GANFilling/GANFilling_percept_results.png filter=lfs diff=lfs merge=lfs -text
+GANFilling/results/context_29TQF_2017-08-16_2018-01-12_5049_5177_441_569_78_158_6_86.npz.png filter=lfs diff=lfs merge=lfs -text
+GANFilling/results/context_30UYV_2018-06-19_2018-11-15_1977_2105_2617_2745_30_110_40_120.npz.png filter=lfs diff=lfs merge=lfs -text
+GANFilling/results/context_33SXD_2019-05-16_2019-10-12_4025_4153_3129_3257_62_142_48_128.npz.png filter=lfs diff=lfs merge=lfs -text
+GANFilling/results/context_34TDP_2018-04-28_2018-09-24_2233_2361_5049_5177_34_114_78_158.npz.png filter=lfs diff=lfs merge=lfs -text

GANFilling/.gitattributes

@@ -0,0 +1 @@
+*.pt filter=lfs diff=lfs merge=lfs -text

GANFilling/.gitignore

@@ -0,0 +1,3 @@
+# pycache
+__pycache__
+*/__pycache__

GANFilling/README.md

@@ -0,0 +1,62 @@
+# Generative Networks for Spatio-Temporal Gap Filling of Sentinel-2 Reflectances
+---------------
+| [Journal ISPRS paper](https://doi.org/10.1016/j.isprsjprs.2025.01.016) |
+
+# Abstract
+
+Earth observation from satellite sensors offers the possibility to monitor natural ecosystems by deriving spatially explicit and temporally resolved biogeophysical parameters. Optical remote sensing, however, suffers from missing data mainly due to the presence of clouds, sensor malfunctioning, and atmospheric conditions. This study proposes a novel deep learning architecture to address gap filling of satellite reflectances, more precisely the visible and near-infrared bands, and illustrates its performance at high-resolution Sentinel-2 data. We introduce GANFilling, a generative adversarial network capable of sequence-to-sequence translation,  which comprises convolutional long short-term memory layers to effectively exploit complete dependencies in space-time series data. We focus on Europe and evaluate the method's performance quantitatively (through distortion and perceptual metrics) and qualitatively (via visual inspection and visual quality metrics). Quantitatively, our model offers the best trade-off between denoising corrupted data and preserving noise-free information, underscoring the importance of considering multiple metrics jointly when assessing gap filling tasks. Qualitatively, it successfully deals with various noise sources, such as clouds and missing data, constituting a robust solution to multiple scenarios and settings. We also illustrate and quantify the quality of the generated product in the relevant downstream application of vegetation greenness forecasting, where using GANFilling enhances forecasting in approximately 70% of the considered regions in Europe. This research contributes to underlining the utility of deep learning for Earth observation data, which allows for improved spatially and temporally resolved monitoring of the Earth surface.
+
+# Usage and Requirements
+
+## Train
+
+`
+python src/train.py --dataPath data/train_data --cleanDataPath src/utils/example_clean_data.json --name model_1
+`
+
+Models are saved to ./trained_models/ (can be changed by passing --modelsPath=your_dir in train.py). 
+
+For training the discriminator, a json file with the noiseless(real) data has to be generated. This can be done with the ./utils/generate_cleanData_file.py. An example of the structure of this file can be found in ./utils/generate_cleanData_file.py.  
+
+The folder with the training data is specified with the `--dataPath` argument. This repository just includes few samples in data/train_data/ as a reference for structure and behaviour checking. 
+
+## Test 
+
+`
+python src/test.py --model_path trained_models/GANFilling.pt --data_path data --results_path results
+`
+
+This will run the GANFilling trained model on a small set of examples and generate the corresponding gap filled time series. 
+
+To test your own model modify the parameter `--model_path` in test.py.
+To test on your own data modify the parameter `--data_path` in test.py.
+
+# Results
+
+![image](GANFilling_percept_results.png)
+
+<b>Example images showing the GANFilling reconstruction on different land covers.</b> For each example, the first row shows the land cover map followed by ten original time steps of a visible (RGB) sequence, while the second row corresponds to its noise-free version. All images are noted with the day-of-year (DOY). The third row illustrates the NDVI maps for the noise-free images. Different types of noise are outlined in red. (A) Complexe scene with predominant herbaceous vegetation and multiple frames with complete loss of information. (B) Sequence with mostly cultivated areas to show the performance on fast changes in the Earth’s surface with heavily occluded frames. (C) Sequence with widespread vines characterized by a rapid evolution of the land cover. (D) Predominant coniferous tree cover with a water body nearby. (E) Sequence with predominant broadleaf tree cover and several consecutive time steps with dense occlusions. Land cover’s legend: <span style="color:rgb(255,255,255)">&#9723;</span> Cloud or No data, <span style="color:rgb(210,0,0)">&#9724;</span> Artificial surfaces and constructions, <span style="color:rgb(253,211,39)">&#9724;</span> Cultivated areas, <span style="color:rgb(176,91,16)">&#9724;</span>  Vineyards, <span style="color:rgb(35,152,0)">&#9724;</span>  Broadleaf tree cover, <span style="color:rgb(8,98,0)">&#9724;</span>  Coniferous tree cover, <span style="color:rgb(249,150,39)">&#9724;</span>  Herbaceous vegetation, <span style="color:rgb(141,139,0)">&#9724;</span>  Moors and Heathland, <span style="color:rgb(95,53,6)">&#9724;</span>  Sclerophyllous vegetation, <span style="color:rgb(149,107,196)">&#9724;</span>  Marshes, <span style="color:rgb(77,37,106)">&#9724;</span>  Peatbogs, <span style="color:rgb(154,154,154)">&#9724;</span>  Natural material surfaces, <span style="color:rgb(106,255,255)">&#9724;</span>  Permanent snow covered surfaces, <span style="color:rgb(20,69,249)">&#9724;</span>  Water bodies.
+
+# How to cite
+
+If you use this code for your research, please cite our paper Generative Networks for Spatio-Temporal Gap Filling of Sentinel-2 Reflectances:
+
+```
+@article{GonzalezCalabuig2025,
+title = {Generative networks for spatio-temporal gap filling of Sentinel-2 reflectances},
+journal = {ISPRS Journal of Photogrammetry and Remote Sensing},
+volume = {220},
+pages = {637-648},
+year = {2025},
+issn = {0924-2716},
+doi = {https://doi.org/10.1016/j.isprsjprs.2025.01.016},
+author = {Maria Gonzalez-Calabuig and Miguel-Ángel Fernández-Torres and Gustau Camps-Valls}}
+```
+
+# Aknowledgments
+
+Authors acknowledge the support from the European Research Council (ERC) under the ERC SynergyGrant USMILE (grant agreement 855187), the European Union’s Horizon 2020 research and innovation program within the projects ‘XAIDA: Extreme Events- Artificial Intelligence for Detection and Attribution,’(grant agreement 101003469), ‘DeepCube: Explainable AI pipelines for big Copernicus data’ (grant agreement 101004188), the ESA AI4Science project ”MultiHazards, Compounds and Cascade events: DeepExtremes”, 2022-2024, the computer resources provided by the J¨ulich Supercomputing Centre (JSC) (Project No.PRACE-DEV-2022D01-048), the computer resources provided by Artemisa (funded by the European Union ERDF and Comunitat Valenciana), as well as the technical support provided by the Instituto de Física Corpuscular, IFIC (CSIC-UV).
+
+# License
+
+[MIT]()

GANFilling/data/landcover_types.csv

@@ -0,0 +1,15 @@
+0,255,255,255,255,Clouds or No data
+62,210,0,0,255,Artificial surfaces and constructions
+73,253,211,39,255,Cultivated areas
+75,176,91,16,255,Vineyards
+82,35,152,0,255,Broadleaf tree cover
+83,8,98,0,255,Coniferous tree cover
+102,249,150,39,255,Herbaceous vegetation
+103,141,139,0,255,Moors and Heathland
+104,95,53,6,255,Sclerophyllous vegetation 
+105,149,107,196,255,Marshes
+106,77,37,106,255,Peatbogs 
+121,154,154,154,255,Natural material surfaces 
+123,106,255,255,255,Permanent snow covered surfaces
+162,20,69,249,255,Water bodies
+255,255,255,255,255,No data

GANFilling/data/test_data/context_29TQF_2017-08-16_2018-01-12_5049_5177_441_569_78_158_6_86.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ff96d1e04dc8f23e1caa09dfb3b65adc4111393729ed3b963a87b1b5185c048
+size 4735175

GANFilling/data/test_data/context_30UYV_2018-06-19_2018-11-15_1977_2105_2617_2745_30_110_40_120.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:603f455b64f8044910fa78b6dee26f335b924f2fb6f9cbf03539bfbae6c4f037
+size 4198884

GANFilling/data/test_data/context_33SXD_2019-05-16_2019-10-12_4025_4153_3129_3257_62_142_48_128.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e15904f0feee4a86d1ef931d97d34a3b5b7d4e855aed218fad0bfd8dd50ca839
+size 3328672

GANFilling/data/test_data/context_34TDP_2018-04-28_2018-09-24_2233_2361_5049_5177_34_114_78_158.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:716193f3fea3641ce5c94b4169fb047f6ac224ae5d13720cf0a413b2520322fe
+size 5594634

GANFilling/data/train_data/33UXQ/33UXQ_2018-03-30_2018-08-26_2873_3001_3513_3641_44_124_54_134.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:372e0bf759d6856b57dd9a5295b9cb03d1ad1507940d4d8b6be582d6d09f4939
+size 8178516

GANFilling/data/train_data/33UXQ/33UXQ_2018-06-18_2018-11-14_2105_2233_3769_3897_32_112_58_138.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7127004d5dd01befad7f786fea024f522a75fb8942a49a220b7843a36e1fb7f
+size 8292551

GANFilling/data/train_data/33UXQ/33UXQ_2018-06-18_2018-11-14_2745_2873_3897_4025_42_122_60_140.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a9bcebacc479854a4715c8096a7ff6d90de208f2e6df0259a3b594390902b28d
+size 8150040

GANFilling/data/train_data/33UXQ/33UXQ_2018-07-08_2018-12-04_2617_2745_3513_3641_40_120_54_134.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9ff886ee2e13a9201a3ec2341bea317ba6eb28bc944dd458e5c51b2c94603d2
+size 8096535

GANFilling/data/train_data/33UXQ/33UXQ_2018-07-08_2018-12-04_697_825_4793_4921_10_90_74_154.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd8ce22607591b366ba2b4ac61a4df19b0fd58c2f7cfaa94504f6d6599134e2f
+size 8291712

GANFilling/data/train_data/33UXQ/33UXQ_2018-07-18_2018-12-14_825_953_4537_4665_12_92_70_150.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d522b678b06e37d7ce62b4d898bed0fec187a72711efe3c7fc4d7dada43b083
+size 8400326

GANFilling/data/train_data/33UXQ/33UXQ_2018-07-28_2018-12-24_4025_4153_2105_2233_62_142_32_112.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:143075cf0da3467a2d25ff4a69f290cf497f43b846147bb9b13167b5c4710068
+size 7970874

GANFilling/src/iterator.py

@@ -0,0 +1,111 @@
+import numpy as np
+import json
+import torch
+import random
+import os
+
+from kornia.augmentation import RandomHorizontalFlip, RandomVerticalFlip
+
+
+class Iterator():
+    def __init__(self, dataPath, cleanDataPath, mode):
+        self.mode = mode
+        self.random_seed_data = 42
+        self.__initData(dataPath, cleanDataPath)
+    
+    def __initData(self, dataPath, cleanDataPath):
+        self.data = {}
+        
+        if self.mode == 'train':
+            self.__loadTrainData(dataPath)
+            self.cleanData = json.load(open(cleanDataPath))
+        
+    def __loadTrainData(self, dataPath):
+        tiles = os.listdir(dataPath)
+        if 'LICENSE' in tiles:
+            tiles.remove('LICENSE')
+        tiles.sort()
+
+        loadedData = []
+        for tile in tiles:
+            in_tile_path = dataPath / tile
+            files = os.listdir(in_tile_path)
+            files.sort()
+            in_files = []
+            for file in files:
+                in_files.append(os.path.join(in_tile_path, file))
+            random.Random(self.random_seed_data).shuffle(in_files)
+            for f in in_files[:int(len(in_files)*0.8)]:
+                loadedData.append(f)
+        self.data = loadedData
+                
+    def __len__(self):
+        return len(self.data)
+
+    def __getMasks(self, sample, index):
+        return sample["highresdynamic"][:,:,-1,:index+1]
+
+    def __find_closest_area(self, condition_1):
+        self.cleanData.keys()
+        if condition_1 in self.cleanData.keys():
+            return condition_1
+        else:
+            if str(int(condition_1[:2])+1) + condition_1[-1] in self.cleanData.keys():
+                return str(int(condition_1[:2])+1) + condition_1[-1]
+            elif str(int(condition_1[:2])-1) + condition_1[-1] in self.cleanData.keys():
+                return str(int(condition_1[:2])-1) + condition_1[-1]
+
+    def __getCleanSequence(self, condition_1, condition_2):
+        time_lenght = 4
+        count_stuck = 0
+        condition_1 = self.__find_closest_area(condition_1)
+        data = self.cleanData[condition_1]
+        rand_number = torch.randint(len(data), (1,)).item()
+        attributes = data[rand_number]
+        while len(attributes['time steps']) < time_lenght:
+            rand_number = torch.randint(len(data), (1,)).item()
+            attributes = data[rand_number]
+            # Check end of the loop:
+            '''count_stuck += 1
+            if count_stuck >= 2:
+                print('Stuck ', count_stuck)'''
+
+        kernelSize = attributes['kernel size']
+        sample = np.load(attributes['path'])
+        x_min = attributes['bbox'][1]
+        x_max = attributes['bbox'][1]+kernelSize
+        y_min = attributes['bbox'][0]
+        y_max = attributes['bbox'][0]+kernelSize
+
+        discriminator_sample = sample["highresdynamic"][y_min:y_max,x_min:x_max, 0:4, attributes['time steps'][0]:attributes['time steps'][0]+time_lenght]
+        discriminator_sample = torch.from_numpy(discriminator_sample)
+
+        ## Data augmentation: Horitzontal Flip
+        transformation_1 = RandomHorizontalFlip(p=0.4)
+        discriminator_sample[:,:,:,0] = transformation_1(discriminator_sample[:,:,:,0])
+        for t in range(1, discriminator_sample.shape[3]):
+            discriminator_sample[:,:,:,t] = transformation_1(discriminator_sample[:,:,:,t], params=transformation_1._params)
+        
+        ## Data augmentation: Vertical Flip
+        transformation_2 = RandomVerticalFlip(p=0.4)
+        discriminator_sample[:,:,:,0] = transformation_2(discriminator_sample[:,:,:,0])
+        for t in range(1, discriminator_sample.shape[3]):
+            discriminator_sample[:,:,:,t] = transformation_2(discriminator_sample[:,:,:,t], params=transformation_2._params)
+
+        return discriminator_sample.numpy()
+
+        
+    def __getitem__(self, index):
+        self.index = index
+        sample = np.load(self.data[index])
+
+        context = 10    
+        noisyImg = sample["highresdynamic"][:,:,0:4,:context]
+
+        masks = self.__getMasks(sample, context-1)
+        cleanImg = self.__getCleanSequence(self.data[index].split('/')[3][:3], None)
+
+        noisyImg = np.nan_to_num(np.clip(noisyImg, 0, 1), nan=1.0)
+        cleanImg = np.nan_to_num(np.clip(cleanImg, 0, 1), nan=1.0)
+
+        return np.transpose(noisyImg, (2,0,1,3)), np.transpose(cleanImg, (2,0,1,3)), masks

GANFilling/src/models/convlstm.py

@@ -0,0 +1,200 @@
+import torch.nn.functional as F
+import torch.nn as nn
+import torch
+
+
+class ConvLSTMCell(nn.Module):
+
+    def __init__(self, input_dim, hidden_dim, kernel_size, bias, device):
+        """
+        Initialize ConvLSTM cell.
+
+        Parameters
+        ----------
+        input_dim: int
+            Number of channels of input tensor.
+        hidden_dim: int
+            Number of channels of hidden state.
+        kernel_size: (int, int)
+            Size of the convolutional kernel.
+        bias: bool
+            Whether or not to add the bias.
+        """
+
+        super(ConvLSTMCell, self).__init__()
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+
+        self.kernel_size = kernel_size
+        self.padding = kernel_size[0] // 2, kernel_size[1] // 2
+        self.bias = bias
+        self.device = device
+
+        self.conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,
+                              out_channels=4*self.hidden_dim,
+                              kernel_size=self.kernel_size,
+                              padding=self.padding,
+                              bias=self.bias)
+
+    def __initStates(self, size):
+        return torch.zeros(size).to(self.device), torch.zeros(size).to(self.device)
+        #return torch.zeros(size).cuda(), torch.zeros(size).cuda()
+
+    def forward(self, input_tensor, cur_state):
+        if cur_state == None:
+            h_cur, c_cur = self.__initStates([input_tensor.shape[0], self.hidden_dim, input_tensor.shape[2], input_tensor.shape[3]])
+        else:
+            h_cur, c_cur = cur_state
+        
+        combined = torch.cat([input_tensor, h_cur], dim=1)  # concatenate along channel axis
+        combined_conv = self.conv(combined)
+        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)
+
+        i = torch.sigmoid(cc_i)
+        f = torch.sigmoid(cc_f)
+        o = torch.sigmoid(cc_o)
+        g = torch.tanh(cc_g)
+
+        c_next = f * c_cur + i * g
+        h_next = o * torch.tanh(c_next)
+
+        return h_next, c_next
+
+    def init_hidden(self, batch_size, image_size):
+        height, width = image_size
+        return (torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device),
+                torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device))
+
+
+class ConvLSTM(nn.Module):
+
+    """
+
+    Parameters:
+        input_dim: Number of channels in input
+        hidden_dim: Number of hidden channels
+        kernel_size: Size of kernel in convolutions
+        num_layers: Number of LSTM layers stacked on each other
+        batch_first: Whether or not dimension 0 is the batch or not
+        bias: Bias or no bias in Convolution
+        return_all_layers: Return the list of computations for all layers
+        Note: Will do same padding.
+
+    Input:
+        A tensor of size B, T, C, H, W or T, B, C, H, W
+    Output:
+        A tuple of two lists of length num_layers (or length 1 if return_all_layers is False).
+            0 - layer_output_list is the list of lists of length T of each output
+            1 - last_state_list is the list of last states
+                    each element of the list is a tuple (h, c) for hidden state and memory
+    Example:
+        >> x = torch.rand((32, 10, 64, 128, 128))
+        >> convlstm = ConvLSTM(64, 16, 3, 1, True, True, False)
+        >> _, last_states = convlstm(x)
+        >> h = last_states[0][0]  # 0 for layer index, 0 for h index
+    """
+
+    def __init__(self, input_dim, hidden_dim, kernel_size, num_layers,
+                 batch_first=False, bias=True, return_all_layers=False):
+        super(ConvLSTM, self).__init__()
+
+        self._check_kernel_size_consistency(kernel_size)
+
+        # Make sure that both `kernel_size` and `hidden_dim` are lists having len == num_layers
+        kernel_size = self._extend_for_multilayer(kernel_size, num_layers)
+        hidden_dim = self._extend_for_multilayer(hidden_dim, num_layers)
+        if not len(kernel_size) == len(hidden_dim) == num_layers:
+            raise ValueError('Inconsistent list length.')
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.kernel_size = kernel_size
+        self.num_layers = num_layers
+        self.batch_first = batch_first
+        self.bias = bias
+        self.return_all_layers = return_all_layers
+
+        cell_list = []
+        for i in range(0, self.num_layers):
+            cur_input_dim = self.input_dim if i == 0 else self.hidden_dim[i - 1]
+
+            cell_list.append(ConvLSTMCell(input_dim=cur_input_dim,
+                                          hidden_dim=self.hidden_dim[i],
+                                          kernel_size=self.kernel_size[i],
+                                          bias=self.bias))
+
+        self.cell_list = nn.ModuleList(cell_list)
+
+    def forward(self, input_tensor, hidden_state=None):
+        """
+
+        Parameters
+        ----------
+        input_tensor: todo
+            5-D Tensor either of shape (t, b, c, h, w) or (b, t, c, h, w)
+        hidden_state: todo
+            None. todo implement stateful
+
+        Returns
+        -------
+        last_state_list, layer_output
+        """
+        if not self.batch_first:
+            # (t, b, c, h, w) -> (b, t, c, h, w)
+            input_tensor = input_tensor.permute(1, 0, 2, 3, 4)
+
+        b, _, _, h, w = input_tensor.size()
+
+        # Implement stateful ConvLSTM
+        if hidden_state is not None:
+            raise NotImplementedError()
+        else:
+            # Since the init is done in forward. Can send image size here
+            hidden_state = self._init_hidden(batch_size=b,
+                                             image_size=(h, w))
+
+        layer_output_list = []
+        last_state_list = []
+
+        seq_len = input_tensor.size(1)
+        cur_layer_input = input_tensor
+
+        for layer_idx in range(self.num_layers):
+
+            h, c = hidden_state[layer_idx]
+            output_inner = []
+            for t in range(seq_len):
+                h, c = self.cell_list[layer_idx](input_tensor=cur_layer_input[:, t, :, :, :],
+                                                 cur_state=[h, c])
+                output_inner.append(h)
+
+            layer_output = torch.stack(output_inner, dim=1)
+            cur_layer_input = layer_output
+
+            layer_output_list.append(layer_output)
+            last_state_list.append([h, c])
+
+        if not self.return_all_layers:
+            layer_output_list = layer_output_list[-1:]
+            last_state_list = last_state_list[-1:]
+
+        return layer_output_list, last_state_list
+
+    def _init_hidden(self, batch_size, image_size):
+        init_states = []
+        for i in range(self.num_layers):
+            init_states.append(self.cell_list[i].init_hidden(batch_size, image_size))
+        return init_states
+
+    @staticmethod
+    def _check_kernel_size_consistency(kernel_size):
+        if not (isinstance(kernel_size, tuple) or
+                (isinstance(kernel_size, list) and all([isinstance(elem, tuple) for elem in kernel_size]))):
+            raise ValueError('`kernel_size` must be tuple or list of tuples')
+
+    @staticmethod
+    def _extend_for_multilayer(param, num_layers):
+        if not isinstance(param, list):
+            param = [param] * num_layers
+        return param

GANFilling/src/models/discriminator.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.parametrizations import spectral_norm
+
+from convlstm import ConvLSTMCell
+
+class Discriminator(nn.Module): #PatchGAN
+    def __init__(self, device, inputChannels=6, d=64):
+        super().__init__()
+        self.device = device
+        self.conv1 = nn.Conv2d(inputChannels, d, 4, 2, 1)
+        self.conv2 = nn.Conv2d(d, d * 2, 4, 2, 1)
+        self.conv2_bn = nn.BatchNorm2d(d * 2, track_running_stats=False)
+        self.conv3 = nn.Conv2d(d * 2, d * 4, 4, 2, 1)
+        self.conv3_bn = nn.BatchNorm2d(d * 4, track_running_stats=False)
+        self.conv4_lstm = ConvLSTMCell(d * 4, d * 4, (3,3), False,self.device)
+        self.conv4 = nn.Conv2d(d * 4, d * 8, 4, 1, 1)
+        self.conv4_bn = nn.BatchNorm2d(d * 8, track_running_stats=False)
+        self.conv5_lstm = ConvLSTMCell(d * 8, d * 8, (3,3), False,self.device)
+        self.conv5 = nn.Conv2d(d * 8, 1, 4, 1, 1)
+
+        self.conv_fusion = nn.Conv2d(4, 1, 1)
+
+        torch.backends.cudnn.deterministic = True
+
+    def weight_init(self, mean, std):
+        for m in self._modules:
+            normal_init(self._modules[m], mean, std)
+    
+    def forward_step(self, input, states):
+        x = F.leaky_relu(self.conv1(input), 0.2)
+        x = F.leaky_relu(self.conv2_bn(self.conv2(x)), 0.2)
+        x = F.leaky_relu(self.conv3_bn(self.conv3(x)), 0.2)
+        states1 = self.conv4_lstm(x, states[0])
+        x = F.leaky_relu(self.conv4_bn(self.conv4(states1[0])), 0.2)
+        states2 = self.conv5_lstm(x, states[1])
+        x = F.leaky_relu(self.conv5(states2[0]), 0.2)
+        return x.squeeze(dim=1), [states1, states2]
+    
+    def forward(self, tensor):
+        output = torch.empty((tensor.shape[0], int(tensor.shape[2]/8)-2,int(tensor.shape[3]/8)-2,tensor.shape[4])).to(self.device)
+        for patch in range(tensor.shape[0]):
+            states = (None,None,None,None)
+            for timeStep in range(tensor.shape[4]):
+                output[patch,:,:,timeStep], states = self.forward_step(tensor[patch,:,:,:,timeStep].unsqueeze(dim=0), states)
+
+        return F.sigmoid(output), states
+
+
+def normal_init(m, mean, std):
+    if isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Conv2d):
+        m.weight.data.normal_(mean, std)
+        m.bias.data.zero_()
+
+if __name__=='__main__':
+    x = torch.zeros((16, 3, 32, 32, 4), dtype=torch.float32) 
+    model = Discriminator('cpu', inputChannels=3)
+    y = model(x)
+    print(y[0].size()) 

GANFilling/src/models/generator.py

@@ -0,0 +1,123 @@
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+sys.path.append('src/models')
+from convlstm import ConvLSTMCell
+
+class Generator(nn.Module):
+    def __init__(self, device, inputChannels = 4, outputChannels=3, d=64):
+        super().__init__()
+        self.d = d
+        self.device = device
+
+        self.conv1 = nn.Conv2d(inputChannels, d, 3, 2, 1)
+        self.conv2 = nn.Conv2d(d, d * 2, 3, 2, 1)
+        self.conv3 = nn.Conv2d(d * 2, d * 4, 3, 2, 1)
+        self.conv4 = nn.Conv2d(d * 4, d * 8, 3, 2, 1)
+        self.conv5 = nn.Conv2d(d * 8, d * 8, 3, 2, 1)
+        self.conv6 = nn.Conv2d(d * 8, d * 8, 3, 2, 1)
+        self.conv7 = nn.Conv2d(d * 8, d * 8, 3, 2, 1)
+
+        self.conv_lstm_d1 = ConvLSTMCell(d * 8, d * 8, (3,3), False, device)
+        self.conv_lstm_d2 = ConvLSTMCell(d * 8 * 2 , d * 8, (3,3), False, device)
+        self.conv_lstm_d3 = ConvLSTMCell(d * 8 * 2 , d * 8, (3,3), False, device)
+        self.conv_lstm_d4 = ConvLSTMCell(d * 8  * 2 , d * 4, (3,3), False, device)
+        self.conv_lstm_d5 = ConvLSTMCell(d * 4 * 2 , d * 2, (3,3), False, device)
+        self.conv_lstm_d6 = ConvLSTMCell(d * 2 * 2 , d, (3,3), False, device)
+        self.conv_lstm_d7 = ConvLSTMCell(d * 2 , d, (3,3), False, device)
+
+        self.conv_lstm_e1 = ConvLSTMCell(d, d, (3,3), False, device)
+        self.conv_lstm_e2 = ConvLSTMCell(d * 2 , d * 2, (3,3), False, device)
+        self.conv_lstm_e3 = ConvLSTMCell(d * 4 , d * 4, (3,3), False, device)
+        self.conv_lstm_e4 = ConvLSTMCell(d * 8 , d * 8, (3,3), False, device)
+        self.conv_lstm_e5 = ConvLSTMCell(d * 8 , d * 8, (3,3), False, device)
+        self.conv_lstm_e6 = ConvLSTMCell(d * 8 , d * 8, (3,3), False, device)
+        self.conv_lstm_e7 = ConvLSTMCell(d * 8 , d * 8, (3,3), False, device)
+
+        self.up = nn.Upsample(scale_factor=2)
+        self.conv_out = nn.Conv2d(d, outputChannels, 3, 1, 1)
+
+        self.slope = 0.2
+
+    def weight_init(self, mean, std):
+        for m in self._modules:
+            normal_init(self._modules[m], mean, std)
+
+    def forward_step(self, input, states_encoder, states_decoder):
+
+        e1 = self.conv1(input)
+        states_e1 = self.conv_lstm_e1(e1, states_encoder[0])
+        e2 = self.conv2(F.leaky_relu(states_e1[0], self.slope))
+        states_e2 = self.conv_lstm_e2(e2, states_encoder[1])
+        e3 = self.conv3(F.leaky_relu(states_e2[0], self.slope))
+        states_e3 = self.conv_lstm_e3(e3, states_encoder[2])
+        e4 = self.conv4(F.leaky_relu(states_e3[0], self.slope))
+        states_e4 = self.conv_lstm_e4(e4, states_encoder[3])
+        e5 = self.conv5(F.leaky_relu(states_e4[0], self.slope))
+        states_e5 = self.conv_lstm_e5(e5, states_encoder[4])
+        e6 = self.conv6(F.leaky_relu(states_e5[0], self.slope))
+        states_e6 = self.conv_lstm_e6(e6, states_encoder[5])
+        e7 = self.conv7(F.leaky_relu(states_e6[0], self.slope))
+        
+    
+        states1 = self.conv_lstm_d1(F.relu(e7), states_decoder[0]) 
+        d1 = self.up(states1[0])
+        d1 = torch.cat([d1, e6], 1)
+
+        states2 = self.conv_lstm_d2(F.relu(d1), states_decoder[1]) 
+        d2 = self.up(states2[0])
+        d2 = torch.cat([d2, e5], 1)
+        
+        states3 = self.conv_lstm_d3(F.relu(d2), states_decoder[2]) 
+        d3 = self.up(states3[0])  
+        d3 = torch.cat([d3, e4], 1)
+
+        states4 = self.conv_lstm_d4(F.relu(d3), states_decoder[3])
+        d4 = self.up(states4[0])
+        d4 = torch.cat([d4, e3], 1)
+
+        states5 = self.conv_lstm_d5(F.relu(d4), states_decoder[4])
+        d5 = self.up(states5[0])
+        d5 = torch.cat([d5, e2], 1)
+
+        states6 = self.conv_lstm_d6(F.relu(d5), states_decoder[5])
+        d6 = self.up(states6[0])
+        d6 = torch.cat([d6, e1], 1)
+
+        states7 = self.conv_lstm_d7(F.relu(d6), states_decoder[6])
+        d7 = self.up(states7[0])
+
+        o = torch.clip(torch.tanh(self.conv_out(d7)), min=-0.0, max = 1)
+
+        states_e = [states_e1, states_e2, states_e3,states_e4, states_e5, states_e6]
+        states_d = [states1, states2, states3,states4, states5, states6, states7]
+
+        return o, (states_e, states_d)
+    
+    def forward(self, tensor):
+        states_encoder = (None,None,None,None,None,None,None)
+        states_decoder = (None,None,None,None,None,None,None)
+        output = torch.empty_like(tensor)
+        for timeStep in range(tensor.shape[4]):
+            output[:,:,:,:,timeStep], states = self.forward_step(tensor[:,:,:,:,timeStep], states_encoder, states_decoder)
+            states_encoder, states_decoder = states[0], states[1]
+        return output, states
+
+
+def normal_init(m, mean, std):
+    if isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Conv2d):
+        m.weight.data.normal_(mean, std)
+        m.bias.data.zero_()
+
+if __name__=='__main__':
+    # batch_size = number of 3D patches
+    # num_channles = BGR+NIR
+    # h,w = spatial resolution 
+    states_encoder = (None,None,None,None,None,None,None)
+    states_decoder = (None,None,None,None,None,None,None)
+    x = torch.zeros((2, 4, 128, 128, 10), dtype=torch.float32) 
+    model = Generator('cpu', inputChannels=4)
+    y, states = model(x)
+    print(y.size()) 

GANFilling/src/test.py

@@ -0,0 +1,138 @@
+import os
+import sys
+import torch 
+import csv
+import numpy as np
+
+import argparse
+from pathlib import Path
+from PIL import Image
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import matplotlib.patches as mpatches
+
+sys.path.append('src/models')
+from generator import *
+
+
+def main(config):
+
+    gap_filling_model_path = config.model_path
+    results_path = Path(config.results_path)
+    data_path = Path(config.data_path)
+    list_samples = [x for x in os.listdir(config.data_path) if x.endswith('.npz')]
+
+
+    landcover_cmap = {}
+    with open('data/landcover_types.csv', 'r') as fd:
+        reader = csv.reader(fd)
+        for row in reader:
+            landcover_cmap[row[0]] = [row[1], row[2], row[3]]
+    
+    for sample in list_samples:
+
+        sample_lc = data_path / "landcover" / sample.replace("context_", "")
+
+        landcover = np.load(sample_lc)['landcover'][0]
+        context = np.load(data_path / sample)['highresdynamic'][:,:,:4, :10]
+
+        model = __load_gapFill_model(gap_filling_model_path)
+        gap_filled_sample = __gapFill_context(context, model)
+
+        
+        n = gap_filled_sample[0,-1]
+        r = gap_filled_sample[0,2]
+        ndvi = (n-r)/(n+r)
+
+        fig, axs = plt.subplots(3, 11, figsize=(15, 4))
+
+        landcover_rgb = np.empty((128,128, 3))   
+        for c in np.unique(landcover):
+            for ch in range(3):
+                landcover_rgb[landcover==c, ch] = landcover_cmap[str(c)][ch]
+
+        landcover_rgb = Image.fromarray(landcover_rgb.astype('uint8'), 'RGB')
+        axs[0, 0].imshow(landcover_rgb)
+        axs[0, 0].axis('off')
+        axs[1, 0].axis('off')
+        axs[2, 0].axis('off')
+
+        import datetime
+        date = datetime.datetime.strptime(sample.split('_')[2], '%Y-%m-%d')
+
+
+        for i in range(10):
+            img_rgb = np.clip(context[:,:,:3, i]*3,0,1)[:,:,[2,1,0]]            
+            img_rgb = (img_rgb * 255).astype('uint8')
+            img_rgb = Image.fromarray(img_rgb)
+            axs[0, i+1].imshow(img_rgb) 
+            axs[0, i+1].text(94, 122, (date+datetime.timedelta(days=5*i)).strftime('%j'), color='black', fontsize=10, bbox=dict(boxstyle='square,pad=0.1', fc='white', ec='none'))
+            axs[0, i+1].axis('off')
+        
+        for i in range(10):
+            img_rgb = np.transpose(gap_filled_sample.cpu().detach().numpy()[0, [2,1,0],:,:, i], (1, 2, 0))
+            img_rgb = np.clip(img_rgb*3, 0, 1)
+            img_rgb = (img_rgb * 255).astype('uint8')
+            img_rgb = Image.fromarray(img_rgb)
+            axs[1, i+1].imshow(img_rgb) 
+            axs[1, i+1].axis('off')
+        
+        cmap = mpl.colormaps.get_cmap('jet')
+        cmap.set_bad(color='black')
+        
+        for i in range(10):
+            ndvi_img = Image.fromarray(ndvi[:,:,i].cpu().detach().numpy())
+            im = axs[2, i+1].imshow(ndvi[:,:,i].cpu().detach().numpy(), cmap=cmap, vmin=0., vmax=1.) 
+            axs[2, i+1].axis('off')
+        
+        cbar_ax = fig.add_axes([0.2, 0.95, 0.7, 0.03])
+        cbar = fig.colorbar(im, cax=cbar_ax, orientation='horizontal')
+        cbar.outline.set_linewidth(0)
+
+        plt.subplots_adjust(wspace=0., hspace=0.05)
+
+        landcover_cmap_legend = {}
+        with open('data/landcover_types.csv', 'r') as fd:
+            reader = csv.reader(fd)
+            for row in reader:
+                landcover_cmap_legend[row[5]] = [int(row[1])/255, int(row[2])/255, int(row[3])/255, 1]
+        
+        landcover_cmap_legend.pop('No data')
+        legend_items =[mpatches.Patch(color=color,label=lc) for lc, color in landcover_cmap_legend.items()]
+
+        path = results_path / sample
+        plt.savefig(f"{path}.png", bbox_inches='tight')
+        plt.close()
+
+
+## GAP FILLING FUNCTIONS
+
+def __load_gapFill_model(gapFilling_model):
+    modelConfiguration = torch.load(gapFilling_model, map_location='cpu')
+    trainConfiguration = modelConfiguration['configuration']
+    model = Generator('cpu' , trainConfiguration.inputChannels, trainConfiguration.outputChannels)
+    model.load_state_dict(modelConfiguration['model'])
+    model.eval()
+    return model
+
+def __prepare_data(context):
+    context = np.nan_to_num(np.clip(context, 0, 1), nan=1.0)
+    context = np.transpose(context, (2,0,1,3))
+    return context
+
+def __gapFill_context(context, model):
+    context = __prepare_data(context)
+    tensor = torch.from_numpy(context).unsqueeze(dim=0).float()
+    output, _ = model(tensor)
+    return output
+
+if __name__=='__main__':
+    parser = argparse.ArgumentParser(description = "ArgParse")
+    parser.add_argument('--model_path', type=str) 
+    parser.add_argument('--data_path', type=str) 
+    parser.add_argument('--results_path', type=str) 
+    
+    config = parser.parse_args()
+
+    main(config)

GANFilling/src/train.py

@@ -0,0 +1,258 @@
+import os
+import argparse
+import torch
+from torch.utils.data import DataLoader 
+import logging
+import numpy as np
+import random
+import torch.nn as nn
+from pathlib import Path
+
+from torch.autograd import Variable
+from torchvision import transforms
+from skimage.metrics import structural_similarity as ssim
+from skimage.metrics import peak_signal_noise_ratio as psnr
+
+from iterator import Iterator
+from models.generator import Generator
+from models.discriminator import Discriminator
+
+class Trainer():
+    def __init__(self, configuration):
+        self.configuration = configuration
+        self.__createDirectories()
+        self.__initCriterions()
+        self.__setLogger()
+        self.__logParser()
+        self.__resetAccumulatedMetrics()
+
+        self.generatorHistLoss = []
+        self.generatorHistLoss_adv = []
+        self.generatorHistLoss_rec = []
+        self.discriminatorHistLoss = []
+
+        self.metricSSIM = []
+        self.metricPSNR = []
+
+        torch.manual_seed(100)
+        # Seed python RNG
+        random.seed(100)
+        # Seed numpy RNG
+        np.random.seed(100)
+        
+    def __logParser(self):
+        for arg, value in vars(self.configuration).items():
+            logging.info("Argument %s: %r", arg, value)
+    
+    def __createDirectories(self):
+        self.modelPath = Path(self.configuration.modelsPath + self.configuration.name )
+        Path.mkdir(self.modelPath / 'Images/', parents=True, exist_ok=True)
+
+    def __initCriterions(self):
+        self.discriminatorCriterion = nn.BCELoss(reduction='none')
+        self.generatorCriterion = [nn.BCELoss(reduction='none'), nn.L1Loss(reduction='none')]
+    
+    def train(self):
+        logging.info('Starting training')
+        self.step = 0
+        self.mode = 'train'
+        for self.epoch in range(self.configuration.maxEpochs): 
+            for noisyTensor, cleanTensor, masksTensor in self.trainDataLoader: 
+
+                noisyTensor = noisyTensor.float().to(self.device)
+                cleanTensor = cleanTensor.float().to(self.device)
+                masksTensor = masksTensor.float().to(self.device)
+
+                ### Update Discriminator
+                self.__setRequiresGrad(self.discriminator, True)
+                self.discriminator.zero_grad()
+
+                outputDiscriminator,_ = self.discriminator(cleanTensor)
+                realLossDiscriminator = self.discriminatorCriterion(outputDiscriminator, Variable(torch.ones(outputDiscriminator.size()).to(self.device)))
+                
+                outputGenerator, _ = self.generator(noisyTensor)
+                outputGenerator_patch, _ = self.__divideInPatches(outputGenerator,  cleanTensor.shape[4])
+                outputDiscriminator,_ = self.discriminator(outputGenerator_patch)
+                fakeLossDiscriminator = self.discriminatorCriterion(outputDiscriminator, Variable(torch.zeros(outputDiscriminator.size()).to(self.device)))
+
+                trainLossDiscriminator = (torch.mean(realLossDiscriminator) + torch.mean(fakeLossDiscriminator))*0.5
+                trainLossDiscriminator.backward()
+                self.discriminatorOptimizer.step()
+                self.discriminator.zero_grad()
+                self.discriminatorHistLoss.append(trainLossDiscriminator)
+
+                ### Update Generator
+                self.__setRequiresGrad(self.discriminator, False)
+                self.generator.zero_grad()
+                
+                outputGenerator, _ = self.generator(noisyTensor)
+                outputGenerator_patch, startTime = self.__divideInPatches(outputGenerator,  4)
+                outputDiscriminator,_ = self.discriminator(outputGenerator_patch)
+
+                generatorAdversarialLoss = torch.mean(self.generatorCriterion[0](outputDiscriminator, Variable(torch.ones(outputDiscriminator.size()).to(self.device))))
+                generatorReconstructionLoss = self.generatorCriterion[1](outputGenerator, noisyTensor) * (1-masksTensor).unsqueeze(dim=1)
+                generatorReconstructionLoss = self.configuration.lambdaL1 * torch.mean(generatorReconstructionLoss)
+
+                trainLossGenerator =  torch.mean(generatorAdversarialLoss) + generatorReconstructionLoss
+                trainLossGenerator.backward()
+                self.generatorOptimizer.step()
+                self.generator.zero_grad()
+                self.generatorHistLoss.append(trainLossGenerator)
+                self.generatorHistLoss_adv.append(torch.mean(generatorAdversarialLoss))
+                self.generatorHistLoss_rec.append(torch.mean(generatorReconstructionLoss))
+                self.__evaluate(outputGenerator, noisyTensor, masksTensor) 
+                self.__logMetrics()
+                self.__plotGeneratorOutput(outputGenerator, noisyTensor, self.modelPath)
+
+                if self.step % int(self.configuration.validateEvery) == 0:
+                    self.__saveModelFreq()
+                
+    def __saveModelFreq(self):
+        self.__saveModel(self.generator, 'step_generator_'+str(self.step))
+        logging.info('Saved model')
+    
+    def __divideInPatches(self, tensor, timeSteps, kernelSize = 64):
+        new_tensor = torch.zeros((1,tensor.shape[1], kernelSize, kernelSize, timeSteps)).float().to(self.device)
+        startTime = random.sample([x+1 for x in range(-1,tensor.shape[4]-timeSteps)],1)[0]
+        x = random.sample([x+1 for x in range(-1,64)],1)[0]
+        y = random.sample([x+1 for x in range(-1,64)],1)[0]
+
+        new_tensor[0] = tensor[0, :, x:x+kernelSize, y:y+kernelSize, startTime:startTime+timeSteps]
+
+        return new_tensor, startTime
+
+    def __unnormalize(self, tensor):
+        if self.configuration.normalize:
+            return tensor/2 + .5
+        else:
+            return tensor
+
+    def __setRequiresGrad(self, net, requires_grad=False):
+        for param in net.parameters():
+            param.requires_grad = requires_grad
+            
+    def __evaluate(self, outputTensor, noisyTensor, masksTensor):
+        outputImages = outputTensor.detach().permute(0,2,3,1,4).cpu().numpy()
+        noisyImages = noisyTensor.detach().permute(0,2,3,1,4).cpu().numpy()
+        maskImages = masksTensor.detach().cpu().numpy()
+
+        for imgBatch in range(outputTensor.shape[0]):
+            maskImg = np.expand_dims( (1-maskImages[imgBatch]), 2) 
+            outputImg = outputImages[imgBatch]*maskImg
+            noisyImg = noisyImages[imgBatch]*maskImg
+
+            ssim_aux = []
+            psnr_aux = []
+            for timeStep in range(noisyImg.shape[3]):
+                metric_1 = ssim(outputImg[:,:,:,timeStep], noisyImg[:,:,:,timeStep], multichannel=True)
+                metric_2 = psnr(outputImg[:,:,:,timeStep], noisyImg[:,:,:,timeStep])
+                if np.isfinite(metric_1):
+                    ssim_aux.append(metric_1)
+                if np.isfinite(metric_2):
+                    psnr_aux.append(metric_2)
+
+            self.metricSSIM.append(sum(ssim_aux)/len(ssim_aux))
+            self.metricPSNR.append(sum(psnr_aux)/len(psnr_aux))
+        self.metricSSIM = sum(self.metricSSIM)/self.configuration.batchSize
+        self.metricPSNR = sum(self.metricPSNR)/self.configuration.batchSize
+        self.accumulatedMetricSSIM.append(self.metricSSIM)
+        self.accumulatedMetricPSNR.append(self.metricPSNR)
+    
+    def __plotGeneratorOutput(self, outputGenerator, noisyImg, path):
+        if self.step % int(self.configuration.plotEvery) == 0:
+            for sampleNumber in range(outputGenerator.shape[0]):
+                for timeStep in range(outputGenerator.shape[4]):
+                    self.__saveImage(self.__unnormalize(outputGenerator[sampleNumber,:,:,:,timeStep])[[2,1,0],:,:], sampleNumber, path, '_'+str(timeStep))
+                    self.__saveImage(self.__unnormalize(noisyImg[sampleNumber,:,:,:,timeStep])[[2,1,0],:,:], sampleNumber, path, 'noisy_'+str(timeStep))
+    
+    def __saveImage(self, tensor, sampleNumber, path, label):
+        tensor = torch.clamp(3*tensor, -1, 1)
+        img = transforms.ToPILImage()(tensor).convert("RGB")
+        img.save(str(path /'Images') + '/e'+str(self.epoch)+'_s'+str(self.step)+'_img'+str(sampleNumber)+'_'+label+'.png')
+
+    def __logMetrics(self):
+        if self.step % int(self.configuration.printEvery) == 0:
+            logging.info('Epoch {epoch} - Step {step} ---> GLoss: {gloss} - adv: {gloss_adv} - rec: {gloss_rec}, DLoss: {dloss}, SSIM: {ssim} PSNR: {psnr}'.format(\
+                            epoch=self.epoch, step=self.step, gloss=torch.mean(torch.FloatTensor(self.generatorHistLoss[-self.configuration.printEvery:])), gloss_adv=torch.mean(torch.FloatTensor(self.generatorHistLoss_adv[-self.configuration.printEvery:])), gloss_rec=torch.mean(torch.FloatTensor(self.generatorHistLoss_rec[-self.configuration.printEvery:])),\
+                            dloss= torch.mean(torch.FloatTensor(self.discriminatorHistLoss[-self.configuration.printEvery:])),\
+                            ssim=torch.mean(torch.FloatTensor(self.accumulatedMetricSSIM)), psnr=torch.mean(torch.FloatTensor(self.accumulatedMetricPSNR))))
+            self.__resetAccumulatedMetrics()
+        self.step += 1
+        self.__resetMetrics()
+    
+    def __resetMetrics(self):
+        self.metricSSIM = []
+        self.metricPSNR = []
+    
+    def __resetAccumulatedMetrics(self):
+        self.accumulatedMetricSSIM = []
+        self.accumulatedMetricPSNR = []
+
+    def __saveModel(self, model, tag):
+            modelDictionary = {'model':model.state_dict(), 'configuration':self.configuration}
+            torch.save(modelDictionary, str(self.modelPath / tag)+'.pt')    
+
+    def setData(self): 
+        self.trainIterator = Iterator(Path(configuration.dataPath), configuration.cleanDataPath, 'train')
+        self.trainDataLoader = DataLoader(self.trainIterator, batch_size=self.configuration.batchSize, shuffle=True)
+
+        self.validIterator = Iterator(Path(configuration.dataPath), configuration.cleanDataPath,'valid')
+        self.validDataLoader = DataLoader(self.validIterator, batch_size=1)
+
+    def setModels(self):
+        self.__setDevice()
+        self.__setGenerator()
+        self.__setDiscriminator()
+    
+    def __setGenerator(self):
+        self.generator = Generator(self.device, self.configuration.inputChannels, self.configuration.outputChannels)
+        self.generator.weight_init(mean=0.0, std=0.02)
+        self.generator.to(self.device)
+        self.generatorOptimizer = torch.optim.Adam(self.generator.parameters(),lr=self.configuration.lrG, betas=(0.5,0.999))
+    
+    def __setDiscriminator(self):
+        self.discriminator = Discriminator(self.device, self.configuration.inputChannels)
+        self.discriminator.weight_init(mean=0.0, std=0.02)
+        self.discriminator.to(self.device)
+        self.discriminatorOptimizer = torch.optim.Adam(self.discriminator.parameters(),lr=self.configuration.lrD, betas=(0.5,0.999))
+
+    def __setDevice(self):
+        print('TORCH AVAILABLE: ', torch.cuda.is_available())
+        if torch.cuda.is_available():
+            self.device = 'cuda:0'
+        else:
+            self.device = 'cpu'
+
+    def __setLogger(self):
+        Path.mkdir(Path(self.configuration.logsPath), exist_ok=True)
+        logging.basicConfig(filename=self.configuration.logsPath+self.configuration.name+'.log', format='%(asctime)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S', level=logging.INFO, filemode='w')
+
+def main(configuration):
+    trainer = Trainer(configuration)
+    trainer.setData()
+    trainer.setModels()
+    trainer.train()
+
+if __name__=='__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dataPath', type=str, default="train_data/", help='Path to the .npz files')
+    parser.add_argument('--cleanDataPath', type=str, default="utils/example_clean_data.json", help='file with the noiseless data paths for the discriminator')
+    parser.add_argument('--logsPath', type=str, default='logs/')
+    parser.add_argument('--modelsPath', type=str, default='trained_models/', help='Path to save the trained models')
+    parser.add_argument('--name', type=str, help='Model name')
+
+    parser.add_argument('--maxEpochs', type=int, default=1)
+    parser.add_argument('--batchSize', type=int, default=1)
+    parser.add_argument('--printEvery',type=int, default=1)
+    parser.add_argument('--plotEvery', type=int, default=5)
+    parser.add_argument('--validateEvery', type=int, default=5)
+
+    parser.add_argument('--inputChannels', type=int, default=4)
+    parser.add_argument('--outputChannels', type=int, default=4)
+    parser.add_argument('--lrD', type=float, default=0.0002, help='learning rate, default=0.0002')
+    parser.add_argument('--lrG', type=float, default=0.0002, help='learning rate, default=0.0002')
+    parser.add_argument('--lambdaL1', type=float, default=100, help='lambda for L1 loss')
+    parser.add_argument('--normalize', action='store_true')
+
+    configuration = parser.parse_args()
+    main(configuration)

GANFilling/src/utils/example_clean_data.json

@@ -0,0 +1,83 @@
+{
+    "33U": [
+        {
+            "path": "data/train_data/33UXQ/33UXQ_2018-06-18_2018-11-14_2105_2233_3769_3897_32_112_58_138.npz",
+            "kernel size": 64,
+            "bbox": [
+                11,
+                60
+            ],
+            "time steps": [
+                7,
+                8,
+                9
+            ]
+        },
+        {
+            "path": "data/train_data/33UXQ/33UXQ_2018-07-28_2018-12-24_4025_4153_2105_2233_62_142_32_112.npz",
+            "kernel size": 64,
+            "bbox": [
+                0,
+                44
+            ],
+            "time steps": [
+                6,
+                7,
+                8,
+                9
+            ]
+        },
+        {
+            "path": "data/train_data/33UXQ/33UXQ_2018-07-08_2018-12-04_2617_2745_3513_3641_40_120_54_134.npz",
+            "kernel size": 64,
+            "bbox": [
+                0,
+                29
+            ],
+            "time steps": [
+                3,
+                4,
+                5
+            ]
+        },
+        {
+            "path": "data/train_data/33UXQ/33UXQ_2018-07-18_2018-12-14_825_953_4537_4665_12_92_70_150.npz",
+            "kernel size": 64,
+            "bbox": [
+                0,
+                35
+            ],
+            "time steps": [
+                1,
+                2,
+                3
+            ]
+        },
+        {
+            "path": "data/train_data/33UXQ/33UXQ_2018-06-18_2018-11-14_2745_2873_3897_4025_42_122_60_140.npz",
+            "kernel size": 64,
+            "bbox": [
+                0,
+                0
+            ],
+            "time steps": [
+                7,
+                8,
+                9
+            ]
+        },
+        {
+            "path": "data/train_data/33UXQ/33UXQ_2018-07-08_2018-12-04_697_825_4793_4921_10_90_74_154.npz",
+            "kernel size": 64,
+            "bbox": [
+                0,
+                49
+            ],
+            "time steps": [
+                3,
+                4,
+                5
+            ]
+        }
+    ]
+}

GANFilling/src/utils/generate_cleanData_file.py

@@ -0,0 +1,84 @@
+import os 
+import sys
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from itertools import groupby
+from operator import itemgetter
+
+import matplotlib.pyplot as plt
+import json
+
+folderName = '/train_data/context/'
+split = 'IID'
+data = {split:[]}
+maxSizeKernel = 128
+minSizeKernel = 127
+
+i = 0
+for fileName in os.listdir(folderName):
+    if fileName == 'LICENSE':
+        pass
+    else:
+        print(fileName)
+        for dataCube in os.listdir(folderName+'/'+fileName):
+            cubePath = folderName + '/' + fileName + '/' +dataCube
+            print('Evaluating data cube ', cubePath)
+            sample = np.load(cubePath)
+            masks = sample['highresdynamic'][:,:,-1,:10]
+
+            clean = {}
+            total_kernels = {}
+            max_lenght_kernels = {}
+            for kernelSize in range(maxSizeKernel,minSizeKernel, -1):
+                clean[kernelSize] = {}
+                total_kernels[kernelSize] = {}
+                max_lenght_kernels[kernelSize] = {}
+            
+            ########################################################################
+            masks = np.expand_dims(masks, axis=(0, 1))
+            validKernel = False
+
+            for kernelSize in range(maxSizeKernel,minSizeKernel, -1):
+                kernel = torch.ones((1,1,kernelSize, kernelSize))
+                for timeStep in range(masks.shape[4]):
+                    conv = F.conv2d(torch.from_numpy(masks[:,:,:,:,timeStep]).float(), kernel)
+                    clean[kernelSize][timeStep]= (conv == 0).nonzero()
+
+                    if len(clean[kernelSize][timeStep]) != 0:
+                        for clean_kernel in clean[kernelSize][timeStep]:
+                            if (clean_kernel[2].item(),clean_kernel[3].item()) not in total_kernels[kernelSize].keys():
+                                total_kernels[kernelSize][(clean_kernel[2].item(),clean_kernel[3].item())] = [timeStep]
+                            else:
+                                total_kernels[kernelSize][(clean_kernel[2].item(),clean_kernel[3].item())].append(timeStep)
+                #print('KS: ', kernelSize)
+                #print(total_kernels[kernelSize])
+                lenght = 3
+                kernel_to_save = []
+                for key, value in  total_kernels[kernelSize].items():
+                    for k, g in groupby(enumerate(value), lambda ix : ix[0] - ix[1]):
+                        consecutives = list(map(itemgetter(1), g))
+                        #print(consecutives)
+                        #print(len(consecutives) >= lenght)
+                        if len(consecutives) >= lenght:
+                            kernel_to_save.append(key)
+                            kernel_to_save.append(consecutives)
+                            lenght = len(consecutives)  
+                            validKernel = True
+                if validKernel:
+                    print('Found valid kernel ', kernelSize, ' with ', len(kernel_to_save[1]), ' consecutive time steps')
+                    data[split].append({
+                        'path':str(cubePath),
+                        'kernel size': kernelSize,
+                        'bbox': kernel_to_save[0],
+                        'time steps': kernel_to_save[1]
+                    })
+                    validKernel = False
+                    break
+
+print('Saving data ...')
+with open('cleanData.json', 'w') as outfile:
+    json.dump(data, outfile)
+
+