spacr 0.0.18__py3-none-any.whl → 0.0.21__py3-none-any.whl

spacr/utils.py

@@ -1,6 +1,8 @@
-import os, re, sqlite3, gc, torch, torchvision, time, random, string, shutil, cv2, tarfile, glob
+import sys, os, re, sqlite3, gc, torch, torchvision, time, random, string, shutil, cv2, tarfile, glob
 
 import numpy as np
+from cellpose import models as cp_models
+from cellpose import denoise
 from skimage import morphology
 from skimage.measure import label, regionprops_table, regionprops
 import skimage.measure as measure
@@ -18,6 +20,8 @@ from functools import reduce
 from IPython.display import display, clear_output
 from multiprocessing import Pool, cpu_count
 from skimage.transform import resize as resizescikit
+from skimage.morphology import dilation, square
+from skimage.measure import find_contours
 import torch.nn as nn
 import torch.nn.functional as F
 #from torchsummary import summary
@@ -29,6 +33,7 @@ from skimage.segmentation import clear_border
 import seaborn as sns
 import matplotlib.pyplot as plt
 import scipy.ndimage as ndi
+from scipy.spatial import distance
 from scipy.stats import fisher_exact
 from scipy.ndimage import binary_erosion, binary_dilation
 from skimage.exposure import rescale_intensity
@@ -36,6 +41,7 @@ from sklearn.metrics import auc, precision_recall_curve
 from sklearn.model_selection import train_test_split
 from sklearn.linear_model import Lasso, Ridge
 from sklearn.preprocessing import OneHotEncoder
+from sklearn.cluster import KMeans
 from torchvision.models.resnet import ResNet18_Weights, ResNet34_Weights, ResNet50_Weights, ResNet101_Weights, ResNet152_Weights
 
 from .logger import log_function_call
@@ -45,6 +51,54 @@ from .logger import log_function_call
 #from .plot import _plot_images_on_grid, plot_masks, _plot_histograms_and_stats, plot_resize, _plot_plates, _reg_v_plot, plot_masks
 #from .core import identify_masks
 
+
+def _gen_rgb_image(image, cahnnels):
+    rgb_image = np.take(image, cahnnels, axis=-1)
+    rgb_image = rgb_image.astype(float)
+    rgb_image -= rgb_image.min()
+    rgb_image /= rgb_image.max()
+    return rgb_image
+
+def _outline_and_overlay(image, rgb_image, mask_dims, outline_colors, outline_thickness):
+    from concurrent.futures import ThreadPoolExecutor
+    import cv2
+
+    outlines = []
+    overlayed_image = rgb_image.copy()
+
+    def process_dim(mask_dim):
+        mask = np.take(image, mask_dim, axis=-1)
+        outline = np.zeros_like(mask, dtype=np.uint8)  # Use uint8 for contour detection efficiency
+
+        # Find and draw contours
+        for j in np.unique(mask):
+        #for j in np.unique(mask)[1:]:
+            contours = find_contours(mask == j, 0.5)
+            # Convert contours for OpenCV format and draw directly to optimize
+            cv_contours = [np.flip(contour.astype(int), axis=1) for contour in contours]
+            cv2.drawContours(outline, cv_contours, -1, color=int(j), thickness=outline_thickness)
+
+        return dilation(outline, square(outline_thickness))
+
+    # Parallel processing
+    with ThreadPoolExecutor() as executor:
+        outlines = list(executor.map(process_dim, mask_dims))
+
+    # Overlay outlines onto the RGB image in a batch/vectorized manner if possible
+    for i, outline in enumerate(outlines):
+        # This part may need to be adapted to your specific use case and available functions
+        # The goal is to overlay each outline with its respective color more efficiently
+        color = outline_colors[i % len(outline_colors)]
+        for j in np.unique(outline)[1:]:
+            mask = outline == j
+            overlayed_image[mask] = color  # Direct assignment with broadcasting
+
+    # Remove mask_dims from image
+    channels_to_keep = [i for i in range(image.shape[-1]) if i not in mask_dims]
+    image = np.take(image, channels_to_keep, axis=-1)
+
+    return overlayed_image, outlines, image
+
 def _convert_cq1_well_id(well_id):
     """
     Converts a well ID to the CQ1 well format.
@@ -114,7 +168,7 @@ def _extract_filename_metadata(filenames, src, images_by_key, regular_expression
                 if metadata_type =='cq1':
                     orig_wellID = wellID
                     wellID = _convert_cq1_well_id(wellID)
-                    clear_output(wait=True)
+                    #clear_output(wait=True)
                     print(f'Converted Well ID: {orig_wellID} to {wellID}', end='\r', flush=True)
 
                 if pick_slice:
@@ -673,9 +727,6 @@ def _crop_center(img, cell_mask, new_width, new_height, normalize=(2,98)):
     img = img[start_y:end_y, start_x:end_x, :]
     return img
     
-    
-
-    
 def _masks_to_masks_stack(masks):
     """
     Convert a list of masks into a stack of masks.
@@ -692,53 +743,50 @@ def _masks_to_masks_stack(masks):
     return mask_stack
     
 def _get_diam(mag, obj):
-    if obj == 'cell':
-        if mag == 20:
-            scale = 6
-        if mag == 40:
-            scale = 4.5
-        if mag == 60:
-            scale = 3
-    elif obj == 'nucleus':
-        if mag == 20:
-            scale = 3
-        if mag == 40:
-            scale = 2
-        if mag == 60:
-            scale = 1.5
-    elif obj == 'pathogen':
-        if mag == 20:
-            scale = 1.5
-        if mag == 40:
-            scale = 1
-        if mag == 60:
-            scale = 1.25
-    elif obj == 'pathogen_nucleus':
-        if mag == 20:
-            scale = 0.25
-        if mag == 40:
-            scale = 0.2
-        if mag == 60:
-            scale = 0.2
+
+    if mag == 20:
+        if obj == 'cell':
+            diamiter = 120
+        elif obj == 'nucleus':
+            diamiter = 60
+        elif obj == 'pathogen':
+            diamiter = 30
+        else:
+            raise ValueError("Invalid magnification: Use 20, 40 or 60")
+
+    elif mag == 40:
+        if obj == 'cell':
+            diamiter = 160
+        elif obj == 'nucleus':
+            diamiter = 80
+        elif obj == 'pathogen':
+            diamiter = 40
+        else:
+            raise ValueError("Invalid magnification: Use 20, 40 or 60")
+
+    elif mag == 60:
+        if obj == 'cell':
+            diamiter = 200
+        if obj == 'nucleus':
+            diamiter = 90
+        if obj == 'pathogen':
+            diamiter = 75
+        else:
+            raise ValueError("Invalid magnification: Use 20, 40 or 60")
     else:
-        raise ValueError("Invalid object type")
-    diamiter = mag*scale
+        raise ValueError("Invalid magnification: Use 20, 40 or 60")
+    
     return diamiter
 
 def _get_object_settings(object_type, settings):
-    
     object_settings = {}
-    object_settings['refine_masks'] = False
-    object_settings['filter_size'] = False
-    object_settings['filter_dimm'] = False
-    print(object_type)
+
     object_settings['diameter'] = _get_diam(settings['magnification'], obj=object_type)
-    object_settings['remove_border_objects'] = False
-    object_settings['minimum_size'] = (object_settings['diameter']**2)/10
-    object_settings['maximum_size'] = object_settings['minimum_size']*50
+    object_settings['minimum_size'] = (object_settings['diameter']**2)/4
+    object_settings['maximum_size'] = (object_settings['diameter']**2)*10
     object_settings['merge'] = False
-    object_settings['net_avg'] = True
     object_settings['resample'] = True
+    object_settings['remove_border_objects'] = False
     object_settings['model_name'] = 'cyto'
     
     if object_type == 'cell':
@@ -746,20 +794,28 @@ def _get_object_settings(object_type, settings):
             object_settings['model_name'] = 'cyto'
         else:
             object_settings['model_name'] = 'cyto2'
-        
+        object_settings['filter_size'] = False
+        object_settings['filter_intensity'] = False
+        object_settings['restore_type'] = settings.get('cell_restore_type', None)
+
     elif object_type == 'nucleus':
         object_settings['model_name'] = 'nuclei'
+        object_settings['filter_size'] = False
+        object_settings['filter_intensity'] = False
+        object_settings['restore_type'] = settings.get('nucleus_restore_type', None)
 
     elif object_type == 'pathogen':
-        object_settings['model_name'] = 'cyto3'
-
-    elif object_type == 'pathogen_nucleus':
-        object_settings['filter_size'] = True
         object_settings['model_name'] = 'cyto'
