biomedisa 24.8.11__py3-none-any.whl → 25.6.1__py3-none-any.whl

biomedisa/features/keras_helper.py

@@ -39,13 +39,13 @@ from tensorflow.keras.layers import (
 from tensorflow.keras import backend as K
 from tensorflow.keras.utils import to_categorical
 from tensorflow.keras.callbacks import Callback, ModelCheckpoint, EarlyStopping
-from biomedisa.features.DataGenerator import DataGenerator
+from biomedisa.features.DataGenerator import DataGenerator, welford_mean_std
 from biomedisa.features.PredictDataGenerator import PredictDataGenerator
-from biomedisa.features.biomedisa_helper import (unique,
+from biomedisa.features.biomedisa_helper import (unique, welford_mean_std,
     img_resize, load_data, save_data, set_labels_to_zero, id_generator, unique_file_path)
 from biomedisa.features.remove_outlier import clean, fill
 from biomedisa.features.active_contour import activeContour
-from tifffile import TiffFile, imread
+from tifffile import TiffFile, imread, imwrite
 import matplotlib.pyplot as plt
 import SimpleITK as sitk
 import tensorflow as tf
@@ -220,11 +220,11 @@ def compute_position(position, zsh, ysh, xsh):
                 position[k,l,m] = x+y+z
     return position
 
-def make_conv_block(nb_filters, input_tensor, block):
+def make_conv_block(nb_filters, input_tensor, block, dtype):
     def make_stage(input_tensor, stage):
         name = 'conv_{}_{}'.format(block, stage)
         x = Conv3D(nb_filters, (3, 3, 3), activation='relu',
-                   padding='same', name=name, data_format="channels_last")(input_tensor)
+                   padding='same', name=name, data_format="channels_last", dtype=dtype)(input_tensor)
         name = 'batch_norm_{}_{}'.format(block, stage)
         try:
             x = BatchNormalization(name=name, synchronized=True)(x)
@@ -266,56 +266,65 @@ def make_conv_block_resnet(nb_filters, input_tensor, block):
 
     return out
 
-def make_unet(input_shape, nb_labels, filters='32-64-128-256-512', resnet=False):
+def make_unet(bm, input_shape, nb_labels):
+    # enable mixed_precision
+    if bm.mixed_precision:
+        dtype = "float16"
+    else:
+        dtype = "float32"
 
+    # input
     nb_plans, nb_rows, nb_cols, _ = input_shape
+    inputs = Input(input_shape, dtype=dtype)
 
-    inputs = Input(input_shape)
-
-    filters = filters.split('-')
+    # configure number of layers and filters
+    filters = bm.network_filters.split('-')
     filters = np.array(filters, dtype=int)
     latent_space_size = filters[-1]
     filters = filters[:-1]
+
+    # initialize blocks
     convs = []
 
+    # encoder
     i = 1
     for f in filters:
         if i==1:
-            if resnet:
+            if bm.resnet:
                 conv = make_conv_block_resnet(f, inputs, i)
             else:
-                conv = make_conv_block(f, inputs, i)
+                conv = make_conv_block(f, inputs, i, dtype)
         else:
-            if resnet:
+            if bm.resnet:
                 conv = make_conv_block_resnet(f, pool, i)
             else:
-                conv = make_conv_block(f, pool, i)
+                conv = make_conv_block(f, pool, i, dtype)
         pool = MaxPooling3D(pool_size=(2, 2, 2))(conv)
         convs.append(conv)
         i += 1
 
-    if resnet:
+    # latent space
+    if bm.resnet:
         conv = make_conv_block_resnet(latent_space_size, pool, i)
     else:
-        conv = make_conv_block(latent_space_size, pool, i)
+        conv = make_conv_block(latent_space_size, pool, i, dtype)
     i += 1
 
+    # decoder
     for k, f in enumerate(filters[::-1]):
         up = Concatenate()([UpSampling3D(size=(2, 2, 2))(conv), convs[-(k+1)]])
-        if resnet:
+        if bm.resnet:
             conv = make_conv_block_resnet(f, up, i)
         else:
-            conv = make_conv_block(f, up, i)
+            conv = make_conv_block(f, up, i, dtype)
         i += 1
 
+    # final layer and output
     conv = Conv3D(nb_labels, (1, 1, 1), name=f'conv_{i}_1')(conv)
-
     x = Reshape((nb_plans * nb_rows * nb_cols, nb_labels))(conv)
     x = Activation('softmax')(x)
     outputs = Reshape((nb_plans, nb_rows, nb_cols, nb_labels))(x)
-
     model = Model(inputs=inputs, outputs=outputs)
-
     return model
 
 def get_labels(arr, allLabels):
@@ -415,6 +424,13 @@ def load_training_data(bm, img_list, label_list, channels, img_in=None, label_in
                 print(f'{os.path.basename(label_names[0])}:', 'Labels:', label_values[1:], 'Sizes:', counts[1:])
                 label = img_resize(label, bm.z_scale, bm.y_scale, bm.x_scale, labels=True)
 
+            # label channel must be 1 or 2 if using ignore mask
+            if len(label.shape)>3 and label.shape[3]>1 and not bm.ignore_mask:
+                InputError.message = 'Training labels must have one channel (gray values).'
+                raise InputError()
+            if len(label.shape)==3:
+                label = label.reshape(label.shape[0], label.shape[1], label.shape[2], 1)
+
             # if header is not single data stream Amira Mesh falling back to Multi-TIFF
             if extension != '.am':
                 extension, header = '.tif', None
@@ -424,7 +440,7 @@ def load_training_data(bm, img_list, label_list, channels, img_in=None, label_in
             else:
                 header = header[0]
 
-            # load first img
+            # load first image
             if any(img_list):
                 img, _ = load_data(img_names[0], 'first_queue')
                 if img is None:
@@ -436,14 +452,15 @@ def load_training_data(bm, img_list, label_list, channels, img_in=None, label_in
             else:
                 img = img_in
                 img_names = ['img_1']
-            if label_dim != img.shape[:3]:
+
+            # label and image dimensions must match
+            if label_dim[:3] != img.shape[:3]:
                 InputError.message = f'Dimensions of "{os.path.basename(img_names[0])}" and "{os.path.basename(label_names[0])}" do not match'
                 raise InputError()
 
-            # ensure images have channels >=1
+            # image channels must be >=1
             if len(img.shape)==3:
-                z_shape, y_shape, x_shape = img.shape
-                img = img.reshape(z_shape, y_shape, x_shape, 1)
+                img = img.reshape(img.shape[0], img.shape[1], img.shape[2], 1)
             if channels is None:
                 channels = img.shape[3]
             if channels != img.shape[3]:
@@ -462,30 +479,30 @@ def load_training_data(bm, img_list, label_list, channels, img_in=None, label_in
             # scale data to the range from 0 to 1
             if not bm.patch_normalization:
                 img = img.astype(np.float32)
-                for c in range(channels):
-                    img[:,:,:,c] -= np.amin(img[:,:,:,c])
-                    img[:,:,:,c] /= np.amax(img[:,:,:,c])
+                for ch in range(channels):
+                    img[...,ch] -= np.amin(img[...,ch])
+                    img[...,ch] /= np.amax(img[...,ch])
 
             # normalize first validation image
             if bm.normalize and np.any(normalization_parameters):
                 img = img.astype(np.float32)
-                for c in range(channels):
-                    mean, std = np.mean(img[:,:,:,c]), np.std(img[:,:,:,c])
-                    img[:,:,:,c] = (img[:,:,:,c] - mean) / std
-                    img[:,:,:,c] = img[:,:,:,c] * normalization_parameters[1,c] + normalization_parameters[0,c]
+                for ch in range(channels):
+                    mean, std = welford_mean_std(img[...,ch])
+                    img[...,ch] = (img[...,ch] - mean) / std
+                    img[...,ch] = img[...,ch] * normalization_parameters[1,ch] + normalization_parameters[0,ch]
 
             # get normalization parameters from first image
             if normalization_parameters is None:
                 normalization_parameters = np.zeros((2,channels))
                 if bm.normalize:
-                    for c in range(channels):
-                        normalization_parameters[0,c] = np.mean(img[:,:,:,c])
-                        normalization_parameters[1,c] = np.std(img[:,:,:,c])
+                    for ch in range(channels):
+                        normalization_parameters[:,ch] = welford_mean_std(img[...,ch])
 
             # pad data
             if not bm.scaling:
                 img_data_list = [img]
                 label_data_list = [label]
+                img_dtype = img.dtype
                 # no-scaling for list of images needs negative values as it encodes padded areas as -1
                 label_dtype = label.dtype
                 if label_dtype==np.uint8:
@@ -499,7 +516,7 @@ def load_training_data(bm, img_list, label_list, channels, img_in=None, label_in
 
                 for k in range(1, number_of_images):
 
-                    # append label
+                    # load label data and pre-process
                     if any(label_list):
                         a, _ = load_data(label_names[k], 'first_queue')
                         if a is None:
@@ -516,11 +533,21 @@ def load_training_data(bm, img_list, label_list, channels, img_in=None, label_in
                         label_values, counts = unique(a, return_counts=True)
                         print(f'{os.path.basename(label_names[k])}:', 'Labels:', label_values[1:], 'Sizes:', counts[1:])
                         a = img_resize(a, bm.z_scale, bm.y_scale, bm.x_scale, labels=True)
+
+                    # label channel must be 1 or 2 if using ignore mask
+                    if len(a.shape)>3 and a.shape[3]>1 and not bm.ignore_mask:
+                        InputError.message = 'Training labels must have one channel (gray values).'
+                        raise InputError()
+                    if len(a.shape)==3:
+                        a = a.reshape(a.shape[0], a.shape[1], a.shape[2], 1)
+
+                    # append label data
+                    if bm.scaling:
                         label = np.append(label, a, axis=0)
                     else:
                         label_data_list.append(a)
 
-                    # append image
+                    # load image data and pre-process
                     if any(img_list):
                         a, _ = load_data(img_names[k], 'first_queue')
                         if a is None:
@@ -528,12 +555,11 @@ def load_training_data(bm, img_list, label_list, channels, img_in=None, label_in
                             raise InputError()
                     else:
                         a = img_in[k]
-                    if label_dim != a.shape[:3]:
+                    if label_dim[:3] != a.shape[:3]:
                         InputError.message = f'Dimensions of "{os.path.basename(img_names[k])}" and "{os.path.basename(label_names[k])}" do not match'
                         raise InputError()
                     if len(a.shape)==3:
-                        z_shape, y_shape, x_shape = a.shape
-                        a = a.reshape(z_shape, y_shape, x_shape, 1)
+                        a = a.reshape(a.shape[0], a.shape[1], a.shape[2], 1)
                     if a.shape[3] != channels:
                         InputError.message = f'Number of channels must be {channels} for "{os.path.basename(img_names[k])}"'
                         raise InputError()
@@ -544,15 +570,17 @@ def load_training_data(bm, img_list, label_list, channels, img_in=None, label_in
                         a = img_resize(a, bm.z_scale, bm.y_scale, bm.x_scale)
                     if not bm.patch_normalization:
                         a = a.astype(np.float32)
-                        for c in range(channels):
-                            a[:,:,:,c] -= np.amin(a[:,:,:,c])
-                            a[:,:,:,c] /= np.amax(a[:,:,:,c])
+                        for ch in range(channels):
+                            a[...,ch] -= np.amin(a[...,ch])
+                            a[...,ch] /= np.amax(a[...,ch])
                     if bm.normalize:
                         a = a.astype(np.float32)
-                        for c in range(channels):
-                            mean, std = np.mean(a[:,:,:,c]), np.std(a[:,:,:,c])
-                            a[:,:,:,c] = (a[:,:,:,c] - mean) / std
-                            a[:,:,:,c] = a[:,:,:,c] * normalization_parameters[1,c] + normalization_parameters[0,c]
+                        for ch in range(channels):
+                            mean, std = welford_mean_std(a[...,ch])
+                            a[...,ch] = (a[...,ch] - mean) / std
+                            a[...,ch] = a[...,ch] * normalization_parameters[1,ch] + normalization_parameters[0,ch]
+
+                    # append image data
                     if bm.scaling:
                         img = np.append(img, a, axis=0)
                     else:
@@ -564,15 +592,14 @@ def load_training_data(bm, img_list, label_list, channels, img_in=None, label_in
         for img in img_data_list:
             target_y = max(target_y, img.shape[1])
             target_x = max(target_x, img.shape[2])
-        img = np.empty((0, target_y, target_x, channels), dtype=np.float32)
-        label = np.empty((0, target_y, target_x), dtype=label_dtype)
+        img = np.empty((0, target_y, target_x, channels), dtype=img_dtype)
+        label = np.empty((0, target_y, target_x, 2 if bm.ignore_mask else 1), dtype=label_dtype)
         for k in range(len(img_data_list)):
             pad_y = target_y - img_data_list[k].shape[1]
             pad_x = target_x - img_data_list[k].shape[2]
             pad_width = [(0, 0), (0, pad_y), (0, pad_x), (0, 0)]
             tmp = np.pad(img_data_list[k], pad_width, mode='constant', constant_values=0)
             img = np.append(img, tmp, axis=0)
-            pad_width = [(0, 0), (0, pad_y), (0, pad_x)]
             tmp = np.pad(label_data_list[k].astype(label_dtype), pad_width, mode='constant', constant_values=-1)
             label = np.append(label, tmp, axis=0)
 
@@ -586,13 +613,13 @@ def load_training_data(bm, img_list, label_list, channels, img_in=None, label_in
     else:
         # get labels
         if allLabels is None:
-            allLabels = unique(label)
+            allLabels = unique(label[...,0])
             index = np.argwhere(allLabels<0)
             allLabels = np.delete(allLabels, index)
 
         # labels must be in ascending order
         for k, l in enumerate(allLabels):
-            label[label==l] = k
+            label[...,0][label[...,0]==l] = k
 
     return img, label, allLabels, normalization_parameters, header, extension, channels
 
@@ -724,10 +751,11 @@ class Metrics(Callback):
                     m = rest % self.dim_img[2]
                     tmp_X = self.img[k:k+self.dim_patch[0],l:l+self.dim_patch[1],m:m+self.dim_patch[2]]
                     if self.patch_normalization:
-                        tmp_X = np.copy(tmp_X, order='C')
-                        for c in range(self.n_channels):
-                            tmp_X[:,:,:,c] -= np.mean(tmp_X[:,:,:,c])
-                            tmp_X[:,:,:,c] /= max(np.std(tmp_X[:,:,:,c]), 1e-6)
+                        tmp_X = tmp_X.copy().astype(np.float32)
+                        for ch in range(self.n_channels):
+                            mean, std = welford_mean_std(tmp_X[...,ch])
+                            tmp_X[...,ch] -= mean
+                            tmp_X[...,ch] /= max(std, 1e-6)
                     X_val[i] = tmp_X
 
                 # Prediction segmentation
@@ -752,6 +780,7 @@ class Metrics(Callback):
             # get result
             result = np.argmax(result, axis=-1)
             result = result.astype(np.uint8)
+            result = result.reshape(*result.shape, 1)
 
             # calculate standard accuracy
             if not self.train:
@@ -771,17 +800,17 @@ class Metrics(Callback):
                 logs['dice'] = dice
             else:
                 # save best model only
-                if epoch == 0 or round(dice,4) > max(self.history['val_dice']):
+                if epoch == 0 or dice > max(self.history['val_dice']):
                     self.model.save(str(self.path_to_model))
 
                 # add accuracy to history
-                self.history['loss'].append(round(logs['loss'],4))
-                self.history['accuracy'].append(round(logs['accuracy'],4))
+                self.history['loss'].append(logs['loss'])
+                self.history['accuracy'].append(logs['accuracy'])
                 if self.train_dice:
-                    self.history['dice'].append(round(logs['dice'],4))
-                self.history['val_accuracy'].append(round(accuracy,4))
-                self.history['val_dice'].append(round(dice,4))
-                self.history['val_loss'].append(round(val_loss,4))
+                    self.history['dice'].append(logs['dice'])
+                self.history['val_accuracy'].append(accuracy)
+                self.history['val_dice'].append(dice)
+                self.history['val_loss'].append(val_loss)
 
                 # tensorflow monitoring variables
                 logs['val_loss'] = val_loss
@@ -798,11 +827,11 @@ class Metrics(Callback):
 
                 # print accuracies
                 print('\nValidation history:')
-                print('train_acc:', self.history['accuracy'])
+                print("train_acc: [" + " ".join(f"{x:.4f}" for x in self.history['accuracy']) + "]")
                 if self.train_dice:
-                    print('train_dice:', self.history['dice'])
-                print('val_acc:', self.history['val_accuracy'])
-                print('val_dice:', self.history['val_dice'])
+                    print("train_dice: [" + " ".join(f"{x:.4f}" for x in self.history['dice']) + "]")
+                print("val_acc: [" + " ".join(f"{x:.4f}" for x in self.history['val_accuracy']) + "]")
+                print("val_dice: [" + " ".join(f"{x:.4f}" for x in self.history['val_dice']) + "]")
                 print('')
 
                 # early stopping
@@ -849,13 +878,13 @@ def categorical_crossentropy(true_labels, predicted_probs):
     # Clip predicted probabilities to avoid log(0) issues
     predicted_probs = np.clip(predicted_probs, 1e-7, 1 - 1e-7)
     predicted_probs = -np.log(predicted_probs)
-    zsh,ysh,xsh = true_labels.shape
+    zsh, ysh, xsh, _ = true_labels.shape
     # Calculate categorical crossentropy
     loss = 0
     for z in range(zsh):
         for y in range(ysh):
             for x in range(xsh):
-                l = true_labels[z,y,x]
+                l = true_labels[z,y,x,0]
                 loss += predicted_probs[z,y,x,l]
     loss = loss / float(zsh*ysh*xsh)
     return loss
@@ -879,6 +908,42 @@ def dice_coef_loss(nb_labels):
         return loss
     return loss_fn
 
+def custom_loss(y_true, y_pred):
+    # Extract labels and ignore mask
+    labels = tf.cast(y_true[..., 0], tf.int32)  # First channel contains class labels
+    ignore_mask = tf.cast(y_true[..., 1], tf.float32)  # Second channel contains mask (0 = ignore, 1 = include)
+
+    # Convert integer labels to one-hot encoding
+    y_true_one_hot = tf.one_hot(labels, depth=2)
+
+    # Clip y_pred to avoid log(0)
+    y_pred = tf.clip_by_value(y_pred, 1e-7, 1.0)
+
+    # Compute categorical cross-entropy
+    loss = -tf.reduce_sum(y_true_one_hot * tf.math.log(y_pred), axis=-1)
+
+    # Apply ignore mask (ignore = 0 → loss is zero, include = 1 → loss is counted)
+    loss = loss * ignore_mask
+
+    # Return mean loss over valid (non-ignored) samples
+    return tf.reduce_sum(loss) / tf.reduce_sum(ignore_mask)
+
+def custom_accuracy(y_true, y_pred):
+    labels = tf.cast(y_true[..., 0], tf.int32)  # Extract actual values
+    ignore_mask = y_true[..., 1]  # Extract mask (1 = include, 0 = ignore)
+
+    # Convert predictions to discrete values (assuming regression: round values)
+    y_pred_class = tf.argmax(y_pred, axis=-1, output_type=tf.int32)
+
+    # Compute correct predictions (1 where correct, 0 where incorrect)
+    correct_predictions = tf.cast(tf.equal(labels, y_pred_class), tf.float32)
+
+    # Apply ignore mask
+    masked_correct_predictions = correct_predictions * ignore_mask
+
+    # Compute accuracy only over valid (non-ignored) pixels
+    return tf.reduce_sum(masked_correct_predictions) / tf.reduce_sum(ignore_mask)
+
 def train_segmentation(bm):
 
     # training data
@@ -988,18 +1053,19 @@ def train_segmentation(bm):
               'n_channels': bm.channels,
               'augment': (bm.flip_x, bm.flip_y, bm.flip_z, bm.swapaxes, bm.rotate, bm.rotate3d),
               'patch_normalization': bm.patch_normalization,
-              'separation': bm.separation}
+              'separation': bm.separation,
+              'ignore_mask': bm.ignore_mask}
 
     # data generator
     validation_generator = None
-    training_generator = DataGenerator(bm.img_data, bm.label_data, list_IDs_fg, list_IDs_bg, True, True, False, **params)
+    training_generator = DataGenerator(bm.img_data, bm.label_data, list_IDs_fg, list_IDs_bg, True, True, **params)
     if bm.val_img_data is not None:
         if bm.val_dice:
             val_metrics = Metrics(bm, bm.val_img_data, bm.val_label_data, list_IDs_val_fg, (zsh_val, ysh_val, xsh_val), nb_labels, False)
         else:
             params['dim_img'] = (zsh_val, ysh_val, xsh_val)
             params['augment'] = (False, False, False, False, 0, False)
-            validation_generator = DataGenerator(bm.val_img_data, bm.val_label_data, list_IDs_val_fg, list_IDs_val_bg, True, False, False, **params)
+            validation_generator = DataGenerator(bm.val_img_data, bm.val_label_data, list_IDs_val_fg, list_IDs_val_bg, True, False, **params)
 
     # monitor dice score on training data
     if bm.train_dice:
@@ -1017,7 +1083,7 @@ def train_segmentation(bm):
     with strategy.scope():
 
         # build model
-        model = make_unet(input_shape, nb_labels, bm.network_filters, bm.resnet)
+        model = make_unet(bm, input_shape, nb_labels)
         model.summary()
 
         # pretrained model
@@ -1036,13 +1102,28 @@ def train_segmentation(bm):
                 layer.trainable = False
 
         # optimizer
-        sgd = SGD(learning_rate=bm.learning_rate, decay=1e-6, momentum=0.9, nesterov=True)
+        optimizer = SGD(learning_rate=bm.learning_rate, decay=1e-6, momentum=0.9, nesterov=True)
+        #optimizer = tf.keras.optimizers.Adam(learning_rate=bm.learning_rate, epsilon=1e-4) lr=0.0001
+        if bm.mixed_precision:
+            optimizer = tf.keras.mixed_precision.LossScaleOptimizer(optimizer, dynamic=False, initial_scale=128)
+
+        # Rename the function to appear as "accuracy" in logs
+        if bm.ignore_mask:
+            custom_accuracy.__name__ = "accuracy"
+            metrics=[custom_accuracy]
+        else:
+            metrics=['accuracy']
+
+        # loss function
+        if bm.ignore_mask:
+            loss=custom_loss
+        else:
+            loss=dice_coef_loss(nb_labels) if bm.dice_loss else 'categorical_crossentropy'
 
         # comile model
-        loss=dice_coef_loss(nb_labels) if bm.dice_loss else 'categorical_crossentropy'
         model.compile(loss=loss,
-                      optimizer=sgd,
-                      metrics=['accuracy'])
+                      optimizer=optimizer,
+                      metrics=metrics)
 
     # save meta data
     meta_data = MetaData(bm.path_to_model, configuration_data, allLabels,
@@ -1081,7 +1162,7 @@ def train_segmentation(bm):
         callbacks=callbacks,
         workers=bm.workers)
 
-def load_prediction_data(bm, channels, normalize, normalization_parameters,
+def load_prediction_data(bm, channels, normalization_parameters,
     region_of_interest, img, img_header, load_blockwise=False, z=None):
 
     # read image data
@@ -1109,10 +1190,9 @@ def load_prediction_data(bm, channels, normalize, normalization_parameters,
     if bm.acwe:
         img_data = img.copy()
 
-    # handle all images using number of channels >=1
+    # image data must have number of channels >=1
     if len(img.shape)==3:
-        z_shape, y_shape, x_shape = img.shape
-        img = img.reshape(z_shape, y_shape, x_shape, 1)
+        img = img.reshape(img.shape[0], img.shape[1], img.shape[2], 1)
     if img.shape[3] != channels:
         InputError.message = f'Number of channels must be {channels}.'
         raise InputError()
@@ -1127,22 +1207,27 @@ def load_prediction_data(bm, channels, normalize, normalization_parameters,
         region_of_interest = np.array([min_z,max_z,min_y,max_y,min_x,max_x,z_shape,y_shape,x_shape])
         z_shape, y_shape, x_shape = max_z-min_z, max_y-min_y, max_x-min_x
 
-    # scale/resize image data
-    img = img.astype(np.float32)
+    # resize image data
     if bm.scaling:
+        img = img.astype(np.float32)
         img = img_resize(img, bm.z_scale, bm.y_scale, bm.x_scale)
 
+    # scale image data
+    if not bm.patch_normalization:
+        img = img.astype(np.float32)
+        for ch in range(channels):
+            img[...,ch] -= np.amin(img[...,ch])
+            img[...,ch] /= np.amax(img[...,ch])
+
     # normalize image data
-    for c in range(channels):
-        img[:,:,:,c] -= np.amin(img[:,:,:,c])
-        img[:,:,:,c] /= np.amax(img[:,:,:,c])
-        if normalize:
-            mean, std = np.mean(img[:,:,:,c]), np.std(img[:,:,:,c])
-            img[:,:,:,c] = (img[:,:,:,c] - mean) / std
-            img[:,:,:,c] = img[:,:,:,c] * normalization_parameters[1,c] + normalization_parameters[0,c]
+    if bm.normalize:
+        img = img.astype(np.float32)
+        for ch in range(channels):
+            mean, std = welford_mean_std(img[...,ch])
+            img[...,ch] = (img[...,ch] - mean) / std
+            img[...,ch] = img[...,ch] * normalization_parameters[1,ch] + normalization_parameters[0,ch]
 
-    # limit intensity range
-    if normalize:
+        # limit intensity range
         img[img<0] = 0
         img[img>1] = 1
 
@@ -1185,6 +1270,20 @@ def gradient(volData):
     grad[grad>0]=1
     return grad
 
+@numba.jit(nopython=True)
+def scale_probabilities(final):
+    zsh, ysh, xsh, nb_labels = final.shape
+    for k in range(zsh):
+        for l in range(ysh):
+            for m in range(xsh):
+                scale_factor = 0
+                for n in range(nb_labels):
+                    scale_factor += final[k,l,m,n]
+                scale_factor = max(1, scale_factor)
+                for n in range(nb_labels):
+                    final[k,l,m,n] /= scale_factor
+    return final
+
 def predict_segmentation(bm, region_of_interest, channels, normalization_parameters):
 
     from mpi4py import MPI
@@ -1192,13 +1291,26 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
     rank = comm.Get_rank()
     ngpus = comm.Get_size()
 
+    # optional result paths
+    if bm.path_to_image:
+        filename, bm.extension = os.path.splitext(bm.path_to_final)
+        if bm.extension == '.gz':
+            bm.extension = '.nii.gz'
+            filename = filename[:-4]
+        path_to_cleaned = filename + '.cleaned' + bm.extension
+        path_to_filled = filename + '.filled' + bm.extension
+        path_to_cleaned_filled = filename + '.cleaned.filled' + bm.extension
+        path_to_refined = filename + '.refined' + bm.extension
+        path_to_acwe = filename + '.acwe' + bm.extension
+        path_to_probs = filename + '.probs.tif'
+
     # Set the visible GPU by ID
     gpus = tf.config.experimental.list_physical_devices('GPU')
     if gpus:
         try:
             # Restrict TensorFlow to only use the specified GPU
-            tf.config.experimental.set_visible_devices(gpus[rank], 'GPU')
-            tf.config.experimental.set_memory_growth(gpus[rank], True)
+            tf.config.experimental.set_visible_devices(gpus[rank % len(gpus)], 'GPU')
+            tf.config.experimental.set_memory_growth(gpus[rank % len(gpus)], True)
         except RuntimeError as e:
             print(e)
 
@@ -1209,7 +1321,7 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
     nb_labels = len(bm.allLabels)
     results['allLabels'] = bm.allLabels
 
-    # load model
+    # custom objects
     if bm.dice_loss:
         def loss_fn(y_true, y_pred):
             dice = 0
@@ -1220,25 +1332,30 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
             loss = 1 - dice
             return loss
         custom_objects = {'dice_coef_loss': dice_coef_loss,'loss_fn': loss_fn}
-        model = load_model(bm.path_to_model, custom_objects=custom_objects)
+    elif bm.ignore_mask:
+        custom_objects={'custom_loss': custom_loss}
     else:
-        model = load_model(bm.path_to_model)
+        custom_objects=None
+
+    # load model
+    model = load_model(bm.path_to_model, custom_objects=custom_objects)
 
     # check if data can be loaded blockwise to save host memory
     load_blockwise = False
     if not bm.scaling and not bm.normalize and bm.path_to_image and not np.any(region_of_interest) and \
       os.path.splitext(bm.path_to_image)[1] in ['.tif', '.tiff'] and not bm.acwe:
+
         # get image shape
         tif = TiffFile(bm.path_to_image)
         zsh = len(tif.pages)
         ysh, xsh = tif.pages[0].shape
 
         # load mask
-        if bm.separation or bm.refinement:
+        '''if bm.separation or bm.refinement:
             mask, _ = load_data(bm.mask)
             mask = mask.reshape(zsh, ysh, xsh, 1)
             mask, _, _, _ = append_ghost_areas(bm, mask)
-            mask = mask.reshape(mask.shape[:-1])
+            mask = mask.reshape(mask.shape[:-1])'''
 
         # determine new image size after appending ghost areas to make image dimensions divisible by patch size
         z_rest = bm.z_patch - (zsh % bm.z_patch)
@@ -1258,7 +1375,7 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
             xsh +=  x_rest
 
         # get Ids of patches
-        list_IDs = []
+        '''list_IDs = []
         for k in range(0, zsh-bm.z_patch+1, bm.stride_size):
             for l in range(0, ysh-bm.y_patch+1, bm.stride_size):
                 for m in range(0, xsh-bm.x_patch+1, bm.stride_size):
@@ -1268,19 +1385,18 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
                         if centerLabel>0 and np.any(patch!=centerLabel):
                             list_IDs.append(k*ysh*xsh+l*xsh+m)
                     elif bm.refinement:
-                        patch = mask[k:k+bm.z_patch, l:l+bm.y_patch, m:m+bm.x_patch]
-                        if np.any(patch==0) and np.any(patch!=0):
+                        if np.any(mask[k:k+bm.z_patch, l:l+bm.y_patch, m:m+bm.x_patch]):
                             list_IDs.append(k*ysh*xsh+l*xsh+m)
                     else:
-                        list_IDs.append(k*ysh*xsh+l*xsh+m)
+                        list_IDs.append(k*ysh*xsh+l*xsh+m)'''
 
         # make length of list divisible by batch size
-        max_i = len(list_IDs)
+        '''max_i = len(list_IDs)
         rest = bm.batch_size - (len(list_IDs) % bm.batch_size)
-        list_IDs = list_IDs + list_IDs[:rest]
+        list_IDs = list_IDs + list_IDs[:rest]'''
 
         # prediction
-        if len(list_IDs) > 400:
+        if zsh*ysh*xsh > 256**3:
             load_blockwise = True
 
     # load image data and calculate patch IDs
@@ -1288,7 +1404,7 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
 
         # load prediction data
         img, bm.img_header, z_shape, y_shape, x_shape, region_of_interest, bm.img_data = load_prediction_data(
-            bm, channels, bm.normalize, normalization_parameters, region_of_interest, bm.img_data, bm.img_header)
+            bm, channels, normalization_parameters, region_of_interest, bm.img_data, bm.img_header)
 
         # append ghost areas
         img, z_rest, y_rest, x_rest = append_ghost_areas(bm, img)
@@ -1314,6 +1430,7 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
 
     # load all patches on GPU memory
     if not load_blockwise and nb_patches < 400:
+      if rank==0:
 
         # parameters
         params = {'dim': (bm.z_patch, bm.y_patch, bm.x_patch),
@@ -1346,7 +1463,11 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
 
         # allocate final probabilities array
         if rank==0 and bm.return_probs:
-            final = np.zeros((zsh, ysh, xsh, nb_labels), dtype=np.float32)
+            if load_blockwise:
+                if not os.path.exists(path_to_probs[:-4]):
+                    os.mkdir(path_to_probs[:-4])
+            else:
+                final = np.zeros((zsh, ysh, xsh, nb_labels), dtype=np.float32)
 
         # allocate final result array
         if rank==0:
@@ -1361,27 +1482,38 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
           else:
             nprocs = ngpus
           if j % ngpus == rank:
+
             # load blockwise from TIFF
             if load_blockwise:
                 img, _, _, _, _, _, _ = load_prediction_data(bm,
-                    channels, bm.normalize, normalization_parameters,
+                    channels, normalization_parameters,
                     region_of_interest, bm.img_data, bm.img_header, load_blockwise, z)
                 img, _, _, _ = append_ghost_areas(bm, img)
 
+                # load mask block
+                if bm.separation or bm.refinement:
+                    mask = imread(bm.mask, key=range(z,min(len(tif.pages),z+bm.z_patch)))
+                    # pad zeros to make dimensions divisible by patch dimensions
+                    pad_z = bm.z_patch - mask.shape[0]
+                    pad_y = (bm.y_patch - (mask.shape[1] % bm.y_patch)) % bm.y_patch
+                    pad_x = (bm.x_patch - (mask.shape[2] % bm.x_patch)) % bm.x_patch
+                    pad_width = [(0, pad_z), (0, pad_y), (0, pad_x)]
+                    mask = np.pad(mask, pad_width, mode='constant', constant_values=0)
+
             # list of IDs
             list_IDs_block = []
 
             # get Ids of patches
+            k = 0 if load_blockwise else z
             for l in range(0, ysh-bm.y_patch+1, bm.stride_size):
                 for m in range(0, xsh-bm.x_patch+1, bm.stride_size):
                     if bm.separation:
-                        centerLabel = mask[z+bm.z_patch//2,l+bm.y_patch//2,m+bm.x_patch//2]
-                        patch = mask[z:z+bm.z_patch, l:l+bm.y_patch, m:m+bm.x_patch]
+                        centerLabel = mask[k+bm.z_patch//2,l+bm.y_patch//2,m+bm.x_patch//2]
+                        patch = mask[k:k+bm.z_patch, l:l+bm.y_patch, m:m+bm.x_patch]
                         if centerLabel>0 and np.any(patch!=centerLabel):
                             list_IDs_block.append(z*ysh*xsh+l*xsh+m)
                     elif bm.refinement:
-                        patch = mask[z:z+bm.z_patch, l:l+bm.y_patch, m:m+bm.x_patch]
-                        if np.any(patch==0) and np.any(patch!=0):
+                        if np.any(mask[k:k+bm.z_patch, l:l+bm.y_patch, m:m+bm.x_patch]):
                             list_IDs_block.append(z*ysh*xsh+l*xsh+m)
                     else:
                         list_IDs_block.append(z*ysh*xsh+l*xsh+m)
@@ -1413,10 +1545,11 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
                     # get patch
                     tmp_X = img[k:k+bm.z_patch,l:l+bm.y_patch,m:m+bm.x_patch]
                     if bm.patch_normalization:
-                        tmp_X = np.copy(tmp_X, order='C')
-                        for c in range(channels):
-                            tmp_X[:,:,:,c] -= np.mean(tmp_X[:,:,:,c])
-                            tmp_X[:,:,:,c] /= max(np.std(tmp_X[:,:,:,c]), 1e-6)
+                        tmp_X = tmp_X.copy().astype(np.float32)
+                        for ch in range(channels):
+                            mean, std = welford_mean_std(tmp_X[...,ch])
+                            tmp_X[...,ch] -= mean
+                            tmp_X[...,ch] /= max(std, 1e-6)
                     X[i] = tmp_X
 
                 # predict batch
@@ -1460,19 +1593,28 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
                         # overlap in z direction
                         if bm.stride_size < bm.z_patch:
                             if j+source>0:
-                                probs[:bm.stride_size] += overlap
+                                probs[:-bm.stride_size] += overlap
                             overlap = probs[bm.stride_size:].copy()
 
-                        # calculate result
+                        # block z dimension
                         block_z = z_indices[j+source]
-                        if j+source==len(z_indices)-1:
-                            label[block_z:block_z+bm.z_patch] = np.argmax(probs, axis=-1).astype(np.uint8)
-                            if bm.return_probs:
-                                final[block_z:block_z+bm.z_patch] = probs
+                        if j+source==len(z_indices)-1: # last block
+                            block_zsh = bm.z_patch
+                            block_z_rest = z_rest if z_rest>0 else -block_zsh
                         else:
                             block_zsh = min(bm.stride_size, bm.z_patch)
-                            label[block_z:block_z+block_zsh] = np.argmax(probs[:block_zsh], axis=-1).astype(np.uint8)
-                            if bm.return_probs:
+                            block_z_rest = -block_zsh
+
+                        # calculate result
+                        label[block_z:block_z+block_zsh] = np.argmax(probs[:block_zsh], axis=-1).astype(np.uint8)
+
+                        # return probabilities
+                        if bm.return_probs:
+                            if load_blockwise:
+                                block_output = scale_probabilities(probs[:block_zsh])
+                                block_output = block_output[:-block_z_rest,:-y_rest,:-x_rest]
+                                imwrite(path_to_probs[:-4] + f"/block-{j+source}.tif", block_output)
+                            else:
                                 final[block_z:block_z+block_zsh] = probs[:block_zsh]
                 else:
                     for i in range(bm.z_patch):
@@ -1480,7 +1622,7 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
     if rank==0:
 
         # refine mask data with result
-        if bm.refinement:
+        '''if bm.refinement:
             # loop over boundary patches
             for i, ID in enumerate(list_IDs):
                 if i < max_i:
@@ -1489,25 +1631,17 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
                     l = rest // xsh
                     m = rest % xsh
                     mask[k:k+bm.z_patch, l:l+bm.y_patch, m:m+bm.x_patch] = label[k:k+bm.z_patch, l:l+bm.y_patch, m:m+bm.x_patch]
-            label = mask
+            label = mask'''
 
-        # remove appendix
-        if bm.return_probs:
+        # remove ghost areas
+        if bm.return_probs and not load_blockwise:
             final = final[:-z_rest,:-y_rest,:-x_rest]
         label = label[:-z_rest,:-y_rest,:-x_rest]
         zsh, ysh, xsh = label.shape
 
         # return probabilities
-        if bm.return_probs:
-            counter = np.zeros((zsh, ysh, xsh, nb_labels), dtype=np.float32)
-            nb = 0
-            for k in range(0, zsh-bm.z_patch+1, bm.stride_size):
-                for l in range(0, ysh-bm.y_patch+1, bm.stride_size):
-                    for m in range(0, xsh-bm.x_patch+1, bm.stride_size):
-                        counter[k:k+bm.z_patch, l:l+bm.y_patch, m:m+bm.x_patch] += 1
-                        nb += 1
-            counter[counter==0] = 1
-            probabilities = final / counter
+        if bm.return_probs and not load_blockwise:
+            probabilities = scale_probabilities(final)
             if bm.scaling:
                 probabilities = img_resize(probabilities, z_shape, y_shape, x_shape)
             if np.any(region_of_interest):
@@ -1570,17 +1704,8 @@ def predict_segmentation(bm, region_of_interest, channels, normalization_paramet
         # save result
         if bm.path_to_image:
             save_data(bm.path_to_final, label, header=bm.header, compress=bm.compression)
-
-            # paths to optional results
-            filename, bm.extension = os.path.splitext(bm.path_to_final)
-            if bm.extension == '.gz':
-                bm.extension = '.nii.gz'
-                filename = filename[:-4]
-            path_to_cleaned = filename + '.cleaned' + bm.extension
-            path_to_filled = filename + '.filled' + bm.extension
-            path_to_cleaned_filled = filename + '.cleaned.filled' + bm.extension
-            path_to_refined = filename + '.refined' + bm.extension
-            path_to_acwe = filename + '.acwe' + bm.extension
+            if bm.return_probs and not load_blockwise:
+                imwrite(path_to_probs, probabilities)
 
         # remove outliers
         if bm.clean: