Filter out NaNs when loading datasets

datasets.py

@@ -188,10 +188,18 @@ def get_dataset(dataset_name, target_folder="./", datasets=DATASETS_CONFIG):
         # Custom dataset
         img, gt, rgb_bands, ignored_labels, label_values, palette = CUSTOM_DATASETS_CONFIG[dataset_name]['loader'](folder)
 
+    # Filter NaN out
+    nan_mask = np.isnan(img.sum(axis=-1))
+    if np.count_nonzero(nan_mask) > 0:
+       print("Warning: NaN have been found in the data. It is preferable to remove them beforehand. Learning on NaN data is disabled.")
+    img[nan_mask] = 0
+    gt[nan_mask] = 0
+    ignored_labels.append(0)
+
+    ignored_labels = list(set(ignored_labels))
     # Normalization
     img = np.asarray(img, dtype='float32')
     img = (img - np.min(img)) / (np.max(img) - np.min(img))
-
     return img, gt, label_values, ignored_labels, rgb_bands, palette