+        object_settings['filter_size'] = True
+        object_settings['filter_intensity'] = False
+        object_settings['restore_type'] = settings.get('pathogen_restore_type', None)
+        object_settings['merge'] = settings['merge_pathogens']
         
     else:
         print(f'Object type: {object_type} not supported. Supported object types are : cell, nucleus and pathogen')
-        print(f'using settings: {object_settings}')
+
+    if settings['verbose']:
+        print(object_settings)
         
     return object_settings
     
@@ -786,6 +842,7 @@ def _pivot_counts_table(db_path):
         return df
 
     def _pivot_dataframe(df):
+
         """
         Pivot the DataFrame.
 
@@ -812,61 +869,32 @@ def _pivot_counts_table(db_path):
     pivoted_df.to_sql('pivoted_counts', conn, if_exists='replace', index=False)
     conn.close()
     
-def _get_cellpose_channels_v1(mask_channels, nucleus_chann_dim, pathogen_chann_dim, cell_chann_dim):
-    cellpose_channels = {}
-    if nucleus_chann_dim in mask_channels:
-        cellpose_channels['nucleus'] = [0, mask_channels.index(nucleus_chann_dim)]
-    if pathogen_chann_dim in mask_channels:
-        cellpose_channels['pathogen'] = [0, mask_channels.index(pathogen_chann_dim)]
-    if cell_chann_dim in mask_channels:
-        cellpose_channels['cell'] = [0, mask_channels.index(cell_chann_dim)]
-    return cellpose_channels
+def _get_cellpose_channels(src, nucleus_channel, pathogen_channel, cell_channel):
 
-def _get_cellpose_channels_v1(cell_channel, nucleus_channel, pathogen_channel):
-    # Initialize a dictionary to hold the new indices for the specified channels
-    cellpose_channels = {}
+    cell_mask_path = os.path.join(src, 'norm_channel_stack', 'cell_mask_stack')
+    nucleus_mask_path = os.path.join(src, 'norm_channel_stack', 'nucleus_mask_stack')
+    pathogen_mask_path = os.path.join(src, 'norm_channel_stack', 'pathogen_mask_stack')
 
-    # Initialize a list to keep track of the channels in their new order
-    new_channel_order = []
-
-    # Add each channel to the new order list if it is not None
-    if cell_channel is not None:
-        new_channel_order.append(('cell', cell_channel))
-    if nucleus_channel is not None:
-        new_channel_order.append(('nucleus', nucleus_channel))
-    if pathogen_channel is not None:
-        new_channel_order.append(('pathogen', pathogen_channel))
-
-    # Sort the list based on the original channel indices to maintain the original order
-    new_channel_order.sort(key=lambda x: x[1])
-    print(new_channel_order)
-    # Assign new indices based on the sorted order
-    for new_index, (channel_name, _) in enumerate(new_channel_order):
-        cellpose_channels[channel_name] = [new_index, 0]
-        
-    if cell_channel is not None and nucleus_channel is not None:
-        cellpose_channels['cell'][1] = cellpose_channels['nucleus'][0]
-        
-    return cellpose_channels
 
-def _get_cellpose_channels(nucleus_channel, pathogen_channel, cell_channel):
+    if os.path.exists(cell_mask_path) or os.path.exists(nucleus_mask_path) or os.path.exists(pathogen_mask_path):
+        if nucleus_channel is None or nucleus_channel is None or nucleus_channel is None:
+            print('Warning: Cellpose masks already exist. Unexpected behaviour when setting any object dimention to None when the object masks have been created.')
+        
     cellpose_channels = {}
     if not nucleus_channel is None:
         cellpose_channels['nucleus'] = [0,0]
         
     if not pathogen_channel is None:
         if not nucleus_channel is None:
-            cellpose_channels['pathogen'] = [0,1]
+            if not pathogen_channel is None:
+                cellpose_channels['pathogen'] = [0,2]
+            else:
+                cellpose_channels['pathogen'] = [0,1]
         else:
             cellpose_channels['pathogen'] = [0,0]
         
     if not cell_channel is None:
         if not nucleus_channel is None:
-            if not pathogen_channel is None:
-                cellpose_channels['cell'] = [0,2]
-            else:
-                cellpose_channels['cell'] = [0,1]
-        elif not pathogen_channel is None:
             cellpose_channels['cell'] = [0,1]
         else:
             cellpose_channels['cell'] = [0,0]
@@ -1027,9 +1055,6 @@ def _group_by_well(df):
     # Apply mean function to numeric columns and first to non-numeric
     df_grouped = df.groupby(['plate', 'row', 'col']).agg({**{col: np.mean for col in numeric_cols}, **{col: 'first' for col in non_numeric_cols}})
     return df_grouped
-    
-
-
 
 ###################################################
 #  Classify
@@ -1044,7 +1069,7 @@ class Cache:
         cache (OrderedDict): The cache data structure.
     """
 
