biapy 3.4.3__tar.gz → 3.4.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: biapy
-Version: 3.4.3
+Version: 3.4.4
 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.3 → biapy-3.4.4}/biapy/config/config.py

@@ -216,6 +216,13 @@ class Config:
         _C.DATA.TRAIN.IN_MEMORY = True
         _C.DATA.TRAIN.PATH = os.path.join("user_data", 'train', 'x')
         _C.DATA.TRAIN.GT_PATH = os.path.join("user_data", 'train', 'y')
+        # Whether if your input Zarr contains the raw images and labels together or not. Use 'DATA.TRAIN.INPUT_ZARR_MULTIPLE_DATA_RAW_PATH'
+        # and 'DATA.TRAIN.INPUT_ZARR_MULTIPLE_DATA_GT_PATH' to determine the tag to find within the Zarr
+        _C.DATA.TRAIN.INPUT_ZARR_MULTIPLE_DATA = False
+        # Paths to the raw and gt within the Zarr file. Only used when 'DATA.TRAIN.INPUT_ZARR_MULTIPLE_DATA' is True.
+        # E.g. 'volumes.raw' for raw and 'volumes.labels.neuron_ids' for GT path. 
+        _C.DATA.TRAIN.INPUT_ZARR_MULTIPLE_DATA_RAW_PATH = ''
+        _C.DATA.TRAIN.INPUT_ZARR_MULTIPLE_DATA_GT_PATH = ''
         # File to load/save data prepared with the appropiate channels in a instance segmentation problem.
         # E.g. _C.PROBLEM.TYPE ='INSTANCE_SEG' and _C.PROBLEM.INSTANCE_SEG.DATA_CHANNELS != 'B'
         _C.DATA.TRAIN.INSTANCE_CHANNELS_DIR = os.path.join("user_data", 'train', 'x_'+_C.PROBLEM.INSTANCE_SEG.DATA_CHANNELS)
@@ -240,9 +247,9 @@ class Config:
         # Minimum foreground percentage that each image loaded need to have to not discard it (only used when TRAIN.IN_MEMORY == True). 
         # This option is only valid for SEMANTIC_SEG, INSTANCE_SEG and DETECTION. 
         _C.DATA.TRAIN.MINIMUM_FOREGROUND_PER = -1.
-        # Order of the axes of the image when using Zarr/H5 images to train
+        # Order of the axes of the image when using Zarr/H5 images in train data.
         _C.DATA.TRAIN.INPUT_IMG_AXES_ORDER = 'TZCYX'
-        # Order of the axes of the mask when using Zarr/H5 images to train
+        # Order of the axes of the mask when using Zarr/H5 images in train data.
         _C.DATA.TRAIN.INPUT_MASK_AXES_ORDER = 'TZCYX'
 
         # PREPROCESSING
@@ -370,6 +377,13 @@ class Config:
         _C.DATA.VAL.PATH = os.path.join("user_data", 'val', 'x')
         # Path to the validation data mask. Used when _C.DATA.VAL.FROM_TRAIN = False
         _C.DATA.VAL.GT_PATH = os.path.join("user_data", 'val', 'y')
+        # Whether if your input Zarr contains the raw images and labels together or not. Use 'DATA.VAL.INPUT_ZARR_MULTIPLE_DATA_RAW_PATH'
+        # and 'DATA.VAL.INPUT_ZARR_MULTIPLE_DATA_GT_PATH' to determine the tag to find within the Zarr
+        _C.DATA.VAL.INPUT_ZARR_MULTIPLE_DATA = False
+        # Paths to the raw and gt within the Zarr file. Only used when 'DATA.VAL.INPUT_ZARR_MULTIPLE_DATA' is True.
+        # E.g. 'volumes.raw' for raw and 'volumes.labels.neuron_ids' for GT path. 
+        _C.DATA.VAL.INPUT_ZARR_MULTIPLE_DATA_RAW_PATH = ''
+        _C.DATA.VAL.INPUT_ZARR_MULTIPLE_DATA_GT_PATH = ''
         # File to load/save data prepared with the appropiate channels in a instance segmentation problem.
         # E.g. _C.PROBLEM.TYPE ='INSTANCE_SEG' and _C.PROBLEM.INSTANCE_SEG.DATA_CHANNELS != 'B'
         _C.DATA.VAL.INSTANCE_CHANNELS_DIR = os.path.join("user_data", 'val', 'x_'+_C.PROBLEM.INSTANCE_SEG.DATA_CHANNELS)
