biapy 3.4.2__tar.gz → 3.4.3__tar.gz

{biapy-3.4.2 → biapy-3.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: biapy
-Version: 3.4.2
+Version: 3.4.3
 Summary: BiaPy: Bioimage analysis pipelines in Python
 Author-email: Daniel Franco-Barranco <daniel.franco@dipc.org>
 Maintainer-email: Daniel Franco-Barranco <daniel.franco@dipc.org>

{biapy-3.4.2 → biapy-3.4.3}/biapy/_biapy.py

@@ -11,7 +11,7 @@ import numpy as np
 import importlib
 import multiprocessing
 
-from biapy.utils.misc import init_devices, is_dist_avail_and_initialized, set_seed, get_rank
+from biapy.utils.misc import init_devices, is_dist_avail_and_initialized, set_seed, get_rank, setup_for_distributed
 from biapy.config.config import Config
 from biapy.engine.check_configuration import check_configuration
 
@@ -420,6 +420,7 @@ class BiaPy():
 
         if is_dist_avail_and_initialized():
             print(f"[Rank {get_rank()} ({os.getpid()})] Process waiting (test finished) . . . ")
+            setup_for_distributed(self.args.rank == 0)
             dist.barrier()
             dist.destroy_process_group()
 

{biapy-3.4.2 → biapy-3.4.3}/biapy/config/config.py

@@ -184,24 +184,26 @@ class Config:
         _C.DATA.W_BACKGROUND = 0.06 # Used when _C.DATA.PROBABILITY_MAP=True
 
         # Whether to reshape the dimensions that does not satisfy the patch shape selected by padding it with reflect. It's not
-        # implemented in super-resolution workflow.
+        # implemented in super-resolution inference phase workflow (as usually the patch size is small).
         _C.DATA.REFLECT_TO_COMPLETE_SHAPE = False
 
         _C.DATA.NORMALIZATION = CN()
+        # Whether to apply or not a percentile clipping before normalizing the data
+        _C.DATA.NORMALIZATION.PERC_CLIP = False
+        # Lower and upper bound for percentile clip. Must be set when DATA.NORMALIZATION.PERC_CLIP = 'True'
+        _C.DATA.NORMALIZATION.PERC_LOWER = -1.0
+        _C.DATA.NORMALIZATION.PERC_UPPER = -1.0
         # Normalization type to use. Possible options:
         #   'div' to divide values from 0/255 (or 0/65535 if uint16) in [0,1] range
         #   'custom' to use DATA.NORMALIZATION.CUSTOM_MEAN and DATA.NORMALIZATION.CUSTOM_STD to normalize
         _C.DATA.NORMALIZATION.TYPE = 'div'
-        # Whether to apply the normalization by sample ("image") or by all dataset statistics ("dataset"). Options: ["image", "dataset"]  
+        # Whether to apply the normalization by sample ("image") or by all dataset statistics ("dataset"). 
+        # Used with 'DATA.NORMALIZATION.PERC_CLIP' == 'True' and/or when 'DATA.NORMALIZATION.TYPE' == 'custom'. 
+        # Options: ["image", "dataset"]  
         _C.DATA.NORMALIZATION.APPLICATION_MODE = "image"
         # Custom normalization variables: mean and std (they are calculated if not provided)
         _C.DATA.NORMALIZATION.CUSTOM_MEAN = -1.0
         _C.DATA.NORMALIZATION.CUSTOM_STD = -1.0
-        # Whether to apply or not a percentile clipping before normalizing the data
-        _C.DATA.NORMALIZATION.PERC_CLIP = False
-        # Lower and upper bound for percentile clip. Must be set when DATA.NORMALIZATION.PERC_CLIP = 'True'
-        _C.DATA.NORMALIZATION.PERC_LOWER = -1.0
-        _C.DATA.NORMALIZATION.PERC_UPPER = -1.0
 
         # If 'DATA.PATCH_SIZE' selected has 3 channels, e.g. RGB images are expected, so will force grayscale images to be
         # converted into RGB (e.g. in ImageNet some of the images are grayscale)
@@ -384,9 +386,6 @@ class Config:
         # Padding to be done in (y,x)/(z,y,x) when cropping validation data. Useful to avoid patch 'border effect'. This 
         # is only used when the validation is loaded from disk, and thus, not extracted from training. 
         _C.DATA.VAL.PADDING = (0,0)
-        # Directory where validation binary masks should be located. This binary mask will be applied only when MW_TH*
-        # optimized values are find, that is, when _C.PROBLEM.INSTANCE_SEG.DATA_MW_OPTIMIZE_THS = True and _C.TEST.POST_PROCESSING.APPLY_MASK = True
-        _C.DATA.VAL.BINARY_MASKS = os.path.join("user_data", 'val', 'bin_mask')
         # Not used yet.
         _C.DATA.VAL.RESOLUTION = (-1,)
         # Order of the axes of the image when using Zarr/H5 images to train
@@ -435,7 +434,7 @@ class Config:
         # How to fill up the new values created with affine transformations (rotations, shear, shift and zoom).
         # Same meaning as in scipy: 'reflect', grid-'mirror', 'constant', 'grid-constant', 'nearest', 'mirror', 
         # 'grid-wrap' and 'wrap'. 
-        _C.AUGMENTOR.AFFINE_MODE = 'constant'
+        _C.AUGMENTOR.AFFINE_MODE = 'reflect'
         # Make vertical flips
         _C.AUGMENTOR.VFLIP = False
         # Make horizontal flips
@@ -843,7 +842,7 @@ class Config:
         _C.TEST.REUSE_PREDICTIONS = False
 
         # If PROBLEM.NDIM = '2D' this can be activated to process each image entirely instead of patch by patch. Only can be done 
-        # if the neural network is fully convolutional
+        # if the neural network is fully convolutional. It's not implemented in super-resolution workflow. 
         _C.TEST.FULL_IMG = False 
 
         ### Instance segmentation
@@ -891,9 +890,12 @@ class Config:
         _C.TEST.POST_PROCESSING.YZ_FILTERING_SIZE = 5
         _C.TEST.POST_PROCESSING.Z_FILTERING = False
         _C.TEST.POST_PROCESSING.Z_FILTERING_SIZE = 5
-        # Apply a binary mask to remove possible segmentation outside it
+        # Apply a binary mask to remove possible segmentation outside it (you need to provide the mask and it must 
+        # contain two values: '1' -> preserve the pixel ; '0' discard pixel ). A mask for each test sample must be 
+        # provided and it will be loaded using 'DATA.TEST.BINARY_MASKS' variable.
         _C.TEST.POST_PROCESSING.APPLY_MASK = False
 
+        ### Instance segmentation
         # Whether to measure morphological features on each instances, i.e. 'circularity' (2D), 'elongation' (2D), 'npixels', 'area', 'diameter', 
         # 'perimeter', 'sphericity' (3D)
         _C.TEST.POST_PROCESSING.MEASURE_PROPERTIES = CN() 
@@ -964,7 +966,6 @@ class Config:
         # "greather equal", e.g. ">=", "less than", e.g. "<", and "less equal" e.g. "<=" comparisons.
         _C.TEST.POST_PROCESSING.MEASURE_PROPERTIES.REMOVE_BY_PROPERTIES.SIGN = []
 
-        ### Instance segmentation
         # Whether to apply Voronoi using 'BC' or 'M' channels need to be present
         _C.TEST.POST_PROCESSING.VORONOI_ON_MASK = False
         # Threshold to be applied to the 'M' channel when expanding the instances with Voronoi. Need to be in [0,1] range.
@@ -986,7 +987,8 @@ class Config:
         _C.TEST.POST_PROCESSING.REMOVE_CLOSE_POINTS_RADIUS = [-1.0]
         # Whether to apply a watershed to grow the points detected 
         _C.TEST.POST_PROCESSING.DET_WATERSHED = False
-        # Structure per each class to dilate the initial seeds before watershed
+        # Structure per each class to dilate the initial seeds before watershed. For instance, with two classes in a 3D problem:
+        # [ [2,2,1], [10,10,4] ]
         _C.TEST.POST_PROCESSING.DET_WATERSHED_FIRST_DILATION = [[-1,-1],]
         # List of classes to be consider as 'donuts'. For those class points, the 'donuts' type cell means that their nucleus is 
         # to big and that the seeds need to be dilated more so the watershed can grow the instances properly.
@@ -1014,6 +1016,7 @@ class Config:
         _C.PATHS.RESULT_DIR.FULL_IMAGE_INSTANCES = os.path.join(_C.PATHS.RESULT_DIR.PATH, 'full_image_instances')
         _C.PATHS.RESULT_DIR.FULL_IMAGE_POST_PROCESSING = os.path.join(_C.PATHS.RESULT_DIR.PATH, 'full_image_post_processing')
         _C.PATHS.RESULT_DIR.AS_3D_STACK = os.path.join(_C.PATHS.RESULT_DIR.PATH, 'as_3d_stack')
+        _C.PATHS.RESULT_DIR.AS_3D_STACK_BIN = os.path.join(_C.PATHS.RESULT_DIR.PATH, 'as_3d_stack_binarized')
         _C.PATHS.RESULT_DIR.AS_3D_STACK_POST_PROCESSING = os.path.join(_C.PATHS.RESULT_DIR.PATH, 'as_3d_stack_post_processing')
         _C.PATHS.RESULT_DIR.DET_LOCAL_MAX_COORDS_CHECK = os.path.join(_C.PATHS.RESULT_DIR.PATH, 'per_image_local_max_check')
         _C.PATHS.RESULT_DIR.DET_ASSOC_POINTS = os.path.join(_C.PATHS.RESULT_DIR.PATH, 'point_associations')
@@ -1023,7 +1026,6 @@ class Config:
         _C.PATHS.PROFILER = os.path.join(_C.PATHS.RESULT_DIR.PATH, 'profiler')
 
         # Name of the folder where the charts of the loss and metrics values while training the network are stored.
-        # Additionally, MW_TH* variable charts are stored if _C.PROBLEM.INSTANCE_SEG.DATA_MW_OPTIMIZE_THS = True
         _C.PATHS.CHARTS = os.path.join(_C.PATHS.RESULT_DIR.PATH, 'charts')
         # Folder where samples of DA will be stored
         _C.PATHS.DA_SAMPLES = os.path.join(_C.PATHS.RESULT_DIR.PATH, 'aug')
@@ -1091,8 +1093,6 @@ class Config:
         self._C.DATA.VAL.INSTANCE_CHANNELS_DIR = self._C.DATA.VAL.PATH+'_'+self._C.PROBLEM.INSTANCE_SEG.DATA_CHANNELS+'_'+self._C.PROBLEM.INSTANCE_SEG.DATA_CONTOUR_MODE
         self._C.DATA.VAL.INSTANCE_CHANNELS_MASK_DIR = self._C.DATA.VAL.GT_PATH+'_'+self._C.PROBLEM.INSTANCE_SEG.DATA_CHANNELS+'_'+self._C.PROBLEM.INSTANCE_SEG.DATA_CONTOUR_MODE
         # If value is not the default
-        if self._C.DATA.VAL.BINARY_MASKS == os.path.join("user_data", 'val', 'bin_mask'):
-            self._C.DATA.VAL.BINARY_MASKS = os.path.join(self._C.DATA.VAL.PATH, '..', 'bin_mask')
         self._C.DATA.VAL.DETECTION_MASK_DIR = self._C.DATA.VAL.GT_PATH+'_detection_masks'
         self._C.DATA.VAL.SSL_SOURCE_DIR = self._C.DATA.VAL.PATH+'_ssl_source'
         self._C.DATA.TEST.INSTANCE_CHANNELS_DIR = self._C.DATA.TEST.PATH+'_'+self._C.PROBLEM.INSTANCE_SEG.DATA_CHANNELS+'_'+self._C.PROBLEM.INSTANCE_SEG.DATA_CONTOUR_MODE

{biapy-3.4.2 → biapy-3.4.3}/biapy/data/data_2D_manipulation.py

@@ -102,7 +102,7 @@ def load_and_prepare_2D_train_data(train_path, train_mask_path, cross_val=False,
     filenames : List of str
         Loaded train filenames.
 
-    val_index : List of ints
+    test_index : List of ints
         Indexes of the samples beloging to the validation.
 
     Examples
@@ -144,7 +144,7 @@ def load_and_prepare_2D_train_data(train_path, train_mask_path, cross_val=False,
     print("### LOAD ###")
 
     # Disable crops when random_crops_in_DA is selected
-    # Disable crops when random_crops_in_DA is selected
+    delay_crop = False
     if random_crops_in_DA:
         crop = False  
     else:
@@ -154,7 +154,6 @@ def load_and_prepare_2D_train_data(train_path, train_mask_path, cross_val=False,
             delay_crop = True  
         else:
             crop = True
-            delay_crop = False  
 
     # Check validation
     if val_split > 0 or cross_val:
@@ -172,6 +171,19 @@ def load_and_prepare_2D_train_data(train_path, train_mask_path, cross_val=False,
         Y_train, _, _, _ = load_data_from_dir(train_mask_path, crop=crop, crop_shape=scrop, overlap=ov, padding=padding, 
             return_filenames=True, check_channel=False, check_drange=False, reflect_to_complete_shape=reflect_to_complete_shape,
             preprocess_cfg=preprocess_cfg, is_mask=is_y_mask, preprocess_f=preprocess_f)
+        
+        # Check that the shape of all images match
+        if isinstance(Y_train, list):
+            for i in range(len(Y_train)):
+                xshape = X_train[i].shape
+                yshape = Y_train[i].shape
+                real_x_shape = (xshape[0]*y_upscaling[0], xshape[1]*y_upscaling[1], xshape[2]) 
+                real_y_shape = (yshape[0]*y_upscaling[0], yshape[1]*y_upscaling[1], yshape[2])
+                print(real_x_shape,real_y_shape)
+                if real_x_shape != real_y_shape:
+                    raise ValueError(f"There is a mismatch between input image and its corresponding ground truth ({real_x_shape} vs "
+                        f"{real_y_shape}). Please check the images. Specifically, the sample that doesn't match is the number {i}"
+                        f" (file: {t_filenames[i]})")
     else:
         Y_train = None
     
@@ -257,22 +269,36 @@ def load_and_prepare_2D_train_data(train_path, train_mask_path, cross_val=False,
             skf = StratifiedKFold(n_splits=cross_val_nsplits, shuffle=shuffle_val,
                 random_state=seed)
             fold = 1
-            train_index, val_index = None, None
+            train_index, test_index = None, None
 
             y_len = len(Y_train) if Y_train is not None else len(X_train)
             for t_index, te_index in skf.split(np.zeros(len(X_train)), np.zeros(y_len)):
                 if cross_val_fold == fold:
-                    X_train, X_val = X_train[t_index], X_train[te_index]
+                    if not isinstance(X_train, list):
+                        X_train, X_val = X_train[t_index], X_train[te_index]
+                    else:
+                        X_val = []
+                        for val_idx in te_index:
+                            X_val.append(X_train[val_idx])
+                        for val_idx in te_index:
+                            del X_val[val_idx]
                     if Y_train is not None:
-                        Y_train, Y_val = Y_train[t_index], Y_train[te_index]
+                        if not isinstance(Y_train, list):
+                            Y_train, Y_val = Y_train[t_index], Y_train[te_index]
+                        else:
+                            Y_val = []
+                            for val_idx in te_index:
+                                Y_val.append(Y_val[val_idx])
+                            for val_idx in te_index:
+                                del Y_val[val_idx]
                     train_index, test_index = t_index.copy(), te_index.copy()
                     break
                 fold+= 1
 
-            if len(val_index) > 5:
-                print("Fold number {}. Printing the first 5 ids: {}".format(fold, val_index[:5]))
+            if len(test_index) > 5:
+                print("Fold number {}. Printing the first 5 ids: {}".format(fold, test_index[:5]))
             else:
-                print("Fold number {}. Indexes used in cross validation: {}".format(fold, val_index))
+                print("Fold number {}. Indexes used in cross validation: {}".format(fold, test_index))
 
             # Then crop after cross validation
             if delay_crop:
@@ -280,8 +306,8 @@ def load_and_prepare_2D_train_data(train_path, train_mask_path, cross_val=False,
                 data = []
                 for img_num in range(len(X_train)):
                     if X_train[img_num].shape != crop_shape[:2]+(X_train[img_num].shape[-1],):
-                        img = X_train[img_num]
-                        img = crop_2D_data_with_overlap(X_train[img_num][0] if isinstance(X_train, list) else X_train[img_num], 
+                        img = crop_data_with_overlap(
+                            X_train[img_num][0] if isinstance(X_train, list) else np.expand_dims(X_train[img_num],0), 
                             crop_shape[:2]+(X_train[img_num].shape[-1],), overlap=ov, padding=padding, verbose=False)
                     data.append(img)
                 X_train = np.concatenate(data)
@@ -293,8 +319,8 @@ def load_and_prepare_2D_train_data(train_path, train_mask_path, cross_val=False,
                     scrop = (crop_shape[0]*y_upscaling[0], crop_shape[1]*y_upscaling[1], crop_shape[2])
                     for img_num in range(len(Y_train)):
                         if Y_train[img_num].shape != scrop[:2]+(Y_train[img_num].shape[-1],):
-                            img = Y_train[img_num]
-                            img = crop_2D_data_with_overlap(Y_train[img_num][0] if isinstance(Y_train, list) else Y_train[img_num],
+                            img = crop_data_with_overlap(
+                                Y_train[img_num][0] if isinstance(Y_train, list) else np.expand_dims(Y_train[img_num],0),
                                 scrop[:2]+(Y_train[img_num].shape[-1],), overlap=ov, padding=padding, verbose=False)
                         data_mask.append(img)
                     Y_train = np.concatenate(data_mask)
@@ -304,8 +330,8 @@ def load_and_prepare_2D_train_data(train_path, train_mask_path, cross_val=False,
                 data = []
                 for img_num in range(len(X_val)):
                     if X_val[img_num].shape != crop_shape[:2]+(X_val[img_num].shape[-1],):
-                        img = X_val[img_num]
-                        img = crop_2D_data_with_overlap(X_val[img_num][0] if isinstance(X_val, list) else X_val[img_num], 
+                        img = crop_data_with_overlap(
+                            X_val[img_num][0] if isinstance(X_val, list) else np.expand_dims(X_val[img_num],0), 
                             crop_shape[:2]+(X_val[img_num].shape[-1],), overlap=ov, padding=padding, verbose=False)
                     data.append(img)
                 X_val = np.concatenate(data)
@@ -317,13 +343,26 @@ def load_and_prepare_2D_train_data(train_path, train_mask_path, cross_val=False,
                     scrop = (crop_shape[0]*y_upscaling[0], crop_shape[1]*y_upscaling[1], crop_shape[2])
                     for img_num in range(len(Y_val)):
                         if Y_val[img_num].shape != scrop[:2]+(Y_val[img_num].shape[-1],):
-                            img = Y_val[img_num]
-                            img = crop_2D_data_with_overlap(Y_val[img_num][0] if isinstance(Y_val, list) else Y_val[img_num],
+                            img = crop_data_with_overlap(
+                                Y_val[img_num][0] if isinstance(Y_val, list) else np.expand_dims(Y_val[img_num],0),
                                 scrop[:2]+(Y_val[img_num].shape[-1],), overlap=ov, padding=padding, verbose=False)
                         data_mask.append(img)
                     Y_val = np.concatenate(data_mask)
                     del data_mask
 
+    # Check that the shape of all images match
+    if Y_train is not None:
+        if not isinstance(X_train, list):
+            if Y_train.shape[0] != X_train.shape[0]:
+                raise ValueError(f"Seems that input images do not correspond to their ground truth in shape ({X_train.shape[0]} samples vs "
+                    f"{Y_train.shape[0]} samples). Please check the images. If you are in super-resolution workflow maybe you did not "
+                    "configured properly 'PROBLEM.SUPER_RESOLUTION.UPSCALING' variable")
+        else:
+            if Y_train[0].shape[0] != X_train[0].shape[0]:
+                raise ValueError(f"Seems that input images do not correspond to their ground truth in shape ({X_train[0].shape[0]} samples vs "
+                    f"{Y_train[0].shape[0]} samples). Please check the images. If you are in super-resolution workflow maybe you did not "
+                    "configured properly 'PROBLEM.SUPER_RESOLUTION.UPSCALING' variable")
+
     s = X_train.shape if not isinstance(X_train, list) else (len(X_train),)+X_train[0].shape[1:]
     if Y_train is not None:
         sm = Y_train.shape if not isinstance(Y_train, list) else (len(Y_train),)+Y_train[0].shape[1:]
@@ -344,7 +383,7 @@ def load_and_prepare_2D_train_data(train_path, train_mask_path, cross_val=False,
         if not cross_val:
             return X_train, Y_train, X_val, Y_val, t_filenames
         else:
-            return X_train, Y_train, X_val, Y_val, t_filenames, val_index
+            return X_train, Y_train, X_val, Y_val, t_filenames, test_index
     else:
         print("*** Loaded train data shape is: {}".format(s))
         print("### END LOAD ###")
@@ -443,7 +482,11 @@ def crop_data_with_overlap(data, crop_shape, data_mask=None, overlap=(0,0), padd
            #     **** New data shape is: (3960, 256, 256, 1)
     """
 
+    if data.ndim != 4:
+        raise ValueError("data expected to be 4 dimensional, given {}".format(data.shape))
     if data_mask is not None:
+        if data.ndim != 4:
+            raise ValueError("data mask expected to be 4 dimensional, given {}".format(data_mask.shape))
         if data.shape[:-1] != data_mask.shape[:-1]:
             raise ValueError("data and data_mask shapes mismatch: {} vs {}".format(data.shape[:-1], data_mask.shape[:-1]))
 
@@ -799,8 +842,8 @@ def merge_data_with_overlap(data, original_shape, data_mask=None, overlap=(0,0),
         return merged_data
 
 
-def load_data_classification(data_dir, patch_shape, convert_to_rgb=True, expected_classes=None, cross_val=False, cross_val_nsplits=5, cross_val_fold=1, 
-    val_split=0.1, seed=0, shuffle_val=True):
+def load_data_classification(data_dir, patch_shape, convert_to_rgb=True, expected_classes=None, cross_val=False, cross_val_nsplits=5, 
+    cross_val_fold=1, val_split=0.1, seed=0, shuffle_val=True, preprocess_cfg=None, preprocess_f=None):
     """
     Load data to train classification methods.
 
@@ -837,6 +880,12 @@ def load_data_classification(data_dir, patch_shape, convert_to_rgb=True, expecte
     shuffle_val : bool, optional
         Take random training examples to create validation data.
 
+    preprocess_cfg : dict, optional
+        Configuration parameters for preprocessing, is necessary in case you want to apply any preprocessing.
+    
+    preprocess_f : function, optional
+        The preprocessing function, is necessary in case you want to apply any preprocessing.
+
     Returns
     -------
     X_data : 4D Numpy array
@@ -854,7 +903,7 @@ def load_data_classification(data_dir, patch_shape, convert_to_rgb=True, expecte
     all_ids : List of str
         Loaded data filenames.
 
-    val_index : List of ints
+    test_index : List of ints
         Indexes of the samples beloging to the validation.
     """
 
@@ -889,7 +938,8 @@ def load_data_classification(data_dir, patch_shape, convert_to_rgb=True, expecte
 
         # Loading images 
         images, _, _, image_ids = load_data_from_dir(f, return_filenames=True, crop_shape=patch_shape, 
-            convert_to_rgb=convert_to_rgb)
+            convert_to_rgb=convert_to_rgb, preprocess_cfg=preprocess_cfg, is_mask=False, 
+            preprocess_f=preprocess_f)
 
         X_data.append(images)
         Y_data.append((c_num,)*len(ids))
@@ -908,26 +958,41 @@ def load_data_classification(data_dir, patch_shape, convert_to_rgb=True, expecte
     if create_val:
         print("Creating validation data")
         if not cross_val:
+            if len(X_data) == 1:
+                raise ValueError("Validation data can not be extracted from training data as it only has one sample. Please check the data.")
             X_data, X_val, Y_data, Y_val = train_test_split(
                 X_data, Y_data, test_size=val_split, shuffle=shuffle_val, random_state=seed,
                 stratify=Y_data)
         else:
+            if len(X_data) < cross_val_nsplits:
+                raise ValueError(f"Validation data can not be extracted from training data as the number of splits ({cross_val_nsplits}) "
+                    f"is greater than the number of samples {len(X_data)}. Please check the data.")
             skf = StratifiedKFold(n_splits=cross_val_nsplits, shuffle=shuffle_val,
                 random_state=seed)
             fold = 1
-            train_index, val_index = None, None
+            train_index, test_index = None, None
 
             for t_index, te_index in skf.split(X_data, Y_data):
                 if cross_val_fold == fold:
-                    X_data, X_val = X_data[t_index], X_data[te_index]
-                    Y_data, Y_val = Y_data[t_index], Y_data[te_index]
-                    train_index, val_index = t_index.copy(), te_index.copy()
+                    if not isinstance(X_data, list):
+                        X_data, X_val = X_data[t_index], X_data[te_index]
+                        Y_data, Y_val = Y_data[t_index], Y_data[te_index]
+                    else:
+                        X_val = []
+                        Y_val = []
+                        for val_idx in te_index:
+                            X_val.append(X_data[val_idx])
+                            Y_val.append(Y_data[val_idx])
+                        for val_idx in te_index:
+                            del X_val[val_idx]
+                            del Y_val[val_idx]
+                    train_index, test_index = t_index.copy(), te_index.copy()
                     break
                 fold+= 1
-            if len(val_index) > 5:
-                print("Fold number {}. Printing the first 5 ids: {}".format(fold, val_index[:5]))
+            if len(test_index) > 5:
+                print("Fold number {}. Printing the first 5 ids: {}".format(fold, test_index[:5]))
             else:
-                print("Fold number {}. Indexes used in cross validation: {}".format(fold, val_index))
+                print("Fold number {}. Indexes used in cross validation: {}".format(fold, test_index))
 
     if create_val:
         print("*** Loaded train data shape is: {}".format(X_data.shape))
@@ -935,7 +1000,7 @@ def load_data_classification(data_dir, patch_shape, convert_to_rgb=True, expecte
         if not cross_val:
             return X_data, Y_data, X_val, Y_val, all_ids
         else:
-            return X_data, Y_data, X_val, Y_val, all_ids, val_index
+            return X_data, Y_data, X_val, Y_val, all_ids, test_index
     else:
         print("*** Loaded train data shape is: {}".format(X_data.shape))
         return X_data, Y_data, all_ids