-    def _init__(self, max_size):
+    def __init__(self, max_size):
         self.cache = OrderedDict()
         self.max_size = max_size
 
@@ -1075,7 +1100,7 @@ class ScaledDotProductAttention(nn.Module):
 
     """
 
-    def _init__(self, d_k):
+    def __init__(self, d_k):
         super(ScaledDotProductAttention, self).__init__()
         self.d_k = d_k
 
@@ -1106,7 +1131,7 @@ class SelfAttention(nn.Module):
         d_k (int): Dimensionality of the key and query vectors.
     """
 
-    def _init__(self, in_channels, d_k):
+    def __init__(self, in_channels, d_k):
         super(SelfAttention, self).__init__()
         self.W_q = nn.Linear(in_channels, d_k)
         self.W_k = nn.Linear(in_channels, d_k)
@@ -1130,7 +1155,7 @@ class SelfAttention(nn.Module):
         return output
 
 class ScaledDotProductAttention(nn.Module):
-    def _init__(self, d_k):
+    def __init__(self, d_k):
         """
         Initializes the ScaledDotProductAttention module.
 
@@ -1167,7 +1192,7 @@ class SelfAttention(nn.Module):
         in_channels (int): Number of input channels.
         d_k (int): Dimensionality of the key and query vectors.
     """
-    def _init__(self, in_channels, d_k):
+    def __init__(self, in_channels, d_k):
         super(SelfAttention, self).__init__()
         self.W_q = nn.Linear(in_channels, d_k)
         self.W_k = nn.Linear(in_channels, d_k)
@@ -1198,7 +1223,7 @@ class EarlyFusion(nn.Module):
     Args:
         in_channels (int): Number of input channels.
     """
-    def _init__(self, in_channels):
+    def __init__(self, in_channels):
         super(EarlyFusion, self).__init__()
         self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=1, stride=1)
         
@@ -1217,7 +1242,7 @@ class EarlyFusion(nn.Module):
 
 # Spatial Attention Mechanism
 class SpatialAttention(nn.Module):
-    def _init__(self, kernel_size=7):
+    def __init__(self, kernel_size=7):
         """
         Initializes the SpatialAttention module.
 