@@ -388,10 +402,9 @@ class Config:
         _C.DATA.VAL.PADDING = (0,0)
         # Not used yet.
         _C.DATA.VAL.RESOLUTION = (-1,)
-        # Order of the axes of the image when using Zarr/H5 images to train
+        # Order of the axes of the image when using Zarr/H5 images in validation data.
         _C.DATA.VAL.INPUT_IMG_AXES_ORDER = 'TZCYX'
-        # Order of the axes of the mask when using Zarr/H5 images to train. 
-        # detection the mask
+        # Order of the axes of the mask when using Zarr/H5 images in validation data.
         _C.DATA.VAL.INPUT_MASK_AXES_ORDER = 'TZCYX'
 
         #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -820,6 +833,16 @@ class Config:
         _C.TEST.BY_CHUNKS.FLUSH_EACH = 100
         # Input Numpy/Zarr/H5 image's axes order. Options: ['TZCYX', 'TZYXC', 'ZCYX', 'ZYXC']
         _C.TEST.BY_CHUNKS.INPUT_IMG_AXES_ORDER = 'TZCYX'
+        # Input Numpy/Zarr/H5 image's axes order. Options: ['TZCYX', 'TZYXC', 'ZCYX', 'ZYXC']
+        _C.TEST.BY_CHUNKS.INPUT_MASK_AXES_ORDER = 'TZCYX'
+        # Whether if your input Zarr contains the raw images and labels together or not. Use 'TEST.BY_CHUNKS.INPUT_ZARR_MULTIPLE_DATA_RAW_PATH'
+        # and 'TEST.BY_CHUNKS.INPUT_ZARR_MULTIPLE_DATA_GT_PATH' to determine the tag to find within the Zarr
+        _C.TEST.BY_CHUNKS.INPUT_ZARR_MULTIPLE_DATA = False
+        # Paths to the raw and gt within the Zarr file. Only used when 'TEST.BY_CHUNKS.INPUT_ZARR_MULTIPLE_DATA' is True.
+        # E.g. 'volumes.raw' for raw and 'volumes.labels.neuron_ids' for GT path. 
+        _C.TEST.BY_CHUNKS.INPUT_ZARR_MULTIPLE_DATA_RAW_PATH = ''
+        _C.TEST.BY_CHUNKS.INPUT_ZARR_MULTIPLE_DATA_GT_PATH = ''
+        
         # Whether if after reconstructing the prediction the pipeline will continue each workflow specific steps. For this process
         # the prediction image needs to be loaded into memory so be sure that it can fit in you memory. E.g. in instance 
         # segmentation the instances will be created from the prediction.

{biapy-3.4.3 → biapy-3.4.4}/biapy/data/data_3D_manipulation.py

@@ -6,7 +6,7 @@ from skimage.io import imread
 from tqdm import tqdm
 from sklearn.model_selection import train_test_split, StratifiedKFold
 
