spacr 0.2.53__py3-none-any.whl → 0.2.61__py3-none-any.whl

spacr/core.py

@@ -16,7 +16,6 @@ import seaborn as sns
 import cellpose
 from skimage.measure import regionprops, label
 from skimage.transform import resize as resizescikit
-from torch.utils.data import DataLoader
 
 from skimage import measure
 from sklearn.model_selection import train_test_split
@@ -43,6 +42,16 @@ import warnings
 warnings.filterwarnings("ignore", message="3D stack used, but stitch_threshold=0 and do_3D=False, so masks are made per plane only")
 
 
+from torchvision import transforms
+from torch.utils.data import DataLoader, random_split
+from collections import defaultdict
+import os
+import random
+from PIL import Image
+from torchvision.transforms import ToTensor
+
+
+
 def analyze_plaques(folder):
     summary_data = []
     details_data = []
@@ -877,22 +886,22 @@ def annotate_results(pred_loc):
     display(df)
     return df
 
-def generate_dataset(src, file_metadata=None, experiment='TSG101_screen', sample=None):
+def generate_dataset(settings={}):
     
     from .utils import initiate_counter, add_images_to_tar
     
-    db_path = os.path.join(src, 'measurements', 'measurements.db')
-    dst = os.path.join(src, 'datasets')
+    db_path = os.path.join(settings['src'], 'measurements', 'measurements.db')
+    dst = os.path.join(settings['src'], 'datasets')
     all_paths = []
 
     # Connect to the database and retrieve the image paths
-    print(f'Reading DataBase: {db_path}')
+    print(f"Reading DataBase: {db_path}")
     try:
         with sqlite3.connect(db_path) as conn:
             cursor = conn.cursor()
-            if file_metadata:
-                if isinstance(file_metadata, str):
-                    cursor.execute("SELECT png_path FROM png_list WHERE png_path LIKE ?", (f"%{file_metadata}%",))
+            if settings['file_metadata']:
+                if isinstance(settings['file_metadata'], str):
+                    cursor.execute("SELECT png_path FROM png_list WHERE png_path LIKE ?", (f"%{settings['file_metadata']}%",))
             else:
                 cursor.execute("SELECT png_path FROM png_list")
 