@@ -1262,7 +1287,7 @@ class MultiScaleBlockWithAttention(nn.Module):
         forward: Forward method for the module.
     """
 
-    def _init__(self, in_channels, out_channels):
+    def __init__(self, in_channels, out_channels):
         super(MultiScaleBlockWithAttention, self).__init__()
         self.dilated_conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, dilation=1, padding=1)
         self.spatial_attention = nn.Conv2d(out_channels, out_channels, kernel_size=1)
@@ -1295,7 +1320,7 @@ class MultiScaleBlockWithAttention(nn.Module):
 
 # Final Classifier
 class CustomCellClassifier(nn.Module):
-    def _init__(self, num_classes, pathogen_channel, use_attention, use_checkpoint, dropout_rate):
+    def __init__(self, num_classes, pathogen_channel, use_attention, use_checkpoint, dropout_rate):
         super(CustomCellClassifier, self).__init__()
         self.early_fusion = EarlyFusion(in_channels=3)
         
@@ -1324,7 +1349,7 @@ class CustomCellClassifier(nn.Module):
 
 #CNN and Transformer class, pick any Torch model.
 class TorchModel(nn.Module):
-    def _init__(self, model_name='resnet50', pretrained=True, dropout_rate=None, use_checkpoint=False):
+    def __init__(self, model_name='resnet50', pretrained=True, dropout_rate=None, use_checkpoint=False):
         super(TorchModel, self).__init__()
         self.model_name = model_name
         self.use_checkpoint = use_checkpoint
@@ -1398,7 +1423,7 @@ class TorchModel(nn.Module):
         return logits
 
 class FocalLossWithLogits(nn.Module):
-    def _init__(self, alpha=1, gamma=2):
+    def __init__(self, alpha=1, gamma=2):
         super(FocalLossWithLogits, self).__init__()
         self.alpha = alpha
         self.gamma = gamma
@@ -1410,7 +1435,7 @@ class FocalLossWithLogits(nn.Module):
         return focal_loss.mean()
     
 class ResNet(nn.Module):
-    def _init__(self, resnet_type='resnet50', dropout_rate=None, use_checkpoint=False, init_weights='imagenet'):
+    def __init__(self, resnet_type='resnet50', dropout_rate=None, use_checkpoint=False, init_weights='imagenet'):
         super(ResNet, self).__init__()
 
         resnet_map = {
@@ -1763,25 +1788,24 @@ def annotate_predictions(csv_loc):
     df['cond'] = df.apply(assign_condition, axis=1)
     return df
 
-def init_globals(counter_, lock_):
+def initiate_counter(counter_, lock_):
     global counter, lock
     counter = counter_
     lock = lock_
 
-def add_images_to_tar(args):
-    global counter, lock, total_images
-    paths_chunk, tar_path = args
+def add_images_to_tar(paths_chunk, tar_path, total_images):
     with tarfile.open(tar_path, 'w') as tar:
-        for img_path in paths_chunk:
+        for i, img_path in enumerate(paths_chunk):
             arcname = os.path.basename(img_path)
             try:
                 tar.add(img_path, arcname=arcname)
                 with lock:
                     counter.value += 1
-                    print(f"\rProcessed: {counter.value}/{total_images}", end='', flush=True)
+                    if counter.value % 100 == 0:  # Print every 100 updates
+                        progress = (counter.value / total_images) * 100
+                        print(f"Progress: {counter.value}/{total_images} ({progress:.2f}%)", end='\r', file=sys.stdout, flush=True)
             except FileNotFoundError:
                 print(f"File not found: {img_path}")
-    return tar_path
 
 def generate_fraction_map(df, gene_column, min_frequency=0.0):
     df['fraction'] = df['count']/df['well_read_sum']
@@ -2230,8 +2254,8 @@ def dice_coefficient(mask1, mask2):
 def extract_boundaries(mask, dilation_radius=1):
     binary_mask = (mask > 0).astype(np.uint8)
     struct_elem = np.ones((dilation_radius*2+1, dilation_radius*2+1))
-    dilated = binary_dilation(binary_mask, footprint=struct_elem)
-    eroded = binary_erosion(binary_mask, footprint=struct_elem)
+    dilated = morphology.binary_dilation(binary_mask, footprint=struct_elem)
+    eroded = morphology.binary_erosion(binary_mask, footprint=struct_elem)
     boundary = dilated ^ eroded
     return boundary
 
@@ -2612,24 +2636,21 @@ def _filter_object(mask, min_value):
     mask[np.isin(mask, to_remove)] = 0
     return mask
 
-def _filter_cp_masks(masks, flows, filter_size, minimum_size, maximum_size, remove_border_objects, merge, filter_dimm, batch, moving_avg_q1, moving_avg_q3, moving_count, plot, figuresize):
+def _filter_cp_masks(masks, flows, filter_size, filter_intensity, minimum_size, maximum_size, remove_border_objects, merge, batch, plot, figuresize):
+    
     """
     Filter the masks based on various criteria such as size, border objects, merging, and intensity.
 
     Args:
         masks (list): List of masks.
         flows (list): List of flows.
-        refine_masks (bool): Flag indicating whether to refine masks.
         filter_size (bool): Flag indicating whether to filter based on size.
+        filter_intensity (bool): Flag indicating whether to filter based on intensity.
         minimum_size (int): Minimum size of objects to keep.
         maximum_size (int): Maximum size of objects to keep.
         remove_border_objects (bool): Flag indicating whether to remove border objects.
         merge (bool): Flag indicating whether to merge adjacent objects.
-        filter_dimm (bool): Flag indicating whether to filter based on intensity.
         batch (ndarray): Batch of images.
-        moving_avg_q1 (float): Moving average of the first quartile of object intensities.
-        moving_avg_q3 (float): Moving average of the third quartile of object intensities.
-        moving_count (int): Count of moving averages.
         plot (bool): Flag indicating whether to plot the masks.
         figuresize (tuple): Size of the figure.
 