-from biapy.utils.util import load_3d_images_from_dir, order_dimensions, read_chunked_data
+from biapy.utils.util import load_3d_images_from_dir, order_dimensions, read_chunked_data, read_chunked_nested_data
 from biapy.utils.misc import is_main_process
 
 def load_and_prepare_3D_data(train_path, train_mask_path, cross_val=False, cross_val_nsplits=5, cross_val_fold=1, 
@@ -381,7 +381,7 @@ def load_and_prepare_3D_data(train_path, train_mask_path, cross_val=False, cross
 
 def load_and_prepare_3D_efficient_format_data(train_path, train_mask_path, input_img_axes, input_mask_axes=None, cross_val=False, 
     cross_val_nsplits=5, cross_val_fold=1, val_split=0.1, seed=0, shuffle_val=True, crop_shape=(80, 80, 80, 1), y_upscaling=(1,1,1), 
-    ov=(0,0,0), padding=(0,0,0), minimum_foreground_perc=-1):
+    ov=(0,0,0), padding=(0,0,0), minimum_foreground_perc=-1, multiple_data_within_zarr=None):
     """
     Load train and validation images from the given paths to create 3D data.
 
@@ -433,6 +433,9 @@ def load_and_prepare_3D_efficient_format_data(train_path, train_mask_path, input
     minimum_foreground_perc : float, optional
         Minimum percetnage of foreground that a sample need to have no not be discarded. 
 
+    multiple_data_within_zarr : dict, optional
+        Additional information of where to find the data within the Zarr files.
+
     Returns
     -------
     X_train : 5D Numpy array
@@ -457,7 +460,9 @@ def load_and_prepare_3D_efficient_format_data(train_path, train_mask_path, input
         create_val = False
 
     print("0) Loading train image information . . .")
-    X_train, X_train_total_patches = load_3D_efficient_files(train_path, input_img_axes, crop_shape, ov, padding)
+    data_within_zarr_path = multiple_data_within_zarr['raw_path'] if multiple_data_within_zarr is not None else None
+    X_train, X_train_total_patches = load_3D_efficient_files(train_path, input_img_axes, crop_shape, ov, padding,
+        data_within_zarr_path=data_within_zarr_path)
 
     if train_mask_path is not None:
         if input_mask_axes is None:
@@ -465,8 +470,11 @@ def load_and_prepare_3D_efficient_format_data(train_path, train_mask_path, input
 
         print("0) Loading train GT information . . .")
         scrop = (crop_shape[0]*y_upscaling[0], crop_shape[1]*y_upscaling[1], crop_shape[2]*y_upscaling[2], crop_shape[3])
+        data_within_zarr_path = None
+        if multiple_data_within_zarr is not None and multiple_data_within_zarr['use_gt_path']:
+            data_within_zarr_path = multiple_data_within_zarr['gt_path']
         Y_train, Y_train_total_patches = load_3D_efficient_files(train_mask_path, input_mask_axes, scrop, ov, padding, 
-            check_channel=False)
+            check_channel=False, data_within_zarr_path=data_within_zarr_path)
 
         for i in range(len(Y_train_total_patches)):
             if Y_train_total_patches[i] != X_train_total_patches[i]:
@@ -629,7 +637,8 @@ def load_and_prepare_3D_efficient_format_data(train_path, train_mask_path, input
             print("*** Loaded train GT shape is: {}".format(yshape))
         return X_train, Y_train
 
-def load_3D_efficient_files(data_path, input_axes, crop_shape, overlap, padding, check_channel=True):
+def load_3D_efficient_files(data_path, input_axes, crop_shape, overlap, padding, check_channel=True, 
+    data_within_zarr_path=None):
     """
     Load information of all patches that can be extracted from all the Zarr/H5 samples in ``data_path``.
 
@@ -654,6 +663,9 @@ def load_3D_efficient_files(data_path, input_axes, crop_shape, overlap, padding,
     check_channel : bool, optional
         Whether to check if the crop_shape channel matches with the loaded images' one. 
         
+    data_within_zarr_path : str, optional
+        Path to find the data within the Zarr file. E.g. 'volumes.labels.neuron_ids'.
+
     Returns
     -------
     data_info : dict
@@ -669,8 +681,11 @@ def load_3D_efficient_files(data_path, input_axes, crop_shape, overlap, padding,
 
     for i, filename in enumerate(data_path):
         print(f"Reading Zarr/H5 file: {filename}")
-        file, data = read_chunked_data(filename)
-
+        if data_within_zarr_path:
+            file, data = read_chunked_nested_data(filename, data_within_zarr_path)
+        else:
+            file, data = read_chunked_data(filename)
+        
         # Modify crop_shape with the channel
         c_index = -1
         try:
@@ -714,6 +729,53 @@ def load_3D_efficient_files(data_path, input_axes, crop_shape, overlap, padding,
     
     return data_info, data_total_patches
 
+def load_img_part_from_efficient_file(filepath, patch_coords, data_axis_order="ZYXC", data_path=None):
+    """
+    Loads from ``filepath`` the patch determined by ``patch_coords``.
+
+    Parameters
+    ----------
+    filepath : str
+        Path to the Zarr/H5 file to read the patch from. 
+
+    patch_coords : Tuple of tuples of ints
+        Coordinates of patch to extract. 
+        
+    data_axis_order : str
+        Order of axes of ``data``. E.g. 'TZCYX', 'TZYXC', 'ZCYX', 'ZYXC'. 
+
+    data_path : str, optional
+        Path to find the data within the Zarr file. E.g. 'volumes.labels.neuron_ids'.
+        
+    Returns
+    -------
+    img : Numpy array
+        Extracted patch. E.g. ``(z, y, x, channels)``.
+    """
+    if data_path is not None:
+        imgfile, img = read_chunked_nested_data(filepath, data_path)
+    else:
+        imgfile, img = read_chunked_data(filepath)
+
+    # Prepare slices to extract the patch
+    slices = []
+    for j in range(len(patch_coords)):
+        if isinstance(patch_coords[j], int):
+            # +1 to prevent 0 length axes that can not be removed with np.squeeze later
+            slices.append(slice(0,patch_coords[j]+1)) 
+        else:
+            slices.append(slice(patch_coords[j][0],patch_coords[j][1]))
+
+    slices = tuple(slices)
+    data_ordered_slices = tuple(order_dimensions(slices, input_order="ZYXC", output_order=data_axis_order,
+            default_value=0))
+    img = np.squeeze(np.array(img[data_ordered_slices]))
+    
+    if isinstance(imgfile, h5py.File):
+        imgfile.close()
+
+    return img 
+
 def crop_3D_data_with_overlap(data, vol_shape, data_mask=None, overlap=(0,0,0), padding=(0,0,0), verbose=True,
     median_padding=False):
     """Crop 3D data into smaller volumes with a defined overlap. The opposite function is :func:`~merge_3D_data_with_overlap`.

{biapy-3.4.3 → biapy-3.4.4}/biapy/data/generators/__init__.py

@@ -121,8 +121,9 @@ def create_train_val_augmentors(cfg, X_train, Y_train, X_val, Y_val, world_size,
     if cfg.DATA.TRAIN.IN_MEMORY:
         data_mode = "in_memory"
     else:
-        if cfg.PROBLEM.NDIM == '3D' and X_train is not None and isinstance(X_train, list) and isinstance(X_train[0], dict) and \
-            'filepath' in X_train[0] and ('.zarr' in X_train[0]['filepath'] or '.h5' in X_train[0]['filepath']):
+        if cfg.PROBLEM.NDIM == '3D' and X_train is not None and (isinstance(X_train, list) or isinstance(X_train, dict)) and \
+            isinstance(X_train[0], dict) and 'filepath' in X_train[0] and ('.zarr' in X_train[0]['filepath'] or \
+            '.h5' in X_train[0]['filepath']):
             data_mode = "chunked_data"
         else:
             data_mode = "not_in_memory"
@@ -226,8 +227,9 @@ def create_train_val_augmentors(cfg, X_train, Y_train, X_val, Y_val, world_size,
         if cfg.DATA.VAL.IN_MEMORY:
             val_data_mode = "in_memory"
         else:
-            if cfg.PROBLEM.NDIM == '3D' and X_val is not None and isinstance(X_val, list) and isinstance(X_val[0], dict) and \
-                'filepath' in X_val[0] and ('.zarr' in X_val[0]['filepath'] or '.h5' in X_val[0]['filepath']):
+            if cfg.PROBLEM.NDIM == '3D' and X_val is not None and (isinstance(X_val, list) or isinstance(X_val, dict)) and \
+                isinstance(X_val[0], dict) and 'filepath' in X_val[0] and ('.zarr' in X_val[0]['filepath'] or \
+                '.h5' in X_val[0]['filepath']):
                 val_data_mode = "chunked_data"
             else:
                 val_data_mode = "not_in_memory"

{biapy-3.4.3 → biapy-3.4.4}/biapy/data/generators/pair_base_data_generator.py

@@ -18,6 +18,7 @@ from biapy.utils.util import img_to_onehot_encoding, pad_and_reflect, read_chunk
 from biapy.data.generators.augmentors import *
 from biapy.data.pre_processing import normalize, norm_range01, percentile_clip
 from biapy.utils.misc import is_main_process
+from biapy.data.data_3D_manipulation import load_img_part_from_efficient_file 
 
 class PairBaseDataGenerator(Dataset, metaclass=ABCMeta):
     """
@@ -832,39 +833,9 @@ class PairBaseDataGenerator(Dataset, metaclass=ABCMeta):
                 if self.Y_provided:
                     mask = np.squeeze(mask)
         else: # self.data_mode == "chunked_data"
-            imgfile, img = read_chunked_data(self.X[idx]['filepath'])
-
-            # Prepare slices to extract the patch
-            slices = []
-            for j in range(len(self.X[idx]['patch_coords'])):
-                if isinstance(self.X[idx]['patch_coords'][j], int):
-                    # +1 to prevent 0 length axes that can not be removed with np.squeeze later
-                    slices.append(slice(0,self.X[idx]['patch_coords'][j]+1)) 
-                else:
-                    slices.append(slice(self.X[idx]['patch_coords'][j][0],self.X[idx]['patch_coords'][j][1]))
-
-            img = np.squeeze(np.array(img[tuple(slices)]))
-            
-            if isinstance(imgfile, h5py.File):
-                imgfile.close()
-
+            img = load_img_part_from_efficient_file(self.X[idx]['filepath'], self.X[idx]['patch_coords'])
             if self.Y_provided:
-                maskfile, mask = read_chunked_data(self.Y[idx]['filepath'])
-
-                # Prepare slices to extract the patch
-                slices = []
-                for j in range(len(self.Y[idx]['patch_coords'])):
-                    if isinstance(self.Y[idx]['patch_coords'][j], int):
-                        # +1 to prevent 0 length axes that can not be removed with np.squeeze later
-                        slices.append(slice(0,self.Y[idx]['patch_coords'][j]+1)) 
-                    else:
-                        slices.append(slice(self.Y[idx]['patch_coords'][j][0],self.Y[idx]['patch_coords'][j][1]))
-
-                mask = np.squeeze(np.array(mask[tuple(slices)]))
-
-                if isinstance(maskfile, h5py.File):
-                    maskfile.close()
-                
+                mask = load_img_part_from_efficient_file(self.Y[idx]['filepath'], self.Y[idx]['patch_coords'])
         if self.Y_provided:
             img, mask = self.ensure_shape(img, mask)
         else:

{biapy-3.4.3 → biapy-3.4.4}/biapy/data/pre_processing.py

@@ -2,6 +2,7 @@ import os
 import torch
 import scipy
 import h5py
+import zarr
 import numpy as np
 from tqdm import tqdm
 import pandas as pd
@@ -17,8 +18,10 @@ from skimage.exposure import equalize_adapthist
 from skimage.color import rgb2gray
 from skimage.filters import gaussian, median
 
-from biapy.utils.util import load_data_from_dir, load_3d_images_from_dir, save_tif, write_chunked_data, read_chunked_data
+from biapy.utils.util import (load_data_from_dir, load_3d_images_from_dir, save_tif, write_chunked_data, read_chunked_data,
+    order_dimensions)
 from biapy.utils.misc import is_main_process
+from biapy.data.data_3D_manipulation import load_3D_efficient_files, load_img_part_from_efficient_file
 
 #########################
 # INSTANCE SEGMENTATION #
@@ -42,10 +45,41 @@ def create_instance_channels(cfg, data_type='train'):
     """
 
     assert data_type in ['train', 'val']
-
-    f_name = load_data_from_dir if cfg.PROBLEM.NDIM == '2D' else load_3d_images_from_dir
     tag = "TRAIN" if data_type == "train" else "VAL"
-    Y, _, _, filenames = f_name(getattr(cfg.DATA, tag).GT_PATH, check_drange=False, return_filenames=True)
+
+    # Checking if the user inputted Zarr/H5 files 
+    if getattr(cfg.DATA, tag).INPUT_ZARR_MULTIPLE_DATA:
+        zarr_files = sorted(next(os.walk(getattr(cfg.DATA, tag).PATH))[1])
+        h5_files = sorted(next(os.walk(getattr(cfg.DATA, tag).PATH))[2])
+    else:
+        zarr_files = sorted(next(os.walk(getattr(cfg.DATA, tag).GT_PATH))[1])
+        h5_files = sorted(next(os.walk(getattr(cfg.DATA, tag).GT_PATH))[2])
+    
+    # Find patches info so we can iterate over them to create the instance mask
+    working_with_zarr_h5_files = False
+    if cfg.PROBLEM.NDIM == '3D' and (len(zarr_files) > 0 and '.zarr' in zarr_files[0]) or \
+        (len(h5_files) > 0 and '.h5' in h5_files[0]):
+        working_with_zarr_h5_files = True
+        # Check if the raw images and labels are within the same file 
+        mult_dat = None
+        data_path = getattr(cfg.DATA, tag).GT_PATH
+        if getattr(cfg.DATA, tag).INPUT_ZARR_MULTIPLE_DATA:
+            data_path = getattr(cfg.DATA, tag).PATH
+        
+        if len(zarr_files) > 0 and '.zarr' in zarr_files[0]:
+            print("Working with Zarr files . . .")
+            img_files = [os.path.join(data_path, x) for x in zarr_files]
+        elif len(h5_files) > 0 and '.h5' in h5_files[0]:
+            print("Working with H5 files . . .")
+            img_files = [os.path.join(data_path, x) for x in h5_files]
+        
+        Y, Y_total_patches = load_3D_efficient_files(img_files, input_axes=getattr(cfg.DATA, tag).INPUT_IMG_AXES_ORDER, 
+            crop_shape=cfg.DATA.PATCH_SIZE, overlap=getattr(cfg.DATA, tag).OVERLAP, padding=getattr(cfg.DATA, tag).PADDING)
+    else:
+        f_name = load_data_from_dir if cfg.PROBLEM.NDIM == '2D' else load_3d_images_from_dir
+        Y, _, _, filenames = f_name(getattr(cfg.DATA, tag).GT_PATH, check_drange=False, return_filenames=True)
+    del zarr_files, h5_files
+
     print("Creating Y_{} channels . . .".format(data_type))
     if isinstance(Y, list):
         for i in tqdm(range(len(Y)), disable=not is_main_process()):
@@ -61,6 +95,78 @@ def create_instance_channels(cfg, data_type='train'):
             else:
                 Y[i] = labels_into_channels(Y[i], mode=cfg.PROBLEM.INSTANCE_SEG.DATA_CHANNELS, save_dir=getattr(cfg.PATHS, tag+'_INSTANCE_CHANNELS_CHECK'),
                     fb_mode=cfg.PROBLEM.INSTANCE_SEG.DATA_CONTOUR_MODE)
+    # Create the mask patch by patch (Zarr/H5)
+    elif working_with_zarr_h5_files and isinstance(Y, dict):
+        savepath = data_path+'_'+cfg.PROBLEM.INSTANCE_SEG.DATA_CHANNELS+'_'+cfg.PROBLEM.INSTANCE_SEG.DATA_CONTOUR_MODE
+        if 'D' in cfg.PROBLEM.INSTANCE_SEG.DATA_CONTOUR_MODE:
+            dtype_str = "float32"
+            raise ValueError("Currently distance creation using Zarr by chunks is not implemented.")
+        else:
+            dtype_str = "uint8"  
+            
+        mask = None
+        last_zarr_file = None 
+        for i in tqdm(range(len(Y.keys())), disable=not is_main_process()):
+            # Extract the patch to process
+            patch_coords = Y[i]['patch_coords']
+            img = load_img_part_from_efficient_file(Y[i]['filepath'], patch_coords, 
+                data_axis_order=getattr(cfg.DATA, tag).INPUT_IMG_AXES_ORDER,
+                data_path=getattr(cfg.DATA, tag).INPUT_ZARR_MULTIPLE_DATA_GT_PATH)
+            if img.ndim == 3: img = np.expand_dims(img,-1)
+            if img.ndim == 4: img = np.expand_dims(img,0)
+
+            # Create the instance mask
+            if cfg.MODEL.N_CLASSES > 2:
+                if img.shape[-1] != 2:
+                    raise ValueError("In instance segmentation, when 'MODEL.N_CLASSES' are more than 2 labels need to have two channels, "
+                        "e.g. (256,256,2), containing the instance segmentation map (first channel) and classification map (second channel).")
+                else:
+                    class_channel = np.expand_dims(img[...,1].copy(),-1)
+                    img = labels_into_channels(img, mode=cfg.PROBLEM.INSTANCE_SEG.DATA_CHANNELS, save_dir=getattr(cfg.PATHS, tag+'_INSTANCE_CHANNELS_CHECK'),
+                        fb_mode=cfg.PROBLEM.INSTANCE_SEG.DATA_CONTOUR_MODE)
+                    img = np.concatenate([img,class_channel],axis=-1)
+            else:
+                img = labels_into_channels(img, mode=cfg.PROBLEM.INSTANCE_SEG.DATA_CHANNELS, save_dir=getattr(cfg.PATHS, tag+'_INSTANCE_CHANNELS_CHECK'),
+                    fb_mode=cfg.PROBLEM.INSTANCE_SEG.DATA_CONTOUR_MODE)
+            img = img[0]
+
+            # Create the Zarr file where the mask will be placed 
+            if mask is None or os.path.basename(Y[i]['filepath']) != last_zarr_file: 
+                last_zarr_file = os.path.basename(Y[i]['filepath'])
+                imgfile, data = read_chunked_data(Y[i]['filepath'])
+                fname = os.path.join(savepath, os.path.basename(Y[i]['filepath']))
+                fid_mask = zarr.open_group(fname, mode="w")
+
+                # Determine data shape
+                out_data_shape = np.array(data.shape)
+                if "C" not in getattr(cfg.DATA, tag).INPUT_IMG_AXES_ORDER:
+                    out_data_shape = tuple(out_data_shape) + (img.shape[-1],)
+                    out_data_order = getattr(cfg.DATA, tag).INPUT_IMG_AXES_ORDER + "C"
+                    channel_pos = -1
+                else:
+                    out_data_shape[getattr(cfg.DATA, tag).INPUT_IMG_AXES_ORDER.index("C")] = 1
+                    out_data_shape = tuple(out_data_shape)
+                    out_data_order = getattr(cfg.DATA, tag).INPUT_IMG_AXES_ORDER
+                    channel_pos = getattr(cfg.DATA, tag).INPUT_IMG_AXES_ORDER.index("C")
+
+                mask = fid_mask.create_dataset("data", shape=out_data_shape, dtype=dtype_str)
+                del data, imgfile, fname
+            
+            # Adjust slices to calculate where to insert the predicted patch. This slice does not have into account the 
+            # channel so any of them can be inserted 
+            slices = (slice(patch_coords[0][0],patch_coords[0][1]), slice(patch_coords[1][0],patch_coords[1][1]),
+                slice(patch_coords[2][0],patch_coords[2][1]), slice(0,out_data_shape[channel_pos]))
+            data_ordered_slices = tuple(order_dimensions(slices, input_order="ZYXC", 
+                output_order=out_data_order, default_value=0))
+
+            # Adjust patch slice to transpose it before inserting intop the final data 
+            current_order = np.array(range(len(img.shape)))
+            transpose_order = order_dimensions(current_order, input_order="ZYXC", output_order=out_data_order,
+                default_value=np.nan)
+            transpose_order = [x for x in transpose_order if not np.isnan(x)]
+
+            # Place the patch into the Zarr
+            mask[data_ordered_slices] = img.transpose(transpose_order)        
     else:
         if cfg.MODEL.N_CLASSES > 2:
             if Y.shape[-1] != 2:
@@ -75,17 +181,18 @@ def create_instance_channels(cfg, data_type='train'):
             Y = labels_into_channels(Y, mode=cfg.PROBLEM.INSTANCE_SEG.DATA_CHANNELS, save_dir=getattr(cfg.PATHS, tag+'_INSTANCE_CHANNELS_CHECK'),
                 fb_mode=cfg.PROBLEM.INSTANCE_SEG.DATA_CONTOUR_MODE)
     
-    save_tif(Y, data_dir=getattr(cfg.DATA, tag).INSTANCE_CHANNELS_MASK_DIR, filenames=filenames, verbose=cfg.TEST.VERBOSE)
-    X, _, _, filenames = f_name(getattr(cfg.DATA, tag).PATH, return_filenames=True)
-    print("Creating X_{} channels . . .".format(data_type))
-    save_tif(X, data_dir=getattr(cfg.DATA, tag).INSTANCE_CHANNELS_DIR, filenames=filenames, verbose=cfg.TEST.VERBOSE)
-    
-    # Save original X data with the labels 
-    for i in range(min(3,len(X))):
-        if isinstance(X, list):
-            save_tif(X[i], getattr(cfg.PATHS, tag+'_INSTANCE_CHANNELS_CHECK'), filenames=['vol'+str(i)+".tif"], verbose=False)
-        else:
-            save_tif(np.expand_dims(X[i],0), getattr(cfg.PATHS, tag+'_INSTANCE_CHANNELS_CHECK'), filenames=['vol'+str(i)+".tif"], verbose=False)
+    if not working_with_zarr_h5_files:
+        save_tif(Y, data_dir=getattr(cfg.DATA, tag).INSTANCE_CHANNELS_MASK_DIR, filenames=filenames, verbose=cfg.TEST.VERBOSE)
+        X, _, _, filenames = f_name(getattr(cfg.DATA, tag).PATH, return_filenames=True)
+        print("Creating X_{} channels . . .".format(data_type))
+        save_tif(X, data_dir=getattr(cfg.DATA, tag).INSTANCE_CHANNELS_DIR, filenames=filenames, verbose=cfg.TEST.VERBOSE)
+
+        # Save original X data with the labels 
+        for i in range(min(3,len(X))):
+            if isinstance(X, list):
+                save_tif(X[i], getattr(cfg.PATHS, tag+'_INSTANCE_CHANNELS_CHECK'), filenames=['vol'+str(i)+".tif"], verbose=False)
+            else:
+                save_tif(np.expand_dims(X[i],0), getattr(cfg.PATHS, tag+'_INSTANCE_CHANNELS_CHECK'), filenames=['vol'+str(i)+".tif"], verbose=False)
 
 def create_test_instance_channels(cfg):
     """Create test new data with appropiate channels based on ``PROBLEM.INSTANCE_SEG.DATA_CHANNELS`` for instance segmentation.