@@ -909,16 +918,16 @@ def generate_dataset(src, file_metadata=None, experiment='TSG101_screen', sample
         print(f"Error: {e}")
         return
 
-    if isinstance(sample, int):
-        selected_paths = random.sample(all_paths, sample)
-        print(f'Random selection of {len(selected_paths)} paths')
+    if isinstance(settings['sample'], int):
+        selected_paths = random.sample(all_paths, settings['sample'])
+        print(f"Random selection of {len(selected_paths)} paths")
     else:
         selected_paths = all_paths
         random.shuffle(selected_paths)
-        print(f'All paths: {len(selected_paths)} paths')
+        print(f"All paths: {len(selected_paths)} paths")
 
     total_images = len(selected_paths)
-    print(f'Found {total_images} images')
+    print(f"Found {total_images} images")
 
     # Create a temp folder in dst
     temp_dir = os.path.join(dst, "temp_tars")
@@ -936,9 +945,9 @@ def generate_dataset(src, file_metadata=None, experiment='TSG101_screen', sample
         paths_chunks.append(selected_paths[start:end])
         start = end
 
-    temp_tar_files = [os.path.join(temp_dir, f'temp_{i}.tar') for i in range(num_procs)]
+    temp_tar_files = [os.path.join(temp_dir, f"temp_{i}.tar") for i in range(num_procs)]
 
-    print(f'Generating temporary tar files in {dst}')
+    print(f"Generating temporary tar files in {dst}")
 
     # Initialize shared counter and lock
     counter = Value('i', 0)
@@ -949,18 +958,18 @@ def generate_dataset(src, file_metadata=None, experiment='TSG101_screen', sample
 
     # Combine the temporary tar files into a final tar
     date_name = datetime.date.today().strftime('%y%m%d')
-    if not file_metadata is None:
-        tar_name = f'{date_name}_{experiment}_{file_metadata}.tar'
+    if not settings['file_metadata'] is None:
+        tar_name = f"{date_name}_{settings['experiment']}_{settings['file_metadata']}.tar"
     else:
-        tar_name = f'{date_name}_{experiment}.tar'
+        tar_name = f"{date_name}_{settings['experiment']}.tar"
     tar_name = os.path.join(dst, tar_name)
     if os.path.exists(tar_name):
         number = random.randint(1, 100)
-        tar_name_2 = f'{date_name}_{experiment}_{file_metadata}_{number}.tar'
-        print(f'Warning: {os.path.basename(tar_name)} exists, saving as {os.path.basename(tar_name_2)} ')
+        tar_name_2 = f"{date_name}_{settings['experiment']}_{settings['file_metadata']}_{number}.tar"
+        print(f"Warning: {os.path.basename(tar_name)} exists, saving as {os.path.basename(tar_name_2)} ")
         tar_name = os.path.join(dst, tar_name_2)
 
-    print(f'Merging temporary files')
+    print(f"Merging temporary files")
 
     with tarfile.open(tar_name, 'w') as final_tar:
         for temp_tar_path in temp_tar_files:
@@ -974,41 +983,43 @@ def generate_dataset(src, file_metadata=None, experiment='TSG101_screen', sample
     shutil.rmtree(temp_dir)
     print(f"\nSaved {total_images} images to {tar_name}")
 
-def apply_model_to_tar(tar_path, model_path, file_type='cell_png', image_size=224, batch_size=64, normalize=True, preload='images', n_jobs=10, threshold=0.5, verbose=False):
+    return tar_name
+
+def apply_model_to_tar(settings={}):
     
     from .io import TarImageDataset
     from .utils import process_vision_results, print_progress
     
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    if normalize:
+    if settings['normalize']:
         transform = transforms.Compose([
             transforms.ToTensor(),
-            transforms.CenterCrop(size=(image_size, image_size)),
+            transforms.CenterCrop(size=(settings['image_size'], settings['image_size'])),
             transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])
     else:
         transform = transforms.Compose([
             transforms.ToTensor(),
-            transforms.CenterCrop(size=(image_size, image_size))])
+            transforms.CenterCrop(size=(settings['image_size'], settings['image_size']))])
     
-    if verbose:
-        print(f'Loading model from {model_path}')
-        print(f'Loading dataset from {tar_path}')
+    if settings['verbose']:
+        print(f"Loading model from {settings['model_path']}")
+        print(f"Loading dataset from {settings['tar_path']}")
         
-    model = torch.load(model_path)
+    model = torch.load(settings['model_path'])
     
-    dataset = TarImageDataset(tar_path, transform=transform)
-    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=n_jobs, pin_memory=True)
+    dataset = TarImageDataset(settings['tar_path'], transform=transform)
+    data_loader = DataLoader(dataset, batch_size=settings['batch_size'], shuffle=True, num_workers=settings['n_jobs'], pin_memory=True)
     
-    model_name = os.path.splitext(os.path.basename(model_path))[0] 
-    dataset_name = os.path.splitext(os.path.basename(tar_path))[0]  
+    model_name = os.path.splitext(os.path.basename(settings['model_path']))[0] 
+    dataset_name = os.path.splitext(os.path.basename(settings['tar_path']))[0]  
     date_name = datetime.date.today().strftime('%y%m%d')
-    dst = os.path.dirname(tar_path)
+    dst = os.path.dirname(settings['tar_path'])
     result_loc = f'{dst}/{date_name}_{dataset_name}_{model_name}_result.csv'
 
     model.eval()
     model = model.to(device)
     
-    if verbose:
+    if settings['verbose']:
         print(model)
         print(f'Generated dataset with {len(dataset)} images')
         print(f'Generating loader from {len(data_loader)} batches')
@@ -1031,13 +1042,13 @@ def apply_model_to_tar(tar_path, model_path, file_type='cell_png', image_size=22
             stop = time.time()
             duration = stop - start
             time_ls.append(duration)
-            files_processed = batch_idx*batch_size
+            files_processed = batch_idx*settings['batch_size']
             files_to_process = len(data_loader)
-            print_progress(files_processed, files_to_process, n_jobs=n_jobs, time_ls=time_ls, batch_size=batch_size, operation_type="Tar dataset")
+            print_progress(files_processed, files_to_process, n_jobs=settings['n_jobs'], time_ls=time_ls, batch_size=settings['batch_size'], operation_type="Tar dataset")
 
     data = {'path':filenames_list, 'pred':prediction_pos_probs}
     df = pd.DataFrame(data, index=None)
-    df = process_vision_results(df, threshold)
+    df = process_vision_results(df, settings['score_threshold'])
 
     df.to_csv(result_loc, index=True, header=True, mode='w')
     torch.cuda.empty_cache()
@@ -1206,19 +1217,19 @@ def generate_dataset_from_lists(dst, class_data, classes, test_split=0.1):
         for path in train_data:
             start = time.time()
             shutil.copy(path, os.path.join(train_class_dir, os.path.basename(path)))
-            processed_files += 1
             duration = time.time() - start
             time_ls.append(duration)
             print_progress(processed_files, total_files, n_jobs=1, time_ls=None, batch_size=None, operation_type="Copying files for Train dataset")
+            processed_files += 1
 
         # Copy test files
         for path in test_data:
             start = time.time()
             shutil.copy(path, os.path.join(test_class_dir, os.path.basename(path)))
-            processed_files += 1
             duration = time.time() - start
             time_ls.append(duration)
             print_progress(processed_files, total_files, n_jobs=1, time_ls=None, batch_size=None, operation_type="Copying files for Test dataset")
+            processed_files += 1
 
     # Print summary
     for cls in classes:
@@ -1226,44 +1237,47 @@ def generate_dataset_from_lists(dst, class_data, classes, test_split=0.1):
         test_class_dir = os.path.join(dst, f'test/{cls}')
         print(f'Train class {cls}: {len(os.listdir(train_class_dir))}, Test class {cls}: {len(os.listdir(test_class_dir))}')
 
-    return
+    return os.path.join(dst, 'train'), os.path.join(dst, 'test')
 
-def generate_training_dataset(src, mode='annotation', annotation_column='test', annotated_classes=[1,2], classes=['nc','pc'], size=200, test_split=0.1, class_metadata=[['c1'],['c2']], metadata_type_by='col', channel_of_interest=3, custom_measurement=None, tables=None, png_type='cell_png'):
+def generate_training_dataset(settings):
     
     from .io import _read_and_merge_data, _read_db
     from .utils import get_paths_from_db, annotate_conditions
+    from .settings import set_generate_training_dataset_defaults
+
+    settings = set_generate_training_dataset_defaults(settings)
     
-    db_path = os.path.join(src, 'measurements','measurements.db')
-    dst = os.path.join(src, 'datasets', 'training')
+    db_path = os.path.join(settings['src'], 'measurements','measurements.db')
+    dst = os.path.join(settings['src'], 'datasets', 'training')
 
     if os.path.exists(dst):
         for i in range(1, 1000):
-            dst = os.path.join(src, 'datasets', f'training_{i}')
+            dst = os.path.join(settings['src'], 'datasets', f'training_{i}')
             if not os.path.exists(dst):
                 print(f'Creating new directory for training: {dst}')
                 break
                 
-    if mode == 'annotation':
+    if settings['dataset_mode'] == 'annotation':
         class_paths_ls_2 = []
-        class_paths_ls = training_dataset_from_annotation(db_path, dst, annotation_column, annotated_classes=annotated_classes)
+        class_paths_ls = training_dataset_from_annotation(db_path, dst, settings['annotation_column'], annotated_classes=settings['annotated_classes'])
         for class_paths in class_paths_ls:
-            class_paths_temp = random.sample(class_paths, size)
+            class_paths_temp = random.sample(class_paths, settings['size'])
             class_paths_ls_2.append(class_paths_temp)
         class_paths_ls = class_paths_ls_2
 
-    elif mode == 'metadata':
+    elif settings['dataset_mode'] == 'metadata':
         class_paths_ls = []
         class_len_ls = []
         [df] = _read_db(db_loc=db_path, tables=['png_list'])
         df['metadata_based_class'] = pd.NA
-        for i, class_ in enumerate(classes):
-            ls = class_metadata[i]
-            df.loc[df[metadata_type_by].isin(ls), 'metadata_based_class'] = class_
+        for i, class_ in enumerate(settings['classes']):
+            ls = settings['class_metadata'][i]
+            df.loc[df[settings['metadata_type_by']].isin(ls), 'metadata_based_class'] = class_
             
-        for class_ in classes:
-            if size == None:
+        for class_ in settings['classes']:
+            if settings['size'] == None:
                 c_s = []
-                for c in classes:
+                for c in settings['classes']:
                     c_s_t_df = df[df['metadata_based_class'] == c]
                     c_s.append(len(c_s_t_df))
                     print(f'Found {len(c_s_t_df)} images for class {c}')
@@ -1273,12 +1287,12 @@ def generate_training_dataset(src, mode='annotation', annotation_column='test',
             class_temp_df = df[df['metadata_based_class'] == class_]
             class_len_ls.append(len(class_temp_df))
             print(f'Found {len(class_temp_df)} images for class {class_}')
-            class_paths_temp = random.sample(class_temp_df['png_path'].tolist(), size)
+            class_paths_temp = random.sample(class_temp_df['png_path'].tolist(), settings['size'])
             class_paths_ls.append(class_paths_temp)
     
-    elif mode == 'recruitment':
+    elif settings['dataset_mode'] == 'recruitment':
         class_paths_ls = []
-        if not isinstance(tables, list):
+        if not isinstance(settings['tables'], list):
             tables = ['cell', 'nucleus', 'pathogen','cytoplasm']
         
         df, _ = _read_and_merge_data(locs=[db_path],
@@ -1290,60 +1304,58 @@ def generate_training_dataset(src, mode='annotation', annotation_column='test',
         
         print('length df 1', len(df))
         
-        df = annotate_conditions(df, cells=['HeLa'], cell_loc=None, pathogens=['pathogen'], pathogen_loc=None, treatments=classes, treatment_loc=class_metadata, types = ['col','col',metadata_type_by])
+        df = annotate_conditions(df, cells=['HeLa'], cell_loc=None, pathogens=['pathogen'], pathogen_loc=None, treatments=settings['classes'], treatment_loc=settings['class_metadata'], types = settings['metadata_type_by'])
         print('length df 2', len(df))
         [png_list_df] = _read_db(db_loc=db_path, tables=['png_list'])
 	    
-        if custom_measurement != None:
+        if settings['custom_measurement'] != None:
         
-            if not isinstance(custom_measurement, list):
+            if not isinstance(settings['custom_measurement'], list):
                  print(f'custom_measurement should be a list, add [ measurement_1,  measurement_2 ] or [ measurement ]')
                  return
         	
-            if isinstance(custom_measurement, list):
-                if len(custom_measurement) == 2:
-                    print(f'Classes will be defined by the Q1 and Q3 quantiles of recruitment ({custom_measurement[0]}/{custom_measurement[1]})')
-                    df['recruitment'] = df[f'{custom_measurement[0]}']/df[f'{custom_measurement[1]}']
-                if len(custom_measurement) == 1:
-                    print(f'Classes will be defined by the Q1 and Q3 quantiles of recruitment ({custom_measurement[0]})')
-                    df['recruitment'] = df[f'{custom_measurement[0]}']
+            if isinstance(settings['custom_measurement'], list):
+                if len(settings['custom_measurement']) == 2:
+                    print(f"Classes will be defined by the Q1 and Q3 quantiles of recruitment ({settings['custom_measurement'][0]}/{settings['custom_measurement'][1]})")
+                    df['recruitment'] = df[f"{settings['custom_measurement'][0]}']/df[f'{settings['custom_measurement'][1]}"]
+                if len(settings['custom_measurement']) == 1:
+                    print(f"Classes will be defined by the Q1 and Q3 quantiles of recruitment ({settings['custom_measurement'][0]})")
+                    df['recruitment'] = df[f"{settings['custom_measurement'][0]}"]
         else:
-            print(f'Classes will be defined by the Q1 and Q3 quantiles of recruitment (pathogen/cytoplasm for channel {channel_of_interest})')
-            df['recruitment'] = df[f'pathogen_channel_{channel_of_interest}_mean_intensity']/df[f'cytoplasm_channel_{channel_of_interest}_mean_intensity']
+            print(f"Classes will be defined by the Q1 and Q3 quantiles of recruitment (pathogen/cytoplasm for channel {settings['channel_of_interest']})")
+            df['recruitment'] = df[f"pathogen_channel_{settings['channel_of_interest']}_mean_intensity']/df[f'cytoplasm_channel_{settings['channel_of_interest']}_mean_intensity"]
 		
         q25 = df['recruitment'].quantile(0.25)
         q75 = df['recruitment'].quantile(0.75)
         df_lower = df[df['recruitment'] <= q25]
         df_upper = df[df['recruitment'] >= q75]
         
-        class_paths_lower = get_paths_from_db(df=df_lower, png_df=png_list_df, image_type=png_type)
+        class_paths_lower = get_paths_from_db(df=df_lower, png_df=png_list_df, image_type=settings['png_type'])
         
-        class_paths_lower = random.sample(class_paths_lower['png_path'].tolist(), size)
+        class_paths_lower = random.sample(class_paths_lower['png_path'].tolist(), settings['size'])
         class_paths_ls.append(class_paths_lower)
         
-        class_paths_upper = get_paths_from_db(df=df_upper, png_df=png_list_df, image_type=png_type)
-        class_paths_upper = random.sample(class_paths_upper['png_path'].tolist(), size)
+        class_paths_upper = get_paths_from_db(df=df_upper, png_df=png_list_df, image_type=settings['png_type'])
+        class_paths_upper = random.sample(class_paths_upper['png_path'].tolist(), settings['size'])
         class_paths_ls.append(class_paths_upper)
     
-    generate_dataset_from_lists(dst, class_data=class_paths_ls, classes=classes, test_split=0.1)
+    train_class_dir, test_class_dir = generate_dataset_from_lists(dst, class_data=class_paths_ls, classes=settings['classes'], test_split=settings['test_split'])
     
-    return
+    return train_class_dir, test_class_dir
 
-def generate_loaders(src, train_mode='erm', mode='train', image_size=224, batch_size=32, classes=['nc','pc'], n_jobs=None, validation_split=0.0, max_show=2, pin_memory=False, normalize=False, channels=[1, 2, 3], augment=False, verbose=False):
+def generate_loaders(src, mode='train', image_size=224, batch_size=32, classes=['nc','pc'], n_jobs=None, validation_split=0.0, pin_memory=False, normalize=False, channels=[1, 2, 3], augment=False, preload_batches=3, verbose=False):
     
     """
     Generate data loaders for training and validation/test datasets.
 
     Parameters:
     - src (str): The source directory containing the data.
-    - train_mode (str): The training mode. Options are 'erm' (Empirical Risk Minimization) or 'irm' (Invariant Risk Minimization).
     - mode (str): The mode of operation. Options are 'train' or 'test'.
     - image_size (int): The size of the input images.
     - batch_size (int): The batch size for the data loaders.
     - classes (list): The list of classes to consider.
     - n_jobs (int): The number of worker threads for data loading.
-    - validation_split (float): The fraction of data to use for validation when train_mode is 'erm'.
-    - max_show (int): The maximum number of images to show when verbose is True.
+    - validation_split (float): The fraction of data to use for validation.
     - pin_memory (bool): Whether to pin memory for faster data transfer.
     - normalize (bool): Whether to normalize the input images.
     - verbose (bool): Whether to print additional information and show images.
@@ -1352,18 +1364,10 @@ def generate_loaders(src, train_mode='erm', mode='train', image_size=224, batch_
     Returns:
     - train_loaders (list): List of data loaders for training datasets.
     - val_loaders (list): List of data loaders for validation datasets.
-    - plate_names (list): List of plate names (only applicable when train_mode is 'irm').
     """
 
-    from .io import MyDataset
-    from .plot import _imshow
-    from torchvision import transforms
-    from torch.utils.data import DataLoader, random_split
-    from collections import defaultdict
-    import os
-    import random
-    from PIL import Image
-    from torchvision.transforms import ToTensor
+    from .io import spacrDataset, spacrDataLoader
+    from .plot import _imshow_gpu
     from .utils import SelectChannels, augment_dataset
 
     chans = []
@@ -1380,12 +1384,9 @@ def generate_loaders(src, train_mode='erm', mode='train', image_size=224, batch_
     if verbose:
         print(f'Training a network on channels: {channels}')
         print(f'Channel 1: Red, Channel 2: Green, Channel 3: Blue')
-
-    plate_to_filenames = defaultdict(list)
-    plate_to_labels = defaultdict(list)
+        
     train_loaders = []
     val_loaders = []
-    plate_names = []
 
     if normalize:
         transform = transforms.Compose([
@@ -1413,157 +1414,114 @@ def generate_loaders(src, train_mode='erm', mode='train', image_size=224, batch_
         print(f'mode:{mode} is not valid, use mode = train or test')
         return
 
-    if train_mode == 'erm':
-        
-        data = MyDataset(data_dir, classes, transform=transform, shuffle=shuffle, pin_memory=pin_memory)
-        
-        if validation_split > 0:
-            train_size = int((1 - validation_split) * len(data))
-            val_size = len(data) - train_size
-            if not augment:
-                print(f'Train data:{train_size}, Validation data:{val_size}')
-            train_dataset, val_dataset = random_split(data, [train_size, val_size])
-
-            if augment:
-
-                print(f'Data before augmentation: Train: {len(train_dataset)}, Validataion:{len(val_dataset)}')
-                train_dataset = augment_dataset(train_dataset, is_grayscale=(len(channels) == 1))
-                #val_dataset = augment_dataset(val_dataset, is_grayscale=(len(channels) == 1))
-                print(f'Data after augmentation: Train: {len(train_dataset)}')#, Validataion:{len(val_dataset)}')
-
-            train_loaders = DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=n_jobs if n_jobs is not None else 0, pin_memory=pin_memory)
-            val_loaders = DataLoader(val_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=n_jobs if n_jobs is not None else 0, pin_memory=pin_memory)
-        else:
-            train_loaders = DataLoader(data, batch_size=batch_size, shuffle=shuffle, num_workers=n_jobs if n_jobs is not None else 0, pin_memory=pin_memory)
-        
-    elif train_mode == 'irm':
-        data = MyDataset(data_dir, classes, transform=transform, shuffle=shuffle, pin_memory=pin_memory)
-        
-        for filename, label in zip(data.filenames, data.labels):
-            plate = data.get_plate(filename)
-            plate_to_filenames[plate].append(filename)
-            plate_to_labels[plate].append(label)
-
-        for plate, filenames in plate_to_filenames.items():
-            labels = plate_to_labels[plate]
-            plate_data = MyDataset(data_dir, classes, specific_files=filenames, specific_labels=labels, transform=transform, shuffle=False, pin_memory=pin_memory)
-            plate_names.append(plate)
-
-            if validation_split > 0:
-                train_size = int((1 - validation_split) * len(plate_data))
-                val_size = len(plate_data) - train_size
-                if not augment:
-                    print(f'Train data:{train_size}, Validation data:{val_size}')
-                train_dataset, val_dataset = random_split(plate_data, [train_size, val_size])
-
-                if augment:
-
-                    print(f'Data before augmentation: Train: {len(train_dataset)}, Validataion:{val_dataset}')
-                    train_dataset = augment_dataset(train_dataset, is_grayscale=(len(channels) == 1))
-                    #val_dataset = augment_dataset(val_dataset, is_grayscale=(len(channels) == 1))
-                    print(f'Data after augmentation: Train: {len(train_dataset)}')#, Validataion:{len(val_dataset)}')
-
-                train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=n_jobs if n_jobs is not None else 0, pin_memory=pin_memory)
-                val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=n_jobs if n_jobs is not None else 0, pin_memory=pin_memory)
-
-                train_loaders.append(train_loader)
-                val_loaders.append(val_loader)
-            else:
-                train_loader = DataLoader(plate_data, batch_size=batch_size, shuffle=shuffle, num_workers=n_jobs if n_jobs is not None else 0, pin_memory=pin_memory)
-                train_loaders.append(train_loader)
-                val_loaders.append(None)
-    
-    else:
-        print(f'train_mode:{train_mode} is not valid, use: train_mode = irm or erm')
-        return
-    
+    data = spacrDataset(data_dir, classes, transform=transform, shuffle=shuffle, pin_memory=pin_memory)
+    num_workers = n_jobs if n_jobs is not None else 0
     
-    if train_mode == 'erm':
-        for idx, (images, labels, filenames) in enumerate(train_loaders):
-            if idx >= max_show:
-                break
-            images = images.cpu()
-            label_strings = [str(label.item()) for label in labels]
-            train_fig = _imshow(images, label_strings, nrow=20, fontsize=12)
-            if verbose:
-                plt.show()
+    if validation_split > 0:
+        train_size = int((1 - validation_split) * len(data))
+        val_size = len(data) - train_size
+        if not augment:
+            print(f'Train data:{train_size}, Validation data:{val_size}')
+        train_dataset, val_dataset = random_split(data, [train_size, val_size])
 
-    elif train_mode == 'irm':
-        for plate_name, train_loader in zip(plate_names, train_loaders):
-            print(f'Plate: {plate_name} with {len(train_loader.dataset)} images')
-            for idx, (images, labels, filenames) in enumerate(train_loader):
-                if idx >= max_show:
-                    break
-                images = images.cpu()
-                label_strings = [str(label.item()) for label in labels]
-                train_fig = _imshow(images, label_strings, nrow=20, fontsize=12)
-                if verbose:
-                    plt.show()
+        if augment:
 
-    return train_loaders, val_loaders, plate_names, train_fig
+            print(f'Data before augmentation: Train: {len(train_dataset)}, Validataion:{len(val_dataset)}')
+            train_dataset = augment_dataset(train_dataset, is_grayscale=(len(channels) == 1))
+            print(f'Data after augmentation: Train: {len(train_dataset)}')
+            
+        print(f'Generating Dataloader with {n_jobs} workers')
+        #train_loaders = spacrDataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory, persistent_workers=True, preload_batches=preload_batches)
+        #train_loaders = spacrDataLoader(val_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory, persistent_workers=True, preload_batches=preload_batches)
 
-def analyze_recruitment(src, metadata_settings={}, advanced_settings={}):
+        train_loaders = DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=1, pin_memory=pin_memory, persistent_workers=True)
+        val_loaders = DataLoader(val_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=1, pin_memory=pin_memory, persistent_workers=True)
+    else:
+        train_loaders = DataLoader(data, batch_size=batch_size, shuffle=shuffle, num_workers=1, pin_memory=pin_memory, persistent_workers=True)
+        #train_loaders = spacrDataLoader(data, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory, persistent_workers=True, preload_batches=preload_batches)
+
+    #dataset (Dataset) – dataset from which to load the data.
+    #batch_size (int, optional) – how many samples per batch to load (default: 1).
+    #shuffle (bool, optional) – set to True to have the data reshuffled at every epoch (default: False).
+    #sampler (Sampler or Iterable, optional) – defines the strategy to draw samples from the dataset. Can be any Iterable with __len__ implemented. If specified, shuffle must not be specified.
+    #batch_sampler (Sampler or Iterable, optional) – like sampler, but returns a batch of indices at a time. Mutually exclusive with batch_size, shuffle, sampler, and drop_last.
+    #num_workers (int, optional) – how many subprocesses to use for data loading. 0 means that the data will be loaded in the main process. (default: 0)
+    #collate_fn (Callable, optional) – merges a list of samples to form a mini-batch of Tensor(s). Used when using batched loading from a map-style dataset.
+    #pin_memory (bool, optional) – If True, the data loader will copy Tensors into device/CUDA pinned memory before returning them. If your data elements are a custom type, or your collate_fn returns a batch that is a custom type, see the example below.
+    #drop_last (bool, optional) – set to True to drop the last incomplete batch, if the dataset size is not divisible by the batch size. If False and the size of dataset is not divisible by the batch size, then the last batch will be smaller. (default: False)
+    #timeout (numeric, optional) – if positive, the timeout value for collecting a batch from workers. Should always be non-negative. (default: 0)
+    #worker_init_fn (Callable, optional) – If not None, this will be called on each worker subprocess with the worker id (an int in [0, num_workers - 1]) as input, after seeding and before data loading. (default: None)
+    #multiprocessing_context (str or multiprocessing.context.BaseContext, optional) – If None, the default multiprocessing context of your operating system will be used. (default: None)
+    #generator (torch.Generator, optional) – If not None, this RNG will be used by RandomSampler to generate random indexes and multiprocessing to generate base_seed for workers. (default: None)
+    #prefetch_factor (int, optional, keyword-only arg) – Number of batches loaded in advance by each worker. 2 means there will be a total of 2 * num_workers batches prefetched across all workers. (default value depends on the set value for num_workers. If value of num_workers=0 default is None. Otherwise, if value of num_workers > 0 default is 2).
+    #persistent_workers (bool, optional) – If True, the data loader will not shut down the worker processes after a dataset has been consumed once. This allows to maintain the workers Dataset instances alive. (default: False)
+    #pin_memory_device (str, optional) – the device to pin_memory to if pin_memory is True.
+
+    #images, labels, filenames = next(iter(train_loaders))
+    #images = images.cpu()
+    #label_strings = [str(label.item()) for label in labels]
+    #train_fig = _imshow_gpu(images, label_strings, nrow=20, fontsize=12)
+    #if verbose:
+    #    plt.show()
+    
+    train_fig = None
+
+    return train_loaders, val_loaders, train_fig
+
+def analyze_recruitment(settings={}):
     """
     Analyze recruitment data by grouping the DataFrame by well coordinates and plotting controls and recruitment data.
 
     Parameters:
-    src (str): The source of the recruitment data.
-    metadata_settings (dict): The settings for metadata.
-    advanced_settings (dict): The advanced settings for recruitment analysis.
+    settings (dict): settings.
 
     Returns:
     None
     """
     
     from .io import _read_and_merge_data, _results_to_csv
-    from .plot import plot_merged, _plot_controls, _plot_recruitment
-    from .utils import _object_filter, annotate_conditions, _calculate_recruitment, _group_by_well
+    from .plot import plot_image_mask_overlay, _plot_controls, _plot_recruitment
+    from .utils import _object_filter, annotate_conditions, _calculate_recruitment, _group_by_well, save_settings
     from .settings import get_analyze_recruitment_default_settings
 
-    settings = get_analyze_recruitment_default_settings(settings)
-    
-    settings_dict = {**metadata_settings, **advanced_settings}
-    settings_df = pd.DataFrame(list(settings_dict.items()), columns=['Key', 'Value'])
-    settings_csv = os.path.join(src,'settings','analyze_settings.csv')
-    os.makedirs(os.path.join(src,'settings'), exist_ok=True)
-    settings_df.to_csv(settings_csv, index=False)
+    settings = get_analyze_recruitment_default_settings(settings=settings)
+    save_settings(settings, name='recruitment')
 
     # metadata settings
-    target = metadata_settings['target']
-    cell_types = metadata_settings['cell_types']
-    cell_plate_metadata = metadata_settings['cell_plate_metadata']
-    pathogen_types = metadata_settings['pathogen_types']
-    pathogen_plate_metadata = metadata_settings['pathogen_plate_metadata']
-    treatments = metadata_settings['treatments']
-    treatment_plate_metadata = metadata_settings['treatment_plate_metadata']
-    metadata_types = metadata_settings['metadata_types']
-    channel_dims = metadata_settings['channel_dims']
-    cell_chann_dim = metadata_settings['cell_chann_dim']
-    cell_mask_dim = metadata_settings['cell_mask_dim']
-    nucleus_chann_dim = metadata_settings['nucleus_chann_dim']
-    nucleus_mask_dim = metadata_settings['nucleus_mask_dim']
-    pathogen_chann_dim = metadata_settings['pathogen_chann_dim']
-    pathogen_mask_dim = metadata_settings['pathogen_mask_dim']
-    channel_of_interest = metadata_settings['channel_of_interest']
+    src = settings['src']
+    target = settings['target']
+    cell_types = settings['cell_types']
+    cell_plate_metadata = settings['cell_plate_metadata']
+    pathogen_types = settings['pathogen_types']
+    pathogen_plate_metadata = settings['pathogen_plate_metadata']
+    treatments = settings['treatments']
+    treatment_plate_metadata = settings['treatment_plate_metadata']
+    metadata_types = settings['metadata_types']
+    channel_dims = settings['channel_dims']
+    cell_chann_dim = settings['cell_chann_dim']
+    cell_mask_dim = settings['cell_mask_dim']
+    nucleus_chann_dim = settings['nucleus_chann_dim']
+    nucleus_mask_dim = settings['nucleus_mask_dim']
+    pathogen_chann_dim = settings['pathogen_chann_dim']
+    pathogen_mask_dim = settings['pathogen_mask_dim']
+    channel_of_interest = settings['channel_of_interest']
     
     # Advanced settings
-    plot = advanced_settings['plot']
-    plot_nr = advanced_settings['plot_nr']
-    plot_control = advanced_settings['plot_control']
-    figuresize = advanced_settings['figuresize']
-    remove_background = advanced_settings['remove_background']
-    backgrounds = advanced_settings['backgrounds']
-    include_noninfected = advanced_settings['include_noninfected']
-    include_multiinfected = advanced_settings['include_multiinfected']
-    include_multinucleated = advanced_settings['include_multinucleated']
-    cells_per_well = advanced_settings['cells_per_well']
-    pathogen_size_range = advanced_settings['pathogen_size_range']
-    nucleus_size_range = advanced_settings['nucleus_size_range']
-    cell_size_range = advanced_settings['cell_size_range']
-    pathogen_intensity_range = advanced_settings['pathogen_intensity_range']
-    nucleus_intensity_range = advanced_settings['nucleus_intensity_range']
-    cell_intensity_range = advanced_settings['cell_intensity_range']
-    target_intensity_min = advanced_settings['target_intensity_min']
+    plot = settings['plot']
+    plot_nr = settings['plot_nr']
+    plot_control = settings['plot_control']
+    figuresize = settings['figuresize']
+    include_noninfected = settings['include_noninfected']
+    include_multiinfected = settings['include_multiinfected']
+    include_multinucleated = settings['include_multinucleated']
+    cells_per_well = settings['cells_per_well']
+    pathogen_size_range = settings['pathogen_size_range']
+    nucleus_size_range = settings['nucleus_size_range']
+    cell_size_range = settings['cell_size_range']
+    pathogen_intensity_range = settings['pathogen_intensity_range']
+    nucleus_intensity_range = settings['nucleus_intensity_range']
+    cell_intensity_range = settings['cell_intensity_range']
+    target_intensity_min = settings['target_intensity_min']
     
     print(f'Cell(s): {cell_types}, in {cell_plate_metadata}')
     print(f'Pathogen(s): {pathogen_types}, in {pathogen_plate_metadata}')
@@ -1581,9 +1539,6 @@ def analyze_recruitment(src, metadata_settings={}, advanced_settings={}):
         else:
             metadata_types = metadata_types
     
-    if isinstance(backgrounds, (int,float)):
-        backgrounds = [backgrounds, backgrounds, backgrounds, backgrounds]
-
     sns.color_palette("mako", as_cmap=True)
     print(f'channel:{channel_of_interest} = {target}')
     overlay_channels = channel_dims
@@ -1593,11 +1548,11 @@ def analyze_recruitment(src, metadata_settings={}, advanced_settings={}):
     db_loc = [src+'/measurements/measurements.db']
     tables = ['cell', 'nucleus', 'pathogen','cytoplasm']
     df, _ = _read_and_merge_data(db_loc, 
-                                         tables, 
-                                         verbose=True, 
-                                         include_multinucleated=include_multinucleated, 
-                                         include_multiinfected=include_multiinfected, 
-                                         include_noninfected=include_noninfected)
+                                 tables, 
+                                 verbose=True, 
+                                 include_multinucleated=include_multinucleated, 
+                                 include_multiinfected=include_multiinfected, 
+                                 include_noninfected=include_noninfected)
     
     df = annotate_conditions(df, 
                              cells=cell_types, 
@@ -1616,48 +1571,31 @@ def analyze_recruitment(src, metadata_settings={}, advanced_settings={}):
     random.shuffle(files)
 
     _max = 10**100
-
-    if cell_size_range is None and nucleus_size_range is None and pathogen_size_range is None:
-        filter_min_max = None
-    else:
-        if cell_size_range is None:
-            cell_size_range = [0,_max]
-        if nucleus_size_range is None:
-            nucleus_size_range = [0,_max]
-        if pathogen_size_range is None:
-            pathogen_size_range = [0,_max]
-
-        filter_min_max = [[cell_size_range[0],cell_size_range[1]],[nucleus_size_range[0],nucleus_size_range[1]],[pathogen_size_range[0],pathogen_size_range[1]]]
+    if cell_size_range is None:
+        cell_size_range = [0,_max]
+    if nucleus_size_range is None:
+        nucleus_size_range = [0,_max]
+    if pathogen_size_range is None:
+        pathogen_size_range = [0,_max]
 
     if plot:
-        plot_settings = {'include_noninfected':include_noninfected, 
-                         'include_multiinfected':include_multiinfected,
-                         'include_multinucleated':include_multinucleated,
-                         'remove_background':remove_background,
-                         'filter_min_max':filter_min_max,
-                         'channel_dims':channel_dims,
-                         'backgrounds':backgrounds,
-                         'cell_mask_dim':mask_dims[0],
-                         'nucleus_mask_dim':mask_dims[1],
-                         'pathogen_mask_dim':mask_dims[2],
-                         'overlay_chans':overlay_channels,
-                         'outline_thickness':3,
-                         'outline_color':'gbr',
-                         'overlay_chans':overlay_channels,
-                         'overlay':True,
-                         'normalization_percentiles':[1,99],
-                         'normalize':True,
-                         'print_object_number':True,
-                         'nr':plot_nr,
-                         'figuresize':20,
-                         'cmap':'inferno',
-                         'verbose':False}
-        
-    if os.path.exists(os.path.join(src,'merged')):
-        try:
-            plot_merged(src=os.path.join(src,'merged'), settings=plot_settings)
-        except Exception as e:
-            print(f'Failed to plot images with outlines, Error: {e}')
+        merged_path = os.path.join(src,'merged')
+        if os.path.exists(merged_path):
+            try:
+                for idx, file in enumerate(os.listdir(merged_path)):
+                    file_path = os.path.join(merged_path,file)
+                    if idx <= plot_nr:
+                        plot_image_mask_overlay(file_path, 
+                                                channel_dims,
+                                                cell_chann_dim,
+                                                nucleus_chann_dim,
+                                                pathogen_chann_dim,
+                                                figuresize=10,
+                                                normalize=True,
+                                                thickness=3,
+                                                save_pdf=True)
+            except Exception as e:
+                print(f'Failed to plot images with outlines, Error: {e}')
         
     if not cell_chann_dim is None:
         df = _object_filter(df, object_type='cell', size_range=cell_size_range, intensity_range=cell_intensity_range, mask_chans=mask_chans, mask_chan=0)
@@ -1695,14 +1633,12 @@ def preprocess_generate_masks(src, settings={}):
 
     from .io import preprocess_img_data, _load_and_concatenate_arrays
     from .plot import plot_image_mask_overlay, plot_arrays
-    from .utils import _pivot_counts_table, check_mask_folder, adjust_cell_masks, print_progress
+    from .utils import _pivot_counts_table, check_mask_folder, adjust_cell_masks, print_progress, save_settings
     from .settings import set_default_settings_preprocess_generate_masks
     
     settings = set_default_settings_preprocess_generate_masks(src, settings)
-    settings_df = pd.DataFrame(list(settings.items()), columns=['Key', 'Value'])
-    settings_csv = os.path.join(src,'settings','preprocess_generate_masks_settings.csv')
-    os.makedirs(os.path.join(src,'settings'), exist_ok=True)
-    settings_df.to_csv(settings_csv, index=False)
+    settings['src'] = src
+    save_settings(settings)
 
     if not settings['pathogen_channel'] is None:
         custom_model_ls = ['toxo_pv_lumen','toxo_cyto']
@@ -1993,7 +1929,7 @@ def generate_cellpose_masks(src, settings, object_type):
         settings_df['setting_value'] = settings_df['setting_value'].apply(str)
         display(settings_df)
         
-    figuresize=25
+    figuresize=10
     timelapse = settings['timelapse']
     
     if timelapse:
@@ -2497,7 +2433,6 @@ def ml_analysis(df, channel_of_interest=3, location_column='col', positive_contr
     df_metadata = df[[location_column]].copy()
     df, features = filter_dataframe_features(df, channel_of_interest, exclude, remove_low_variance_features, remove_highly_correlated_features, verbose)
     
-    
     if verbose:
         print(f'Found {len(features)} numerical features in the dataframe')
         print(f'Features used in training: {features}')
@@ -2642,7 +2577,6 @@ def check_index(df, elements=5, split_char='_'):
             print(idx)
         raise ValueError(f"Found {len(problematic_indices)} problematic indices that do not split into {elements} parts.")
 
-#def plate_heatmap(src, model_type='xgboost', variable='predictions', grouping='mean', min_max='allq', cmap='viridis', channel_of_interest=3, min_count=25, n_estimators=100, col_to_compare='col', pos='c2', neg='c1', exclude=None, n_repeats=10, clean=True, nr_to_plot=20, verbose=False, n_jobs=-1):
 def generate_ml_scores(src, settings):
     
     from .io import _read_and_merge_data
@@ -2680,7 +2614,7 @@ def generate_ml_scores(src, settings):
                                settings['top_features'],
                                settings['n_estimators'],
                                settings['test_size'],
-                               settings['model_type'],
+                               settings['model_type_ml'],
                                settings['n_jobs'],
                                settings['remove_low_variance_features'],
                                settings['remove_highly_correlated_features'],
@@ -2701,7 +2635,7 @@ def generate_ml_scores(src, settings):
                                 min_count=settings['minimum_cell_count'],
                                 verbose=settings['verbose'])
 
-    data_path, permutation_path, feature_importance_path, model_metricks_path, permutation_fig_path, feature_importance_fig_path, shap_fig_path, plate_heatmap_path, settings_csv = get_ml_results_paths(src, settings['model_type'], settings['channel_of_interest'])
+    data_path, permutation_path, feature_importance_path, model_metricks_path, permutation_fig_path, feature_importance_fig_path, shap_fig_path, plate_heatmap_path, settings_csv = get_ml_results_paths(src, settings['model_type_ml'], settings['channel_of_interest'])
     df, permutation_df, feature_importance_df, _, _, _, _, _, metrics_df = output
 
     settings_df.to_csv(settings_csv, index=False)
@@ -2858,6 +2792,7 @@ def generate_image_umap(settings={}):
         settings['plot_outlines'] = False
         settings['smooth_lines'] = False
 
+    print(f'Generating Image UMAP ...')
     settings_df = pd.DataFrame(list(settings.items()), columns=['Key', 'Value'])
     settings_dir = os.path.join(settings['src'][0],'settings')
     settings_csv = os.path.join(settings_dir,'embedding_settings.csv')