@@ -2641,51 +2662,66 @@ def _filter_cp_masks(masks, flows, filter_size, minimum_size, maximum_size, remo
     
     mask_stack = []
     for idx, (mask, flow, image) in enumerate(zip(masks, flows[0], batch)):
+        
         if plot and idx == 0:
             num_objects = mask_object_count(mask)
             print(f'Number of objects before filtration: {num_objects}')
             plot_masks(batch=image, masks=mask, flows=flow, cmap='inferno', figuresize=figuresize, nr=1, file_type='.npz', print_object_number=True)
 
-        if filter_size:
-            props = measure.regionprops_table(mask, properties=['label', 'area'])  # Measure properties of labeled image regions.
-            valid_labels = props['label'][np.logical_and(props['area'] > minimum_size, props['area'] < maximum_size)]  # Select labels of valid size.
-            masks[idx] = np.isin(mask, valid_labels) * mask  # Keep only valid objects.
+        if merge:
+            mask = merge_touching_objects(mask, threshold=0.66)
             if plot and idx == 0:
                 num_objects = mask_object_count(mask)
-                print(f'Number of objects after size filtration >{minimum_size} and <{maximum_size} : {num_objects}')
+                print(f'Number of objects after merging adjacent objects, : {num_objects}')
                 plot_masks(batch=image, masks=mask, flows=flow, cmap='inferno', figuresize=figuresize, nr=1, file_type='.npz', print_object_number=True)
-        if remove_border_objects:
-            mask = clear_border(mask)
+
+        if filter_size:
+            props = measure.regionprops_table(mask, properties=['label', 'area'])
+            valid_labels = props['label'][np.logical_and(props['area'] > minimum_size, props['area'] < maximum_size)] 
+            mask = np.isin(mask, valid_labels) * mask
             if plot and idx == 0:
                 num_objects = mask_object_count(mask)
-                print(f'Number of objects after removing border objects, : {num_objects}')
+                print(f'Number of objects after size filtration >{minimum_size} and <{maximum_size} : {num_objects}')
                 plot_masks(batch=image, masks=mask, flows=flow, cmap='inferno', figuresize=figuresize, nr=1, file_type='.npz', print_object_number=True)
-        if merge:
-            mask = merge_touching_objects(mask, threshold=0.25)
+
+        if filter_intensity:
+            intensity_image = image[:, :, 1]  
+            props = measure.regionprops_table(mask, intensity_image=intensity_image, properties=['label', 'mean_intensity'])
+            mean_intensities = np.array(props['mean_intensity']).reshape(-1, 1)
+
+            if mean_intensities.shape[0] >= 2:
+                kmeans = KMeans(n_clusters=2, random_state=0).fit(mean_intensities)
+                centroids = kmeans.cluster_centers_
+            
+                # Calculate the Euclidean distance between the two centroids
+                dist_between_centroids = distance.euclidean(centroids[0], centroids[1])
+                
+                # Set a threshold for the minimum distance to consider clusters distinct
+                distance_threshold = 0.25 
+                
+                if dist_between_centroids > distance_threshold:
+                    high_intensity_cluster = np.argmax(centroids)
+                    valid_labels = np.array(props['label'])[kmeans.labels_ == high_intensity_cluster]
+                    mask = np.isin(mask, valid_labels) * mask
+
             if plot and idx == 0:
                 num_objects = mask_object_count(mask)
-                print(f'Number of objects after merging adjacent objects, : {num_objects}')
+                props_after = measure.regionprops_table(mask, intensity_image=intensity_image, properties=['label', 'mean_intensity'])
+                mean_intensities_after = np.mean(np.array(props_after['mean_intensity']))
+                average_intensity_before = np.mean(mean_intensities)
+                print(f'Number of objects after potential intensity clustering: {num_objects}. Mean intensity before:{average_intensity_before:.4f}. After:{mean_intensities_after:.4f}.')
                 plot_masks(batch=image, masks=mask, flows=flow, cmap='inferno', figuresize=figuresize, nr=1, file_type='.npz', print_object_number=True)
-        if filter_dimm:
-            unique_labels = np.unique(mask)
-            if len(unique_labels) == 1 and unique_labels[0] == 0:
-                continue
-            object_intensities = [np.mean(batch[idx, :, :, 1][mask == label]) for label in unique_labels if label != 0]
-            object_q1s = [np.percentile(intensities, 25) for intensities in object_intensities if intensities.size > 0]
-            object_q3s = [np.percentile(intensities, 75) for intensities in object_intensities if intensities.size > 0]
-            if object_q1s:
-                object_q1_mean = np.mean(object_q1s)
-                object_q3_mean = np.mean(object_q3s)
-                moving_avg_q1 = (moving_avg_q1 * moving_count + object_q1_mean) / (moving_count + 1)
-                moving_avg_q3 = (moving_avg_q3 * moving_count + object_q3_mean) / (moving_count + 1)
-                moving_count += 1
-            mask = remove_intensity_objects(batch[idx, :, :, 1], mask, intensity_threshold=moving_avg_q1, mode='low')
-            mask = remove_intensity_objects(batch[idx, :, :, 1], mask, intensity_threshold=moving_avg_q3, mode='high')
+
+
+        if remove_border_objects:
+            mask = clear_border(mask)
             if plot and idx == 0:
                 num_objects = mask_object_count(mask)
-                print(f'Objects after intensity filtration > {moving_avg_q1} and <{moving_avg_q3}: {num_objects}')
+                print(f'Number of objects after removing border objects, : {num_objects}')
                 plot_masks(batch=image, masks=mask, flows=flow, cmap='inferno', figuresize=figuresize, nr=1, file_type='.npz', print_object_number=True)
+        
         mask_stack.append(mask)
+
     return mask_stack
     
 def _object_filter(df, object_type, size_range, intensity_range, mask_chans, mask_chan):
@@ -2721,6 +2757,71 @@ def _object_filter(df, object_type, size_range, intensity_range, mask_chans, mas
                 print(f'After {object_type} maximum mean intensity filter: {len(df)}')
     return df
 
-###################################################
-#  Classify
-###################################################
+def _run_test_mode(src, regex, timelapse=False):
+    if timelapse:
+        test_images = 1  # Use only 1 set for timelapse to ensure full sequence inclusion
+    else:
+        test_images = 10  # Use 10 sets for non-timelapse scenarios
+
+    test_folder_path = os.path.join(src, 'test')
+    os.makedirs(test_folder_path, exist_ok=True)
+    regular_expression = re.compile(regex)
+    
+    all_filenames = [filename for filename in os.listdir(src) if regular_expression.match(filename)]
+    print(f'Found {len(all_filenames)} files')
+    images_by_set = defaultdict(list)
+
+    for filename in all_filenames:
+        match = regular_expression.match(filename)
+        if match:
+            plate = match.group('plateID') if 'plateID' in match.groupdict() else os.path.basename(src)
+            well = match.group('wellID')
+            field = match.group('fieldID')
+            # For timelapse experiments, group images by plate, well, and field only
+            if timelapse:
+                set_identifier = (plate, well, field)
+            else:
+                # For non-timelapse, you might want to distinguish sets more granularly
+                # Here, assuming you're grouping by plate, well, and field for simplicity
+                set_identifier = (plate, well, field)
+            images_by_set[set_identifier].append(filename)
+    
+    # Prepare for random selection
+    set_identifiers = list(images_by_set.keys())
+    random.seed(42)
+    random.shuffle(set_identifiers)  # Randomize the order
+    
+    # Select a subset based on the test_images count
+    selected_sets = set_identifiers[:test_images]
+
+    # Print information about the number of sets used
+    print(f'Using {test_images} random image set(s) for test model')
+
+    # Copy files for selected sets to the test folder
+    for set_identifier in selected_sets:
+        for filename in images_by_set[set_identifier]:
+            shutil.copy(os.path.join(src, filename), test_folder_path)
+
+    return test_folder_path
+
+def _choose_model(model_name, device, object_type='cell', restore_type=None):
+    restore_list = ['denoise', 'deblur', 'upsample', None]
+    if restore_type not in restore_list:
+        print(f"Invalid restore type. Choose from {restore_list} defaulting to None")
+        restore_type = None
+
+    if restore_type == None:
+        model = cp_models.Cellpose(gpu=True, model_type=model_name, device=device)
+    else:
+        if object_type == 'nucleus':
+            restore = f'{type}_nuclei'
+            model = denoise.CellposeDenoiseModel(gpu=True, model_type="nuclei",restore_type=restore, chan2_restore=False, device=device)
+        else:
+            restore = f'{type}_cyto3'
+            if model_name =='cyto2':
+                chan2_restore = True
+            if model_name =='cyto':
+                chan2_restore = False
+            model = denoise.CellposeDenoiseModel(gpu=True, model_type="cyto3",restore_type=restore, chan2_restore=chan2_restore, device=device)
+    
+    return model