biomedisa 24.8.5__tar.gz → 24.8.7__tar.gz

{biomedisa-24.8.5 → biomedisa-24.8.7}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: biomedisa
-Version: 24.8.5
+Version: 24.8.7
 Summary: Segmentation of 3D volumetric image data
 Author: Philipp Lösel
 Author-email: philipp.loesel@anu.edu.au
@@ -10,7 +10,7 @@ Project-URL: GitHub, https://github.com/biomedisa/biomedisa
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: European Union Public Licence 1.2 (EUPL 1.2)
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.8
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 
@@ -19,6 +19,7 @@ License-File: LICENSE
 - [Overview](#overview)
 - [Hardware Requirements](#hardware-requirements)
 - [Installation (command-line based)](#installation-command-line-based)
+- [Installation (3D Slicer extension)](#installation-3d-slicer-extension)
 - [Installation (browser based)](#installation-browser-based)
 - [Download Data](#download-data)
 - [Revisions](#revisions)
@@ -33,20 +34,21 @@ License-File: LICENSE
 - [License](#license)
 
 ## Overview
-Biomedisa (https://biomedisa.info) is a free and easy-to-use open-source application for segmenting large 3D volumetric images such as CT and MRI scans, developed at [The Australian National University CTLab](https://ctlab.anu.edu.au/). Biomedisa's smart interpolation of sparsely pre-segmented slices enables accurate semi-automated segmentation by considering the complete underlying image data. Additionally, Biomedisa enables deep learning for fully automated segmentation across similar samples and structures. It is compatible with segmentation tools like Amira/Avizo, ImageJ/Fiji and 3D Slicer. If you are using Biomedisa or the data for your research please cite: Lösel, P.D. et al. [Introducing Biomedisa as an open-source online platform for biomedical image segmentation.](https://www.nature.com/articles/s41467-020-19303-w) *Nat. Commun.* **11**, 5577 (2020).
+Biomedisa (https://biomedisa.info) is a free and easy-to-use open-source application for segmenting large 3D volumetric images such as CT and MRI scans, developed at [The Australian National University CTLab](https://ctlab.anu.edu.au/). Biomedisa's smart interpolation of sparsely pre-segmented slices enables accurate semi-automated segmentation by considering the complete underlying image data. Additionally, Biomedisa enables deep learning for fully automated segmentation across similar samples and structures. It is compatible with segmentation tools like Amira/Avizo, ImageJ/Fiji, and 3D Slicer. If you are using Biomedisa or the data for your research please cite: Lösel, P.D. et al. [Introducing Biomedisa as an open-source online platform for biomedical image segmentation.](https://www.nature.com/articles/s41467-020-19303-w) *Nat. Commun.* **11**, 5577 (2020).
 
 ## Hardware Requirements
-+ One or more NVIDIA GPUs with compute capability 3.0 or higher.
++ One or more NVIDIA GPUs.
 
 ## Installation (command-line based)
-+ [Ubuntu 22.04 + Smart Interpolation](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu2204_interpolation_cli.md)
-+ [Ubuntu 22.04 + Smart Interpolation + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu2204_cuda11.8_gpu_cli.md)
-+ [Windows 10 + Smart Interpolation + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/windows10_cuda_gpu_cli.md)
++ [Ubuntu 22/24 + Smart Interpolation](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu_interpolation_cli.md)
++ [Ubuntu 22/24 + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu_deeplearning_cli.md)
++ [Ubuntu 22/24 + Smart Interpolation + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu_cli.md)
++ [Windows 10/11 + Smart Interpolation + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/windows10_cuda_gpu_cli.md)
 + [Windows (WSL) + Smart Interpolation + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/windows_wsl.md)
 
 ## Installation (3D Slicer extension)
 + [Ubuntu 22.04 + Smart Interpolation + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu2204_cuda11.8_gpu_slicer.md)
-+ [Windows 10 + Smart Interpolation](https://github.com/biomedisa/biomedisa/blob/master/README/windows10_cuda_gpu_slicer.md)
++ [Windows 10/11 + Smart Interpolation](https://github.com/biomedisa/biomedisa/blob/master/README/windows10_cuda_gpu_slicer.md)
 
 ## Installation (browser based)
 + [Ubuntu 22.04](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu2204_cuda11.8.md)
@@ -165,7 +167,7 @@ save_data('final.Head5.am', results['regular'], results['header'])
 
 #### Command-line based (prediction)
 ```
-python -m biomedisa.deeplearning C:\Users\%USERNAME%\Downloads\testing_axial_crop_pat13.nii.gz C:\Users\%USERNAME%\Downloads\heart.h5 -p
+python -m biomedisa.deeplearning C:\Users\%USERNAME%\Downloads\testing_axial_crop_pat13.nii.gz C:\Users\%USERNAME%\Downloads\heart.h5
 ```
 
 ## Mesh Generator

{biomedisa-24.8.5 → biomedisa-24.8.7}/README.md

@@ -3,6 +3,7 @@
 - [Overview](#overview)
 - [Hardware Requirements](#hardware-requirements)
 - [Installation (command-line based)](#installation-command-line-based)
+- [Installation (3D Slicer extension)](#installation-3d-slicer-extension)
 - [Installation (browser based)](#installation-browser-based)
 - [Download Data](#download-data)
 - [Revisions](#revisions)
@@ -17,20 +18,21 @@
 - [License](#license)
 
 ## Overview
-Biomedisa (https://biomedisa.info) is a free and easy-to-use open-source application for segmenting large 3D volumetric images such as CT and MRI scans, developed at [The Australian National University CTLab](https://ctlab.anu.edu.au/). Biomedisa's smart interpolation of sparsely pre-segmented slices enables accurate semi-automated segmentation by considering the complete underlying image data. Additionally, Biomedisa enables deep learning for fully automated segmentation across similar samples and structures. It is compatible with segmentation tools like Amira/Avizo, ImageJ/Fiji and 3D Slicer. If you are using Biomedisa or the data for your research please cite: Lösel, P.D. et al. [Introducing Biomedisa as an open-source online platform for biomedical image segmentation.](https://www.nature.com/articles/s41467-020-19303-w) *Nat. Commun.* **11**, 5577 (2020).
+Biomedisa (https://biomedisa.info) is a free and easy-to-use open-source application for segmenting large 3D volumetric images such as CT and MRI scans, developed at [The Australian National University CTLab](https://ctlab.anu.edu.au/). Biomedisa's smart interpolation of sparsely pre-segmented slices enables accurate semi-automated segmentation by considering the complete underlying image data. Additionally, Biomedisa enables deep learning for fully automated segmentation across similar samples and structures. It is compatible with segmentation tools like Amira/Avizo, ImageJ/Fiji, and 3D Slicer. If you are using Biomedisa or the data for your research please cite: Lösel, P.D. et al. [Introducing Biomedisa as an open-source online platform for biomedical image segmentation.](https://www.nature.com/articles/s41467-020-19303-w) *Nat. Commun.* **11**, 5577 (2020).
 
 ## Hardware Requirements
-+ One or more NVIDIA GPUs with compute capability 3.0 or higher.
++ One or more NVIDIA GPUs.
 
 ## Installation (command-line based)
-+ [Ubuntu 22.04 + Smart Interpolation](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu2204_interpolation_cli.md)
-+ [Ubuntu 22.04 + Smart Interpolation + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu2204_cuda11.8_gpu_cli.md)
-+ [Windows 10 + Smart Interpolation + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/windows10_cuda_gpu_cli.md)
++ [Ubuntu 22/24 + Smart Interpolation](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu_interpolation_cli.md)
++ [Ubuntu 22/24 + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu_deeplearning_cli.md)
++ [Ubuntu 22/24 + Smart Interpolation + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu_cli.md)
++ [Windows 10/11 + Smart Interpolation + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/windows10_cuda_gpu_cli.md)
 + [Windows (WSL) + Smart Interpolation + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/windows_wsl.md)
 
 ## Installation (3D Slicer extension)
 + [Ubuntu 22.04 + Smart Interpolation + Deep Learning](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu2204_cuda11.8_gpu_slicer.md)
-+ [Windows 10 + Smart Interpolation](https://github.com/biomedisa/biomedisa/blob/master/README/windows10_cuda_gpu_slicer.md)
++ [Windows 10/11 + Smart Interpolation](https://github.com/biomedisa/biomedisa/blob/master/README/windows10_cuda_gpu_slicer.md)
 
 ## Installation (browser based)
 + [Ubuntu 22.04](https://github.com/biomedisa/biomedisa/blob/master/README/ubuntu2204_cuda11.8.md)
@@ -149,7 +151,7 @@ save_data('final.Head5.am', results['regular'], results['header'])
 
 #### Command-line based (prediction)
 ```
-python -m biomedisa.deeplearning C:\Users\%USERNAME%\Downloads\testing_axial_crop_pat13.nii.gz C:\Users\%USERNAME%\Downloads\heart.h5 -p
+python -m biomedisa.deeplearning C:\Users\%USERNAME%\Downloads\testing_axial_crop_pat13.nii.gz C:\Users\%USERNAME%\Downloads\heart.h5
 ```
 
 ## Mesh Generator

{biomedisa-24.8.5 → biomedisa-24.8.7}/pyproject.toml

@@ -4,13 +4,13 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "biomedisa"
-version = "24.8.5"
+version = "24.8.7"
 authors = [
   { name="Philipp Lösel"}, {email="philipp.loesel@anu.edu.au" },
 ]
 description = "Segmentation of 3D volumetric image data"
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.10"
 classifiers = [
     "Programming Language :: Python :: 3",
     "License :: OSI Approved :: European Union Public Licence 1.2 (EUPL 1.2)",

{biomedisa-24.8.5 → biomedisa-24.8.7}/src/biomedisa/deeplearning.py

@@ -71,7 +71,7 @@ def deep_learning(img_data, label_data=None, val_img_data=None, val_label_data=N
     learning_rate=0.01, stride_size=32, validation_stride_size=32, validation_freq=1,
     batch_size=None, x_scale=256, y_scale=256, z_scale=256, scaling=True, early_stopping=0,
     pretrained_model=None, fine_tune=False, workers=1, cropping_epochs=50,
-    x_range=None, y_range=None, z_range=None, header=None, extension='.tif',
+    x_range=None, y_range=None, z_range=None, header=None, extension=None,
     img_header=None, img_extension='.tif', average_dice=False, django_env=False,
     path=None, success=True, return_probs=False, patch_normalization=False,
     z_patch=64, y_patch=64, x_patch=64, path_to_logfile=None, img_id=None, label_id=None,
@@ -228,9 +228,11 @@ def deep_learning(img_data, label_data=None, val_img_data=None, val_label_data=N
             bm.scaling = bool(meta['scaling'][()])
 
         # check if amira header is available in the network
-        if bm.header is None and meta.get('header') is not None:
+        if bm.extension is None and bm.header is None and meta.get('header') is not None:
             bm.header = [np.array(meta.get('header'))]
             bm.extension = '.am'
+        if bm.extension is None:
+            bm.extension = '.tif'
 
         # crop data
         crop_data = True if 'cropping_weights' in hf else False
@@ -301,6 +303,12 @@ def deep_learning(img_data, label_data=None, val_img_data=None, val_label_data=N
                 results, bm = predict_segmentation(bm, region_of_interest,
                     channels, normalization_parameters)
 
+                from mpi4py import MPI
+                comm = MPI.COMM_WORLD
+                rank = comm.Get_rank()
+                if rank>0:
+                    return 0
+
                 # results
                 if cropped_volume is not None:
                     results['cropped_volume'] = cropped_volume
@@ -488,8 +496,8 @@ if __name__ == '__main__':
                         help='Location of mask')
     parser.add_argument('-rf','--refinement', action='store_true', default=False,
                         help='Refine segmentation on full size data')
-    parser.add_argument('-ext','--extension', type=str, default='.tif',
-                        help='Save data for example as NRRD file using --extension=".nrrd"')
+    parser.add_argument('-ext','--extension', type=str, default=None,
+                        help='Save data in formats like NRRD or TIFF using --extension=".nrrd"')
     bm = parser.parse_args()
     bm.success = True
 

{biomedisa-24.8.5 → biomedisa-24.8.7}/src/biomedisa/features/DataGenerator.py

@@ -129,39 +129,25 @@ class DataGenerator(tf.keras.utils.Sequence):
     def __len__(self):
         'Denotes the number of batches per epoch'
         if len(self.list_IDs_bg) > 0:
-            len_IDs = 2 * max(len(self.list_IDs_fg), len(self.list_IDs_bg))
+            len_IDs = max(len(self.list_IDs_fg), len(self.list_IDs_bg))
+            n_batches = len_IDs // (self.batch_size // 2)
         else:
             len_IDs = len(self.list_IDs_fg)
-        n_batches = int(np.floor(len_IDs / self.batch_size))
+            n_batches = len_IDs // self.batch_size
         return n_batches
 
     def __getitem__(self, index):
         'Generate one batch of data'
-
         if len(self.list_IDs_bg) > 0:
-
-            # len IDs
-            len_IDs = max(len(self.list_IDs_fg), len(self.list_IDs_bg))
-
-            # upsample lists of indexes to the same size
-            repetitions = int(np.floor(len_IDs / len(self.list_IDs_fg))) + 1
-            upsampled_indexes_fg = np.tile(self.indexes_fg, repetitions)
-            upsampled_indexes_fg = upsampled_indexes_fg[:len_IDs]
-
-            repetitions = int(np.floor(len_IDs / len(self.list_IDs_bg))) + 1
-            upsampled_indexes_bg = np.tile(self.indexes_bg, repetitions)
-            upsampled_indexes_bg = upsampled_indexes_bg[:len_IDs]
-
             # Generate indexes of the batch
-            tmp_batch_size = int(self.batch_size / 2)
-            indexes_fg = upsampled_indexes_fg[index*tmp_batch_size:(index+1)*tmp_batch_size]
-            indexes_bg = upsampled_indexes_bg[index*tmp_batch_size:(index+1)*tmp_batch_size]
+            half_batch_size = self.batch_size // 2
+            indexes_fg = self.indexes_fg[index*half_batch_size:(index+1)*half_batch_size]
+            indexes_bg = self.indexes_bg[index*half_batch_size:(index+1)*half_batch_size]
 
             # Find list of IDs
             list_IDs_temp = [self.list_IDs_fg[k] for k in indexes_fg] + [self.list_IDs_bg[k] for k in indexes_bg]
 
         else:
-
             # Generate indexes of the batch
             indexes_fg = self.indexes_fg[index*self.batch_size:(index+1)*self.batch_size]
 
@@ -175,11 +161,22 @@ class DataGenerator(tf.keras.utils.Sequence):
 
     def on_epoch_end(self):
         'Updates indexes after each epoch'
-        self.indexes_fg = np.arange(len(self.list_IDs_fg))
-        self.indexes_bg = np.arange(len(self.list_IDs_bg))
+        if len(self.list_IDs_bg) > 0:
+            # upsample lists of indexes
+            indexes_fg = np.arange(len(self.list_IDs_fg))
+            indexes_bg = np.arange(len(self.list_IDs_bg))
+            len_IDs = max(len(self.list_IDs_fg), len(self.list_IDs_bg))
+            repetitions = len_IDs // len(self.list_IDs_fg) + 1
+            self.indexes_fg = np.tile(indexes_fg, repetitions)
+            repetitions = len_IDs // len(self.list_IDs_bg) + 1
+            self.indexes_bg = np.tile(indexes_bg, repetitions)
+        else:
+            self.indexes_fg = np.arange(len(self.list_IDs_fg))
+        # shuffle indexes
         if self.shuffle == True:
             np.random.shuffle(self.indexes_fg)
-            np.random.shuffle(self.indexes_bg)
+            if len(self.list_IDs_bg) > 0:
+                np.random.shuffle(self.indexes_bg)
 
     def __data_generation(self, list_IDs_temp):
         'Generates data containing batch_size samples' # X : (n_samples, *dim, n_channels)

{biomedisa-24.8.5 → biomedisa-24.8.7}/src/biomedisa/features/DataGeneratorCrop.py

@@ -28,22 +28,25 @@
 
 import numpy as np
 import tensorflow as tf
-from scipy.ndimage import gaussian_filter, map_coordinates
+from scipy.ndimage import gaussian_filter, map_coordinates, rotate
+import random
 
 def elastic_transform(image, alpha=100, sigma=20):
-    zsh, ysh, xsh = image.shape
+    ysh, xsh, csh = image.shape
     dx = gaussian_filter((np.random.rand(ysh, xsh) * 2 - 1) * alpha, sigma)
     dy = gaussian_filter((np.random.rand(ysh, xsh) * 2 - 1) * alpha, sigma)
     y, x = np.meshgrid(np.arange(ysh), np.arange(xsh), indexing='ij')
     indices = np.reshape(y+dy, (-1, 1)), np.reshape(x+dx, (-1, 1))
-    for k in range(zsh):
-        image[k] = map_coordinates(image[k], indices, order=1, mode='reflect').reshape(ysh, xsh)
+    for k in range(csh):
+        image[:,:,k] = map_coordinates(image[:,:,k], indices, order=0, mode='reflect').reshape(ysh, xsh)
     return image
 
 class DataGeneratorCrop(tf.keras.utils.Sequence):
     'Generates data for Keras'
-    def __init__(self, img, label, list_IDs_fg, list_IDs_bg, batch_size=32, dim=(32,32,32),
-                 n_channels=3, n_classes=2, shuffle=True):
+    def __init__(self, img, label, list_IDs_fg, list_IDs_bg, batch_size=32,
+            dim=(32,32,32), n_channels=3, n_classes=2, shuffle=True,
+            augment=(False,False,False,0), train=True):
+
         'Initialization'
         self.dim = dim
         self.list_IDs_fg = list_IDs_fg
@@ -54,6 +57,8 @@ class DataGeneratorCrop(tf.keras.utils.Sequence):
         self.n_channels = n_channels
         self.n_classes = n_classes
         self.shuffle = shuffle
+        self.augment = augment
+        self.train = train
         self.on_epoch_end()
 
     def __len__(self):
@@ -108,14 +113,37 @@ class DataGeneratorCrop(tf.keras.utils.Sequence):
     def __data_generation(self, list_IDs_temp):
         'Generates data containing batch_size samples' # X : (n_samples, *dim, n_channels)
 
+        # get augmentation parameter
+        flip_x, flip_y, flip_z, rotation = self.augment
+        elastic = False
+
         # Initialization
         X = np.empty((self.batch_size, *self.dim, self.n_channels), dtype=np.uint8)
         y = np.empty((self.batch_size,), dtype=np.int32)
 
         # Generate data
         for i, ID in enumerate(list_IDs_temp):
-            X[i,...] = self.img[ID,...]
-            y[i] = self.label[ID]
+            tmp_X = self.img[ID,...].copy()
+
+            # augmentation
+            if self.train and (any(self.augment) or elastic):
+                if flip_x and np.random.randint(2) and abs(self.label[ID])!=3:
+                    tmp_X = np.flip(tmp_X, 1)
+                if flip_y and np.random.randint(2):
+                    if abs(self.label[ID])==1:
+                        tmp_X = np.flip(tmp_X, 0)
+                    elif abs(self.label[ID])==3:
+                        tmp_X = np.flip(tmp_X, 1)
+                if flip_z and np.random.randint(2) and abs(self.label[ID])!=1:
+                    tmp_X = np.flip(tmp_X, 0)
+                if rotation:
+                    angle = random.uniform(-rotation, rotation)
+                    tmp_X = rotate(tmp_X, angle, order=0, mode='reflect', reshape=False)
+                if elastic:
+                    tmp_X = elastic_transform(tmp_X)
+
+            X[i,...] = tmp_X
+            y[i] = 0 if self.label[ID] < 0 else 1
 
         return X, y
 

{biomedisa-24.8.5 → biomedisa-24.8.7}/src/biomedisa/features/crop_helper.py

@@ -44,6 +44,7 @@ import h5py
 import tarfile
 import matplotlib.pyplot as plt
 import tempfile
+import copy
 
 class InputError(Exception):
     def __init__(self, message=None):
@@ -117,9 +118,12 @@ def load_cropping_training_data(normalize, img_list, label_list, x_scale, y_scal
                 a = label_in
             a = a.astype(np.uint8)
             a = set_labels_to_zero(a, labels_to_compute, labels_to_remove)
-            label_z = np.any(a,axis=(1,2))
-            label_y = np.any(a,axis=(0,2))
-            label_x = np.any(a,axis=(0,1))
+            label_z = np.any(a,axis=(1,2)).astype(np.int8) * 1
+            label_y = np.any(a,axis=(0,2)).astype(np.int8) * 2
+            label_x = np.any(a,axis=(0,1)).astype(np.int8) * 3
+            label_z[label_z==0] = -1
+            label_y[label_y==0] = -2
+            label_x[label_x==0] = -3
             label = np.append(label_z,label_y,axis=0)
             label = np.append(label,label_x,axis=0)
 
@@ -178,12 +182,15 @@ def load_cropping_training_data(normalize, img_list, label_list, x_scale, y_scal
                         a = label_in[k]
                     a = a.astype(np.uint8)
                     a = set_labels_to_zero(a, labels_to_compute, labels_to_remove)
-                    next_label_z = np.any(a,axis=(1,2))
-                    next_label_y = np.any(a,axis=(0,2))
-                    next_label_x = np.any(a,axis=(0,1))
-                    label = np.append(label,next_label_z,axis=0)
-                    label = np.append(label,next_label_y,axis=0)
-                    label = np.append(label,next_label_x,axis=0)
+                    label_z = np.any(a,axis=(1,2)).astype(np.int8) * 1
+                    label_y = np.any(a,axis=(0,2)).astype(np.int8) * 2
+                    label_x = np.any(a,axis=(0,1)).astype(np.int8) * 3
+                    label_z[label_z==0] = -1
+                    label_y[label_y==0] = -2
+                    label_x[label_x==0] = -3
+                    label = np.append(label,label_z,axis=0)
+                    label = np.append(label,label_y,axis=0)
+                    label = np.append(label,label_x,axis=0)
 
                     # append image
                     if any(img_list):
@@ -228,20 +235,20 @@ def load_cropping_training_data(normalize, img_list, label_list, x_scale, y_scal
     return img_rgb, label, normalization_parameters, channels
 
 def train_cropping(img, label, path_to_model, epochs, batch_size,
-                    validation_split, flip_x, flip_y, flip_z, rotate,
-                    img_val, label_val):
+    validation_split, flip_x, flip_y, flip_z, rotate,
+    img_val, label_val):
 
     # img shape
     zsh, ysh, xsh, channels = img.shape
 
     # list of IDs
-    list_IDs_fg = list(np.where(label)[0])
-    list_IDs_bg = list(np.where(label==False)[0])
+    list_IDs_fg = list(np.where(label>0)[0])
+    list_IDs_bg = list(np.where(label<0)[0])
 
     # validation data
     if np.any(img_val):
-        list_IDs_val_fg = list(np.where(label_val)[0])
-        list_IDs_val_bg = list(np.where(label_val==False)[0])
+        list_IDs_val_fg = list(np.where(label_val>0)[0])
+        list_IDs_val_bg = list(np.where(label_val<0)[0])
     elif validation_split:
         split_fg = int(len(list_IDs_fg) * validation_split)
         split_bg = int(len(list_IDs_bg) * validation_split)
@@ -287,14 +294,13 @@ def train_cropping(img, label, path_to_model, epochs, batch_size,
               'batch_size': batch_size,
               'n_classes': 2,
               'n_channels': channels,
-              'shuffle': True}
+              'shuffle': True,
+              'augment': (flip_x, flip_y, flip_z, rotate),
+              'train': True}
 
     # validation parameters
-    params_val = {'dim': (ysh, xsh),
-                  'batch_size': batch_size,
-                  'n_classes': 2,
-                  'n_channels': channels,
-                  'shuffle': False}
+    params_val = copy.deepcopy(params)
+    params_val['train'] = False
 
     # data generator
     training_generator = DataGeneratorCrop(img, label, list_IDs_fg, list_IDs_bg, **params)
@@ -497,7 +503,7 @@ def crop_volume(img, path_to_model, path_to_final, z_shape, y_shape, x_shape, ba
 # main functions
 #=====================
 
-def load_and_train(bm, x_scale=256, y_scale=256, z_scale=256):
+def load_and_train(bm, x_scale=256, y_scale=256, z_scale=256, batch_size=24):
 
     # load training data
     img, label, normalization_parameters, channels = load_cropping_training_data(bm.normalize,
@@ -512,8 +518,8 @@ def load_and_train(bm, x_scale=256, y_scale=256, z_scale=256):
             bm.only, bm.ignore, bm.val_img_data, bm.val_label_data, normalization_parameters, channels)
 
     # train cropping
-    train_cropping(img, label, bm.path_to_model, bm.cropping_epochs,
-        bm.batch_size, bm.validation_split,
+    train_cropping(img, label, bm.path_to_model,
+        bm.cropping_epochs, batch_size, bm.validation_split,
         bm.flip_x, bm.flip_y, bm.flip_z, bm.rotate,
         img_val, label_val)
 
@@ -531,7 +537,7 @@ def load_and_train(bm, x_scale=256, y_scale=256, z_scale=256):
 
     return cropping_weights, cropping_config, normalization_parameters
 
-def crop_data(bm):
+def crop_data(bm, batch_size=32):
 
     # get meta data
     hf = h5py.File(bm.path_to_model, 'r')
@@ -554,7 +560,7 @@ def crop_data(bm):
 
     # make prediction
     z_lower, z_upper, y_lower, y_upper, x_lower, x_upper, cropped_volume = crop_volume(img, bm.path_to_model,
-        bm.path_to_cropped_image, z_shape, y_shape, x_shape, bm.batch_size, bm.debug_cropping, bm.save_cropped, bm.img_data)
+        bm.path_to_cropped_image, z_shape, y_shape, x_shape, batch_size, bm.debug_cropping, bm.save_cropped, bm.img_data)
 
     # region of interest
     region_of_interest = np.array([z_lower, z_upper, y_lower, y_upper, x_lower, x_upper])