spacr 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

spacr/io.py

@@ -583,20 +583,25 @@ def _rename_and_organize_image_files(src, regex, batch_size=100, pick_slice=Fals
         None
     """
     
-    from .utils import _extract_filename_metadata
+    from .utils import _extract_filename_metadata, print_progress
     
     regular_expression = re.compile(regex)
     images_by_key = defaultdict(list)
     stack_path = os.path.join(src, 'stack')
+    files_processed = 0
     if not os.path.exists(stack_path) or (os.path.isdir(stack_path) and len(os.listdir(stack_path)) == 0):
         all_filenames = [filename for filename in os.listdir(src) if filename.endswith(img_format)]
-        print(f'All_files:{len(all_filenames)} in {src}')
-        for i in range(0, len(all_filenames), batch_size):
-            batch_filenames = all_filenames[i:i+batch_size]
+        print(f'All_files: {len(all_filenames)} in {src}')
+        time_ls = []
+
+        for idx in range(0, len(all_filenames), batch_size):
+            start = time.time()
+            batch_filenames = all_filenames[idx:idx+batch_size]
             for filename in batch_filenames:
                 images_by_key = _extract_filename_metadata(batch_filenames, src, images_by_key, regular_expression, metadata_type, pick_slice, skip_mode)
+            
             if pick_slice:
-                for key in images_by_key:
+                for i, key in enumerate(images_by_key):
                     plate, well, field, channel, mode = key
                     max_intensity_slice = max(images_by_key[key], key=lambda x: np.percentile(x, 90))
                     mip_image = Image.fromarray(max_intensity_slice)
@@ -604,13 +609,19 @@ def _rename_and_organize_image_files(src, regex, batch_size=100, pick_slice=Fals
                     os.makedirs(output_dir, exist_ok=True)
                     output_filename = f'{plate}_{well}_{field}.tif'
                     output_path = os.path.join(output_dir, output_filename)
+                    files_processed += 1
+                    stop = time.time()
+                    duration = stop - start
+                    time_ls.append(duration)
+                    files_to_process = len(all_filenames)
+                    print_progress(files_processed, files_to_process, n_jobs=1, time_ls=time_ls, batch_size=batch_size, operation_type='Preprocessing filenames')
 
                     if os.path.exists(output_path):                        
                         print(f'WARNING: A file with the same name already exists at location {output_filename}')
                     else:
                         mip_image.save(output_path)
             else:
-                for key, images in images_by_key.items():
+                for i, (key, images) in enumerate(images_by_key.items()):
                     mip = np.max(np.stack(images), axis=0)
                     mip_image = Image.fromarray(mip)
                     plate, well, field, channel = key[:4]
@@ -618,12 +629,17 @@ def _rename_and_organize_image_files(src, regex, batch_size=100, pick_slice=Fals
                     os.makedirs(output_dir, exist_ok=True)
                     output_filename = f'{plate}_{well}_{field}.tif'
                     output_path = os.path.join(output_dir, output_filename)
+                    files_processed += 1
+                    stop = time.time()
+                    duration = stop - start
+                    time_ls.append(duration)
+                    files_to_process = len(all_filenames)
+                    print_progress(files_processed, files_to_process, n_jobs=1, time_ls=time_ls, batch_size=batch_size, operation_type='Preprocessing filenames')
 
                     if os.path.exists(output_path):                        
                         print(f'WARNING: A file with the same name already exists at location {output_filename}')
                     else:
                         mip_image.save(output_path)
-
             images_by_key.clear()
 
         # Move original images to a new directory
@@ -637,6 +653,7 @@ def _rename_and_organize_image_files(src, regex, batch_size=100, pick_slice=Fals
                     print(f'WARNING: A file with the same name already exists at location {move}')
                 else:
                     shutil.move(os.path.join(src, filename), move)
+    files_processed = 0
     return
 
 def _merge_file(chan_dirs, stack_dir, file_name):
@@ -826,6 +843,7 @@ def _merge_channels(src, plot=False):
     """
 
     from .plot import plot_arrays
+    from .utils import print_progress
     
     stack_dir = os.path.join(src, 'stack')
     allowed_names = ['01', '02', '03', '04', '00', '1', '2', '3', '4', '0']
@@ -835,8 +853,7 @@ def _merge_channels(src, plot=False):
     chan_dirs.sort()
 
     print(f'List of folders in src: {chan_dirs}. Single channel folders.')
-    start_time = time.time()
-
+    
     # Assuming chan_dirs[0] is not empty and exists, adjust according to your logic
     first_dir_path = os.path.join(src, chan_dirs[0])
     dir_files = os.listdir(first_dir_path)
@@ -847,14 +864,18 @@ def _merge_channels(src, plot=False):
     print(f'Generated folder with merged arrays: {stack_dir}')
 
     if _is_dir_empty(stack_dir):
-        for file_name in dir_files:
+        time_ls = []
+        files_to_process = len(dir_files)
+        for i, file_name in enumerate(dir_files):
+            start_time = time.time()
             full_file_path = os.path.join(first_dir_path, file_name)
             if os.path.isfile(full_file_path):
                 _merge_file([os.path.join(src, d) for d in chan_dirs], stack_dir, file_name)
-
-    elapsed_time = time.time() - start_time
-    avg_time = elapsed_time / len(dir_files) if dir_files else 0
-    print(f'Average Time: {avg_time:.3f} sec, Total Elapsed Time: {elapsed_time:.3f} sec')
+            stop_time = time.time()
+            duration = stop_time - start_time
+            time_ls.append(duration)
+            files_processed = i + 1
+            print_progress(files_processed, files_to_process, n_jobs=1, time_ls=time_ls, batch_size=None, operation_type='Merging channels into npy stacks')
 
     if plot:
         plot_arrays(os.path.join(src, 'stack'))
@@ -911,6 +932,7 @@ def _mip_all(src, include_first_chan=True):
 
 #@log_function_call
 def _concatenate_channel(src, channels, randomize=True, timelapse=False, batch_size=100):
+    from .utils import print_progress
     """
     Concatenates channel data from multiple files and saves the concatenated data as numpy arrays.
 
@@ -926,6 +948,7 @@ def _concatenate_channel(src, channels, randomize=True, timelapse=False, batch_s
     """
     channels = [item for item in channels if item is not None]
     paths = []
+    time_ls = []
     index = 0
     channel_stack_loc = os.path.join(os.path.dirname(src), 'channel_stack')
     os.makedirs(channel_stack_loc, exist_ok=True)
@@ -944,8 +967,13 @@ def _concatenate_channel(src, channels, randomize=True, timelapse=False, batch_s
                     array = np.take(array, channels, axis=2)
                     stack_region.append(array)
                     filenames_region.append(os.path.basename(path))
-                #clear_output(wait=True)
-                print(f'Region {i+1}/ {len(time_stack_path_lists)}', end='\r', flush=True)
+
+                stop = time.time()
+                duration = stop - start
+                time_ls.append(duration)
+                files_processed = i+1
+                files_to_process = time_stack_path_lists
+                print_progress(files_processed, files_to_process, n_jobs=1, time_ls=time_ls, batch_size=batch_size, operation_type="Concatinating")
                 stack = np.stack(stack_region)
                 save_loc = os.path.join(channel_stack_loc, f'{name}.npz')
                 np.savez(save_loc, data=stack, filenames=filenames_region)
@@ -964,23 +992,24 @@ def _concatenate_channel(src, channels, randomize=True, timelapse=False, batch_s
         nr_files = len(paths)
         batch_index = 0  # Added this to name the output files
         stack_ls = []
-        filenames_batch = []  # to hold filenames of the current batch
+        filenames_batch = []
         for i, path in enumerate(paths):
+            start = time.time()
             array = np.load(path)
             array = np.take(array, channels, axis=2)
             stack_ls.append(array)
             filenames_batch.append(os.path.basename(path))  # store the filename
-            #clear_output(wait=True)
-            print(f'Concatenated: {i+1}/{nr_files} files')
-            #print(f'Concatenated: {i+1}/{nr_files} files', end='\r', flush=True)
-
+            stop = time.time()
+            duration = stop - start
+            time_ls.append(duration)
+            files_processed = i+1
+            files_to_process = nr_files
+            print_progress(files_processed, files_to_process, n_jobs=1, time_ls=time_ls, batch_size=batch_size, operation_type="Concatinating")
             if (i+1) % batch_size == 0 or i+1 == nr_files:
                 unique_shapes = {arr.shape[:-1] for arr in stack_ls}
                 if len(unique_shapes) > 1:
                     max_dims = np.max(np.array(list(unique_shapes)), axis=0)
-                    #clear_output(wait=True)
                     print(f'Warning: arrays with multiple shapes found in batch {i+1}. Padding arrays to max X,Y dimentions {max_dims}')
-                    #print(f'Warning: arrays with multiple shapes found in batch {i+1}. Padding arrays to max X,Y dimentions {max_dims}', end='\r', flush=True)
                     padded_stack_ls = []
                     for arr in stack_ls:
                         pad_width = [(0, max_dim - dim) for max_dim, dim in zip(max_dims, arr.shape[:-1])]
@@ -1002,6 +1031,7 @@ def _concatenate_channel(src, channels, randomize=True, timelapse=False, batch_s
 
 def _normalize_img_batch(stack, channels, save_dtype, settings):
     
+    from .utils import print_progress
     """
     Normalize the stack of images.
 
@@ -1018,7 +1048,9 @@ def _normalize_img_batch(stack, channels, save_dtype, settings):
     normalized_stack = np.zeros_like(stack, dtype=np.float32)
     
     #for channel in range(stack.shape[-1]):
-    for channel in channels:
+    time_ls = []
+    for i, channel in enumerate(channels):
+        start = time.time()
         if channel == settings['nucleus_channel']:
             background = settings['nucleus_background']
             signal_threshold = settings['nucleus_Signal_to_noise']*settings['nucleus_background']
@@ -1065,10 +1097,18 @@ def _normalize_img_batch(stack, channels, save_dtype, settings):
             arr_2d_normalized = exposure.rescale_intensity(arr_2d, in_range=(global_lower, global_upper), out_range=(0, 1))
             normalized_stack[array_index, :, :, channel] = arr_2d_normalized
 
+        stop = time.time()
+        duration = stop - start
+        time_ls.append(duration)
+        files_processed = i+1
+        files_to_process = len(channels)
+        print_progress(files_processed, files_to_process, n_jobs=1, time_ls=time_ls, batch_size=None, operation_type=f"Normalizing")
+
     return normalized_stack.astype(save_dtype)
 
 def concatenate_and_normalize(src, channels, save_dtype=np.float32, settings={}):
-    
+    from .utils import print_progress
+
     """
     Concatenates and normalizes channel data from multiple files and saves the normalized data.
 
@@ -1088,10 +1128,11 @@ def concatenate_and_normalize(src, channels, save_dtype=np.float32, settings={})
     Returns:
         str: The directory path where the concatenated and normalized channel data is saved.
     """
-    # n c p
+
     channels = [item for item in channels if item is not None]
 
     paths = []
+    time_ls = []
     output_fldr = os.path.join(os.path.dirname(src), 'norm_channel_stack')
     os.makedirs(output_fldr, exist_ok=True)
 
@@ -1099,6 +1140,7 @@ def concatenate_and_normalize(src, channels, save_dtype=np.float32, settings={})
         try:
             time_stack_path_lists = _generate_time_lists(os.listdir(src))
             for i, time_stack_list in enumerate(time_stack_path_lists):
+                start = time.time()
                 stack_region = []
                 filenames_region = []
                 for idx, file in enumerate(time_stack_list):
@@ -1107,10 +1149,14 @@ def concatenate_and_normalize(src, channels, save_dtype=np.float32, settings={})
                         parts = file.split('_')
                         name = parts[0] + '_' + parts[1] + '_' + parts[2]
                     array = np.load(path)
-                    #array = np.take(array, channels, axis=2)
                     stack_region.append(array)
                     filenames_region.append(os.path.basename(path))
-                print(f'Region {i + 1}/ {len(time_stack_path_lists)}', end='\r', flush=True)
+                stop = time.time()
+                duration = stop - start
+                time_ls.append(duration)
+                files_processed = i+1
+                files_to_process = len(time_stack_path_lists)
+                print_progress(files_processed, files_to_process, n_jobs=1, time_ls=time_ls, batch_size=None, operation_type="Concatinating")
                 stack = np.stack(stack_region)
 
                 normalized_stack = _normalize_img_batch(stack=stack,
@@ -1138,13 +1184,19 @@ def concatenate_and_normalize(src, channels, save_dtype=np.float32, settings={})
         batch_index = 0
         stack_ls = []
         filenames_batch = []
-
+        time_ls = []
+        files_processed = 0
         for i, path in enumerate(paths):
+            start = time.time()
             array = np.load(path)
-            #array = np.take(array, channels, axis=2)
             stack_ls.append(array)
             filenames_batch.append(os.path.basename(path))
-            print(f'Concatenated: {i + 1}/{nr_files} files')
+            stop = time.time()
+            duration = stop - start
+            time_ls.append(duration)
+            files_processed += 1
+            files_to_process = nr_files
+            print_progress(files_processed, files_to_process, n_jobs=1, time_ls=time_ls, batch_size=None, operation_type="Concatinating")
 
             if (i + 1) % settings['batch_size'] == 0 or i + 1 == nr_files:
                 unique_shapes = {arr.shape[:-1] for arr in stack_ls}
@@ -1219,6 +1271,7 @@ def _get_lists_for_normalization(settings):
     return backgrounds, signal_to_noise, signal_thresholds, remove_background
 
 def _normalize_stack(src, backgrounds=[100, 100, 100], remove_backgrounds=[False, False, False], lower_percentile=2, save_dtype=np.float32, signal_to_noise=[5, 5, 5], signal_thresholds=[1000, 1000, 1000]):
+    from .utils import print_progress
     """
     Normalize the stack of images.
 
@@ -1266,10 +1319,11 @@ def _normalize_stack(src, backgrounds=[100, 100, 100], remove_backgrounds=[False
                 global_upper = np.percentile(non_zero_single_channel, upper_p)
                 if global_upper >= signal_threshold:
                     break
-
+            
             # Normalize the pixels in each image to the global percentiles and then dtype.
             arr_2d_normalized = np.zeros_like(single_channel, dtype=single_channel.dtype)
             signal_to_noise_ratio_ls = []
+            time_ls = []
             for array_index in range(single_channel.shape[0]):
                 start = time.time()
                 arr_2d = single_channel[array_index, :, :]
@@ -1291,8 +1345,15 @@ def _normalize_stack(src, backgrounds=[100, 100, 100], remove_backgrounds=[False
                 duration = (stop - start) * single_channel.shape[0]
                 time_ls.append(duration)
                 average_time = np.mean(time_ls) if len(time_ls) > 0 else 0
-                print(f'Progress: files {file_index + 1}/{len(paths)}, channels:{chan_index}/{stack.shape[-1] - 1}, arrays:{array_index + 1}/{single_channel.shape[0]}, Signal:{upper:.1f}, noise:{lower:.1f}, Signal-to-noise:{average_stnr:.1f}, Time/channel:{average_time:.2f}sec')
-
+                print(f'channels:{chan_index}/{stack.shape[-1] - 1}, arrays:{array_index + 1}/{single_channel.shape[0]}, Signal:{upper:.1f}, noise:{lower:.1f}, Signal-to-noise:{average_stnr:.1f}, Time/channel:{average_time:.2f}sec')
+
+                #stop = time.time()
+                #duration = stop - start
+                #time_ls.append(duration)
+                #files_processed = file_index + 1
+                #files_to_process = len(paths)
+                #print_progress(files_processed, files_to_process, n_jobs=1, time_ls=time_ls, batch_size=None, operation_type="Normalizing")
+                
             normalized_stack[:, :, :, channel] = arr_2d_normalized
         
         save_loc = os.path.join(output_fldr, f'{name}_norm_stack.npz')
@@ -1303,6 +1364,7 @@ def _normalize_stack(src, backgrounds=[100, 100, 100], remove_backgrounds=[False
     return print(f'Saved stacks: {output_fldr}')
 
 def _normalize_timelapse(src, lower_percentile=2, save_dtype=np.float32):
+    from .utils import print_progress
     """
     Normalize the timelapse data by rescaling the intensity values based on percentiles.
 
@@ -1326,8 +1388,9 @@ def _normalize_timelapse(src, lower_percentile=2, save_dtype=np.float32):
 
         for chan_index in range(stack.shape[-1]):
             single_channel = stack[:, :, :, chan_index]
-
+            time_ls = []
             for array_index in range(single_channel.shape[0]):
+                start = time.time()
                 arr_2d = single_channel[array_index]
                 # Calculate the 1% and 98% percentiles for this specific image
                 q_low = np.percentile(arr_2d[arr_2d != 0], lower_percentile)
@@ -1337,7 +1400,14 @@ def _normalize_timelapse(src, lower_percentile=2, save_dtype=np.float32):
                 arr_2d_rescaled = exposure.rescale_intensity(arr_2d, in_range=(q_low, q_high), out_range='dtype')
                 normalized_stack[array_index, :, :, chan_index] = arr_2d_rescaled
 
-                print(f'Progress: files {file_index+1}/{len(paths)}, channels:{chan_index+1}/{stack.shape[-1]}, arrays:{array_index+1}/{single_channel.shape[0]}', end='\r')
+                print(f'channels:{chan_index+1}/{stack.shape[-1]}, arrays:{array_index+1}/{single_channel.shape[0]}', end='\r')
+
+                #stop = time.time()
+                #duration = stop - start
+                #time_ls.append(duration)
+                #files_processed = file_index+1
+                #files_to_process = len(paths)
+                #print_progress(files_processed, files_to_process, n_jobs=1, time_ls=time_ls, batch_size=None, operation_type="Normalizing")
 
         save_loc = os.path.join(output_fldr, f'{name}_norm_timelapse.npz')
         np.savez(save_loc, data=normalized_stack, filenames=filenames)
@@ -1547,8 +1617,8 @@ def preprocess_img_data(settings):
                               save_dtype=np.float32,
                               settings=settings)
 
-    if plot:
-        _plot_4D_arrays(src+'/norm_channel_stack', nr_npz=1, nr=nr)
+    #if plot:
+    #    _plot_4D_arrays(src+'/norm_channel_stack', nr_npz=1, nr=nr)
 
     return settings, src
     
@@ -1602,6 +1672,7 @@ def _get_avg_object_size(masks):
         return 0  # Return 0 if no objects are found
     
 def _save_figure(fig, src, text, dpi=300, i=1, all_folders=1):
+    from .utils import print_progress
     """
     Save a figure to a specified location.
 
@@ -1611,6 +1682,7 @@ def _save_figure(fig, src, text, dpi=300, i=1, all_folders=1):
     text (str): The text to be included in the figure name.
     dpi (int, optional): The resolution of the saved figure. Defaults to 300.
     """
+
     save_folder = os.path.dirname(src)
     obj_type = os.path.basename(src)
     name = os.path.basename(save_folder)
@@ -1619,10 +1691,11 @@ def _save_figure(fig, src, text, dpi=300, i=1, all_folders=1):
     fig_name = f'{obj_type}_{name}_{text}.pdf'        
     save_location = os.path.join(save_folder, fig_name)
     fig.savefig(save_location, bbox_inches='tight', dpi=dpi)
-    clear_output(wait=True)
-    print(f'Progress: {i}/{all_folders}, Saved single cell figure: {os.path.basename(save_location)}')
-    #print(f'Progress: {i}/{all_folders}, Saved single cell figure: {os.path.basename(save_location)}', end='\r', flush=True)
-    # Close and delete the figure to free up memory
+
+    files_processed = i
+    files_to_process = all_folders
+    print_progress(files_processed, files_to_process, n_jobs=1, time_ls=None, batch_size=None, operation_type="Saving Figures")
+    print(f'Saved single cell figure: {os.path.basename(save_location)}')
     plt.close(fig)
     del fig
     gc.collect()
@@ -1842,6 +1915,9 @@ def _create_database(db_path):
             conn.close() 
     
 def _load_and_concatenate_arrays(src, channels, cell_chann_dim, nucleus_chann_dim, pathogen_chann_dim):
+
+    from .utils import print_progress
+
     """
     Load and concatenate arrays from multiple folders.
 
@@ -1870,9 +1946,10 @@ def _load_and_concatenate_arrays(src, channels, cell_chann_dim, nucleus_chann_di
 
     count=0
     all_imgs = len(os.listdir(reference_folder))
-
+    time_ls = []
     # Iterate through each file in the reference folder
-    for filename in os.listdir(reference_folder):
+    for idx, filename in enumerate(os.listdir(reference_folder)):
+        start = time.time()
         stack_ls = []
         if filename.endswith('.npy'):
             count += 1
@@ -1929,9 +2006,13 @@ def _load_and_concatenate_arrays(src, channels, cell_chann_dim, nucleus_chann_di
                     output_path = os.path.join(output_folder, filename)
                     np.save(output_path, stack)
         
-        #clear_output(wait=True)
-        print(f'Files merged: {count}/{all_imgs}')
-        #print(f'Files merged: {count}/{all_imgs}', end='\r', flush=True)
+        stop = time.time()
+        duration = stop - start
+        time_ls.append(duration)
+        files_processed = idx+1
+        files_to_process = all_imgs
+        print_progress(files_processed, files_to_process, n_jobs=1, time_ls=time_ls, batch_size=None, operation_type="Merging Arrays")
+
     return
     
 def _read_db(db_loc, tables):
@@ -2211,7 +2292,7 @@ def _save_model(model, model_type, results_df, dst, epoch, epochs, intermedeate_
     
     def save_model_at_threshold(threshold, epoch, suffix=""):
         percentile = str(threshold * 100)
-        print(f'\rfound: {percentile}% accurate model', end='\r', flush=True)
+        print(f'\rfound: {percentile}% accurate model')#, end='\r', flush=True)
         torch.save(model, f'{dst}/{model_type}_epoch_{str(epoch)}{suffix}_acc_{percentile}_channels_{channels_str}.pth')
 
     if epoch % 100 == 0 or epoch == epochs:
@@ -2466,7 +2547,6 @@ def _read_mask(mask_path):
         mask = img_as_uint(mask)
     return mask
 
-
 def convert_numpy_to_tiff(folder_path, limit=None):
     """
     Converts all numpy files in a folder to TIFF format and saves them in a subdirectory 'tiff'.
@@ -2545,5 +2625,5 @@ def generate_cellpose_train_test(src, test_split=0.1):
             new_mask_path = os.path.join(dst_mask, filename)            
             shutil.copy(img_path, new_img_path)
             shutil.copy(mask_path, new_mask_path)
-            print(f'Copied {idx+1}/{len(ls)} images to {_type} set', end='\r', flush=True)
+            print(f'Copied {idx+1}/{len(ls)} images to {_type} set')#, end='\r', flush=True)