spacr 0.2.81__py3-none-any.whl → 0.3.0__py3-none-any.whl

spacr/__init__.py

@@ -23,9 +23,9 @@ from . import app_measure
 from . import app_classify
 from . import app_sequencing
 from . import app_umap
+from . import mediar
 from . import logger
 
-
 __all__ = [
     "core",
     "io",
@@ -48,6 +48,7 @@ __all__ = [
     "app_classify",
     "app_sequencing",
     "app_umap",
+    "mediar",
     "logger"
 ]
 

spacr/core.py

@@ -50,8 +50,6 @@ import random
 from PIL import Image
 from torchvision.transforms import ToTensor
 
-
-
 def analyze_plaques(folder):
     summary_data = []
     details_data = []
@@ -1674,7 +1672,10 @@ def preprocess_generate_masks(src, settings={}):
             time_ls=[]
             if check_mask_folder(src, 'cell_mask_stack'):
                 start = time.time()
-                generate_cellpose_masks(mask_src, settings, 'cell')
+                if settings['segmantation_model'] == 'cellpose':
+                    generate_cellpose_masks(mask_src, settings, 'cell')
+                elif settings['segmantation_model'] == 'mediar':
+                    generate_mediar_masks(mask_src, settings, 'cell')
                 stop = time.time()
                 duration = (stop - start)
                 time_ls.append(duration)
@@ -1685,7 +1686,10 @@ def preprocess_generate_masks(src, settings={}):
             time_ls=[]
             if check_mask_folder(src, 'nucleus_mask_stack'):
                 start = time.time()
-                generate_cellpose_masks(mask_src, settings, 'nucleus')
+                if settings['segmantation_model'] == 'cellpose':
+                    generate_cellpose_masks(mask_src, settings, 'nucleus')
+                elif settings['segmantation_model'] == 'mediar':
+                    generate_mediar_masks(mask_src, settings, 'nucleus')
                 stop = time.time()
                 duration = (stop - start)
                 time_ls.append(duration)
@@ -1696,7 +1700,10 @@ def preprocess_generate_masks(src, settings={}):
             time_ls=[]
             if check_mask_folder(src, 'pathogen_mask_stack'):
                 start = time.time()
-                generate_cellpose_masks(mask_src, settings, 'pathogen')
+                if settings['segmantation_model'] == 'cellpose':
+                    generate_cellpose_masks(mask_src, settings, 'pathogen')
+                elif settings['segmantation_model'] == 'mediar':
+                    generate_mediar_masks(mask_src, settings, 'pathogen')
                 stop = time.time()
                 duration = (stop - start)
                 time_ls.append(duration)
@@ -1898,7 +1905,7 @@ def all_elements_match(list1, list2):
     # Check if all elements in list1 are in list2
     return all(element in list2 for element in list1)
 
-def prepare_batch_for_cellpose(batch):
+def prepare_batch_for_segmentation(batch):
     # Ensure the batch is of dtype float32
     if batch.dtype != np.float32:
         batch = batch.astype(np.float32)
@@ -2021,7 +2028,7 @@ def generate_cellpose_masks(src, settings, object_type):
             if batch.size == 0:
                 continue
             
-            batch = prepare_batch_for_cellpose(batch)
+            batch = prepare_batch_for_segmentation(batch)
 
             if timelapse:
                 movie_path = os.path.join(os.path.dirname(src), 'movies')
@@ -2180,7 +2187,7 @@ def generate_masks_from_imgs(src, model, model_name, batch_size, diameter, cellp
         image_files = all_image_files[i:i+batch_size]
 
         if normalize:
-            images, _, image_names, _ = _load_normalized_images_and_labels(image_files, None, channels, percentiles,  circular, invert, plot, remove_background, background, Signal_to_noise)
+            images, _, image_names, _, orig_dims = _load_normalized_images_and_labels(image_files, None, channels, percentiles,  circular, invert, plot, remove_background, background, Signal_to_noise, target_height, target_width)
             images = [np.squeeze(img) if img.shape[-1] == 1 else img for img in images]
             orig_dims = [(image.shape[0], image.shape[1]) for image in images]
         else:
@@ -2300,7 +2307,7 @@ def compare_cellpose_masks(src, verbose=False, processes=None, save=True):
     from .io import _read_mask
 
     dirs = [os.path.join(src, d) for d in os.listdir(src) if os.path.isdir(os.path.join(src, d)) and d != 'results']
-    dirs.sort()  # Optional: sort directories if needed
+    dirs.sort()
     conditions = [os.path.basename(d) for d in dirs]
 
     # Get common files in all directories
@@ -2316,7 +2323,7 @@ def compare_cellpose_masks(src, verbose=False, processes=None, save=True):
 
     # Filter out None results (from skipped files)
     results = [res for res in results if res is not None]
-    #print(results)
+    print(results)
     if verbose:
         for result in results:
             filename = result['filename']
@@ -3102,4 +3109,92 @@ def reducer_hyperparameter_search(settings={}, reduction_params=None, dbscan_par
     else:
         plt.show()
 
-    return
+    return
+
+def generate_mediar_masks(src, settings, object_type):
+    """
+    Generates masks using the MEDIARPredictor.
+    
+    :param src: Source folder containing images or npz files.
+    :param settings: Dictionary of settings for generating masks.
+    :param object_type: Type of object to detect (e.g., 'cell', 'nucleus', etc.).
+    """
+    from .mediar import MEDIARPredictor
+    from .io import _create_database, _save_object_counts_to_database
+    from .plot import plot_masks
+    from .settings import set_default_settings_preprocess_generate_masks, _get_object_settings
+
+    # Clear CUDA cache and check if CUDA is available
+    gc.collect()
+    if not torch.cuda.is_available():
+        print(f'Torch CUDA is not available, using CPU')
+
+    # Preprocess settings
+    settings = set_default_settings_preprocess_generate_masks(src, settings)
+
+    if settings['verbose']:
+        settings_df = pd.DataFrame(list(settings.items()), columns=['setting_key', 'setting_value'])
+        settings_df['setting_value'] = settings_df['setting_value'].apply(str)
+        display(settings_df)
+
+    figuresize = 10
+    timelapse = settings['timelapse']
+    batch_size = settings['batch_size']
+
+    # Get object settings and initialize MEDIARPredictor
+    mediar_predictor = MEDIARPredictor(input_path=None, output_path=None, normalize=settings['normalize'], use_tta=False)
+
+    # Paths to input npz files
+    paths = [os.path.join(src, file) for file in os.listdir(src) if file.endswith('.npz')]
+
+    # Initialize a database for saving measurements
+    count_loc = os.path.join(os.path.dirname(src), 'measurements', 'measurements.db')
+    os.makedirs(os.path.dirname(src) + '/measurements', exist_ok=True)
+    _create_database(count_loc)
+
+    for file_index, path in enumerate(paths):
+        name = os.path.basename(path)
+        name, ext = os.path.splitext(name)
+        output_folder = os.path.join(os.path.dirname(path), f'{object_type}_mask_stack')
+        os.makedirs(output_folder, exist_ok=True)
+
+        with np.load(path) as data:
+            stack = data['data']
+            filenames = data['filenames']
+
+            for i, filename in enumerate(filenames):
+                output_path = os.path.join(output_folder, filename)
+                if os.path.exists(output_path):
+                    print(f"File {filename} already exists. Skipping...")
+                    continue
+
+        # Process each batch of images in the stack
+        for i in range(0, stack.shape[0], batch_size):
+            batch = stack[i: i + batch_size]
+            batch_filenames = filenames[i: i + batch_size]
+
+            # Prepare batch for MEDIARPredictor (optional)
+            batch = prepare_batch_for_segmentation(batch)
+
+            # Predict masks using MEDIARPredictor
+            predicted_masks = mediar_predictor.predict_batch(batch)
+
+            # Save predicted masks
+            for j, mask in enumerate(predicted_masks):
+                output_filename = os.path.join(output_folder, batch_filenames[j])
+                mask = mask.astype(np.uint16)
+                np.save(output_filename, mask)
+
+            # Optional: Plot the masks
+            if settings['plot']:
+                for idx, mask in enumerate(predicted_masks):
+                    plot_masks(batch[idx], mask, cmap='inferno', figuresize=figuresize)
+
+            # Save object counts to database
+            _save_object_counts_to_database(predicted_masks, object_type, batch_filenames, count_loc)
+
+        # Clear CUDA cache after each file
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    print("Mask generation completed.")

spacr/gui.py

@@ -32,7 +32,7 @@ class MainApp(tk.Tk):
         # Set the window size to the dimensions of the monitor where it is located
         self.geometry(f"{width}x{height}")
         self.title("SpaCr GUI Collection")
-        self.configure(bg='#333333')  # Set window background to dark gray
+        self.configure(bg='#333333')
 
         style = ttk.Style()
         self.color_settings = set_dark_style(style, parent_frame=self)
@@ -55,7 +55,8 @@ class MainApp(tk.Tk):
             "Cellpose All": (lambda frame: initiate_root(self, 'cellpose_all'), "Run Cellpose on all images."),
             "Map Barcodes": (lambda frame: initiate_root(self, 'map_barcodes'), "Map barcodes to data."),
             "Regression": (lambda frame: initiate_root(self, 'regression'), "Perform regression analysis."),
-            "Recruitment": (lambda frame: initiate_root(self, 'recruitment'), "Analyze recruitment data.")
+            "Recruitment": (lambda frame: initiate_root(self, 'recruitment'), "Analyze recruitment data."),
+            "Plaque": (lambda frame: initiate_root(self, 'analyze_plaques'), "Analyze plaque data.")
         }
 
         self.selected_app = tk.StringVar()

spacr/gui_core.py

@@ -440,6 +440,8 @@ def import_settings(settings_type='mask'):
         settings = set_default_umap_image_settings(settings={})
     elif settings_type == 'recruitment':
         settings = get_analyze_recruitment_default_settings(settings={})
+    elif settings_type == 'analyze_plaques':
+        settings = {}
     else:
         raise ValueError(f"Invalid settings type: {settings_type}")
     
@@ -493,6 +495,8 @@ def setup_settings_panel(vertical_container, settings_type='mask'):
         settings = get_perform_regression_default_settings(settings={})
     elif settings_type == 'recruitment':
         settings = get_analyze_recruitment_default_settings(settings={})
+    elif settings_type == 'analyze_plaques':
+        settings = {}
     else:
         raise ValueError(f"Invalid settings type: {settings_type}")
 
@@ -645,7 +649,7 @@ def setup_usage_panel(horizontal_container, btn_col, uppdate_frequency):
     widgets = [usage_scrollable_frame.scrollable_frame]
     
     usage_bars = []
-    max_elements_per_column = 6
+    max_elements_per_column = 5
     row = 0
     col = 0
 
@@ -773,9 +777,9 @@ def start_process(q=None, fig_queue=None, settings_type='mask'):
     initialize_cuda()
 
     process_args = (settings_type, settings, q, fig_queue, stop_requested)
-    if settings_type in ['mask', 'umap', 'measure', 'simulation', 'sequencing', 
-                         'classify', 'cellpose_dataset', 'train_cellpose', 'ml_analyze', 'cellpose_masks', 'cellpose_all', 'map_barcodes', 
-                         'regression', 'recruitment', 'plaques', 'cellpose_compare', 'vision_scores', 'vision_dataset']:
+    if settings_type in ['mask', 'umap', 'measure', 'simulation', 'sequencing', 'classify', 'analyze_plaques', 
+                         'cellpose_dataset', 'train_cellpose', 'ml_analyze', 'cellpose_masks', 'cellpose_all', 'map_barcodes', 
+                         'regression', 'recruitment', 'cellpose_compare', 'vision_scores', 'vision_dataset']:
 
         # Start the process
         process = Process(target=run_function_gui, args=process_args)

spacr/gui_utils.py

@@ -486,7 +486,7 @@ def function_gui_wrapper(function=None, settings={}, q=None, fig_queue=None, imp
 def run_function_gui(settings_type, settings, q, fig_queue, stop_requested):
     
     from .gui_utils import process_stdout_stderr
-    from .core import generate_image_umap, preprocess_generate_masks, generate_ml_scores, identify_masks_finetune, check_cellpose_models, analyze_recruitment, train_cellpose, compare_cellpose_masks, analyze_plaques, generate_dataset, apply_model_to_tar
+    from .core import generate_image_umap, preprocess_generate_masks, generate_ml_scores, identify_masks_finetune, check_cellpose_models, analyze_recruitment, train_cellpose, analyze_plaques, compare_cellpose_masks, generate_dataset, apply_model_to_tar
     from .io import generate_cellpose_train_test
     from .measure import measure_crop
     from .sim import run_multiple_simulations
@@ -532,6 +532,9 @@ def run_function_gui(settings_type, settings, q, fig_queue, stop_requested):
     elif settings_type == 'umap':
         function = generate_image_umap
         imports = 1
+    elif settings_type == 'analyze_plaques':
+        function = analyze_plaques
+        imports = 1
     else:
         raise ValueError(f"Invalid settings type: {settings_type}")
     try:

spacr/io.py

@@ -26,7 +26,7 @@ import atexit
 
 from .logger import log_function_call
 
-def _load_images_and_labels(image_files, label_files, circular=False, invert=False, image_extension="*.tif", label_extension="*.tif"):
+def _load_images_and_labels(image_files, label_files, circular=False, invert=False):
     
     from .utils import invert_image, apply_mask
     

spacr/measure.py

@@ -998,10 +998,10 @@ def measure_crop(settings):
         src_fldr = os.path.join(src_fldr, 'merged')
         print(f"Changed source folder to: {src_fldr}")
 
-    if settings['save_measurements']:
-        source_folder = os.path.dirname(settings['src'])
-        os.makedirs(source_folder+'/measurements', exist_ok=True)
-        _create_database(source_folder+'/measurements/measurements.db')
+    #if settings['save_measurements']:
+        #source_folder = os.path.dirname(settings['src'])
+        #os.makedirs(source_folder+'/measurements', exist_ok=True)
+        #_create_database(source_folder+'/measurements/measurements.db')
     
     if settings['cell_mask_dim'] is None:
         settings['include_uninfected'] = True

spacr/mediar.py

@@ -0,0 +1,366 @@
+import os, sys, gdown, cv2, torch
+import numpy as np
+import matplotlib.pyplot as plt
+from monai.inferers import sliding_window_inference
+import skimage.io as io
+
+# Path to the MEDIAR directory
+mediar_path = os.path.join(os.path.dirname(__file__), 'resources', 'MEDIAR')
+
+print('mediar path', mediar_path)
+
+# Temporarily create __init__.py to make MEDIAR a package
+init_file = os.path.join(mediar_path, '__init__.py')
+if not os.path.exists(init_file):
+    with open(init_file, 'w'):  # Create the __init__.py file
+        pass
+
+# Add MEDIAR to sys.path
+sys.path.insert(0, mediar_path)
+
+try:
+    # Now import the dependencies from MEDIAR
+    from core.MEDIAR import Predictor, EnsemblePredictor
+    from train_tools.models import MEDIARFormer
+
+    print("Imports successful.")
+finally:
+    # Remove the temporary __init__.py file after the import
+    if os.path.exists(init_file):
+        os.remove(init_file)  # Remove the __init__.py file
+
+def display_imgs_in_list(lists_of_imgs, cmaps=None):
+    """
+    Displays images from multiple lists side by side. 
+    Each row will display one image from each list (lists_of_imgs[i][j] is the j-th image in the i-th list).
+    
+    :param lists_of_imgs: A list of lists, where each inner list contains images.
+    :param cmaps: List of colormaps to use for each list (optional). If not provided, defaults to 'gray' for all lists.
+    """
+    num_lists = len(lists_of_imgs)
+    num_images = len(lists_of_imgs[0])
+    
+    # Ensure that all lists have the same number of images
+    for img_list in lists_of_imgs:
+        assert len(img_list) == num_images, "All inner lists must have the same number of images"
+
+    # Use 'gray' as the default colormap if cmaps are not provided
+    if cmaps is None:
+        cmaps = ['gray'] * num_lists
+    else:
+        assert len(cmaps) == num_lists, "The number of colormaps must match the number of lists"
+
+    plt.figure(figsize=(15, 5 * num_images))
+    
+    for j in range(num_images):
+        for i, img_list in enumerate(lists_of_imgs):
+            img = img_list[j]
+            plt.subplot(num_images, num_lists, j * num_lists + i + 1)
+            
+            if len(img.shape) == 2:  # Grayscale image
+                plt.imshow(img, cmap=cmaps[i])
+            elif len(img.shape) == 3 and img.shape[0] == 3:  # 3-channel image (C, H, W)
+                plt.imshow(img.transpose(1, 2, 0))  # Change shape to (H, W, C) for displaying
+            else:
+                plt.imshow(img)
+            
+            plt.axis('off')
+            plt.title(f'Image {j+1} from list {i+1}')
+
+    plt.tight_layout()
+    plt.show()
+
+def get_weights(finetuned_weights=False):
+    if finetuned_weights:
+        model_path1 = os.path.join(os.path.dirname(__file__), 'resources', 'MEDIAR_weights', 'from_phase1.pth')
+        if not os.path.exists(model_path1):
+            print("Downloading finetuned model 1...")
+            gdown.download('https://drive.google.com/uc?id=1JJ2-QKTCk-G7sp5ddkqcifMxgnyOrXjx', model_path1, quiet=False)
+    else:
+        model_path1 = os.path.join(os.path.dirname(__file__), 'resources', 'MEDIAR_weights', 'phase1.pth')
+        if not os.path.exists(model_path1):
+            print("Downloading model 1...")
+            gdown.download('https://drive.google.com/uc?id=1v5tYYJDqiwTn_mV0KyX5UEonlViSNx4i', model_path1, quiet=False)
+            
+    if finetuned_weights:
+        model_path2 = os.path.join(os.path.dirname(__file__), 'resources', 'MEDIAR_weights', 'from_phase2.pth')
+        if not os.path.exists(model_path2):
+            print("Downloading finetuned model 2...")
+            gdown.download('https://drive.google.com/uc?id=168MtudjTMLoq9YGTyoD2Rjl_d3Gy6c_L', model_path2, quiet=False)
+    else:
+        model_path2 = os.path.join(os.path.dirname(__file__), 'resources', 'MEDIAR_weights', 'phase2.pth')
+        if not os.path.exists(model_path2):
+            print("Downloading model 2...")
+            gdown.download('https://drive.google.com/uc?id=1NHDaYvsYz3G0OCqzegT-bkNcly2clPGR', model_path2, quiet=False)
+    
+    return model_path1, model_path2
+
+def normalize_image(image, lower_percentile=0.0, upper_percentile=99.5):
+    """
+    Normalize an image based on the 0.0 and 99.5 percentiles.
+    
+    :param image: Input image (numpy array).
+    :param lower_percentile: Lower percentile (default is 0.0).
+    :param upper_percentile: Upper percentile (default is 99.5).
+    :return: Normalized image (numpy array).
+    """
+    lower_bound = np.percentile(image, lower_percentile)
+    upper_bound = np.percentile(image, upper_percentile)
+    
+    # Clip image values to the calculated percentiles
+    image = np.clip(image, lower_bound, upper_bound)
+    
+    # Normalize to [0, 1]
+    image = (image - lower_bound) / (upper_bound - lower_bound + 1e-5)  # Add small epsilon to avoid division by zero
+    
+    return image
+
+class MEDIARPredictor:
+    def __init__(self, input_path=None, output_path=None, device=None, model="ensemble", roi_size=512, overlap=0.6, finetuned_weights=False, test=False, use_tta=False, normalize=True, quantiles=[0.0, 99.5]):
+        if device is None:
+            device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+        self.test = test
+        self.model = model
+        self.normalize = normalize
+        self.quantiles = quantiles
+
+        # Paths to model weights
+        self.model1_path, self.model2_path = get_weights(finetuned_weights)
+
+        # Load main models
+        self.model1 = self.load_model(self.model1_path, device=self.device)
+        self.model2 = self.load_model(self.model2_path, device=self.device) if model == "ensemble" or model == "model2" else None
+        if self.test:
+            # Define input and output paths for running test
+            self.input_path = os.path.join(os.path.dirname(__file__), 'resources/images')
+            self.output_path = os.path.join(os.path.dirname(__file__), 'resources/MEDIAR/results')
+        else:
+            self.input_path = input_path
+            self.output_path = output_path
+
+        # If using a single model
+        if self.model == "model1":
+            self.predictor = Predictor(
+                model=self.model1,
+                device=self.device,
+                input_path=self.input_path,
+                output_path=self.output_path,
+                algo_params={"use_tta": use_tta}
+            )
+
+        # If using a single model
+        if self.model == "model2":
+            self.predictor = Predictor(
+                model=self.model2,
+                device=self.device,
+                input_path=self.input_path,
+                output_path=self.output_path,
+                algo_params={"use_tta": use_tta}
+            )
+
+        # If using two models
+        elif self.model == "ensemble":
+            self.predictor = EnsemblePredictor(
+                model=self.model1,  # Pass model1 as model
+                model_aux=self.model2,  # Pass model2 as model_aux
+                device=self.device,
+                input_path=self.input_path,
+                output_path=self.output_path,
+                algo_params={"use_tta": use_tta}
+            )
+
+        if self.test:
+            self.run_test()
+
+        if not self.model in ["model1", "model2", "ensemble"]:
+            raise ValueError("Invalid model type. Choose from 'model1', 'model2', or 'ensemble'.")
+
+    def load_model(self, model_path, device):
+        model_args = {
+            "classes": 3,
+            "decoder_channels": [1024, 512, 256, 128, 64],
+            "decoder_pab_channels": 256,
+            "encoder_name": 'mit_b5',
+            "in_channels": 3
+        }
+        model = MEDIARFormer(**model_args)
+        weights = torch.load(model_path, map_location=device)
+        model.load_state_dict(weights, strict=False)
+        model.to(device)
+        model.eval()
+        return model
+    
+    def display_image_and_mask(self, img, mask):
+
+        from .plot import generate_mask_random_cmap
+        """
+        Displays the normalized input image and the predicted mask side by side.
+        """
+        # If img is a tensor, convert it to NumPy for display
+        if isinstance(img, torch.Tensor):
+            img = img.cpu().numpy()
+
+        # If mask is a tensor, convert it to NumPy for display
+        if isinstance(mask, torch.Tensor):
+            mask = mask.cpu().numpy()
+
+        # Transpose the image to have (H, W, C) format for display if needed
+        if len(img.shape) == 3 and img.shape[0] == 3:
+            img = img.transpose(1, 2, 0)
+
+        # Scale the normalized image back to [0, 255] for proper display
+        img_display = (img * 255).astype(np.uint8)
+
+        plt.figure(figsize=(10, 5))
+
+        # Display normalized image
+        plt.subplot(1, 2, 1)
+        plt.imshow(img_display)
+        plt.title("Normalized Image")
+        plt.axis("off")
+
+        r_cmap = generate_mask_random_cmap(mask)
+
+        # Display predicted mask
+        plt.subplot(1, 2, 2)
+        plt.imshow(mask, cmap=r_cmap)
+        plt.title("Predicted Mask")
+        plt.axis("off")
+
+        plt.tight_layout()
+        plt.show()
+
+    def predict_batch(self, imgs):
+        """
+        Predict masks for a batch of images.
+        
+        :param imgs: List of input images as NumPy arrays (each in (H, W, C) format).
+        :return: List of predicted masks as NumPy arrays.
+        """
+        processed_imgs = []
+
+        # Preprocess and normalize each image
+        for img in imgs:
+            if self.normalize:
+                # Normalize the image using the specified quantiles
+                img_normalized = normalize_image(img, lower_percentile=self.quantiles[0], upper_percentile=self.quantiles[1])
+            else:
+                img_normalized = img
+
+            # Convert image to tensor and send to device
+            img_tensor = torch.tensor(img_normalized.astype(np.float32).transpose(2, 0, 1)).to(self.device)  # (C, H, W)
+            processed_imgs.append(img_tensor)
+
+        # Stack all processed images into a batch tensor
+        batch_tensor = torch.stack(processed_imgs)
+
+        # Run inference to get predicted masks
+        pred_masks = self.predictor._inference(batch_tensor)
+
+        # Ensure pred_masks is always treated as a batch
+        if len(pred_masks.shape) == 3:  # If single image, add batch dimension
+            pred_masks = pred_masks.unsqueeze(0)
+
+        # Convert predictions to NumPy arrays and post-process each mask
+        predicted_masks = []
+        for pred_mask in pred_masks:
+            pred_mask_np = pred_mask.cpu().numpy()
+
+            # Extract dP and cellprob from pred_mask
+            dP = pred_mask_np[:2]  # First two channels as dP (displacement field)
+            cellprob = pred_mask_np[2]  # Third channel as cell probability
+
+            # Concatenate dP and cellprob along axis 0 to pass a single array
+            combined_pred_mask = np.concatenate([dP, np.expand_dims(cellprob, axis=0)], axis=0)
+
+            # Post-process the predicted mask
+            mask = self.predictor._post_process(combined_pred_mask)
+
+            # Append the processed mask to the list
+            predicted_masks.append(mask.astype(np.uint16))
+
+        return predicted_masks
+
+    def run_test(self):
+        """
+        Run the model on test images if the test flag is True.
+        """
+        # List of input images
+        imgs = []
+        img_names = []
+        
+        for img_file in os.listdir(self.input_path):
+            img_path = os.path.join(self.input_path, img_file)
+            img = io.imread(img_path)
+            
+            # Check if the image is grayscale (2D) or RGB (3D), and convert grayscale to RGB
+            if len(img.shape) == 2:  # Grayscale image (H, W)
+                img = np.repeat(img[:, :, np.newaxis], 3, axis=2)  # Convert grayscale to RGB
+
+            # Normalize the image if the normalize flag is True
+            if self.normalize:
+                img_normalized = normalize_image(img, lower_percentile=self.quantiles[0], upper_percentile=self.quantiles[1])
+            else:
+                img_normalized = img
+            
+            # Convert image to tensor and send directly to device
+            img_tensor = torch.tensor(img_normalized.astype(np.float32).transpose(2, 0, 1)).to(self.device)  # (C, H, W)
+
+            imgs.append(img_tensor)
+            img_names.append(os.path.splitext(img_file)[0])
+
+        # Stack all images into a batch (ensure it's always treated as a batch)
+        batch_tensor = torch.stack(imgs)
+
+        # Predict using the predictor (or ensemble predictor)
+        pred_masks = self.predictor._inference(batch_tensor)
+
+        # Ensure pred_masks is always treated as a batch
+        if len(pred_masks.shape) == 3:  # If single image, add batch dimension
+            pred_masks = pred_masks.unsqueeze(0)
+
+        # Convert predictions to NumPy arrays and post-process each mask
+        for i, pred_mask in enumerate(pred_masks):
+            # Ensure the dimensions of pred_mask remain consistent
+            pred_mask_np = pred_mask.cpu().numpy()
+
+            # Extract dP and cellprob from pred_mask
+            dP = pred_mask_np[:2]  # First two channels as dP (displacement field)
+            cellprob = pred_mask_np[2]  # Third channel as cell probability
+
+            # Concatenate dP and cellprob along axis 0 to pass a single array
+            combined_pred_mask = np.concatenate([dP, np.expand_dims(cellprob, axis=0)], axis=0)
+
+            # Post-process the predicted mask
+            mask = self.predictor._post_process(combined_pred_mask)
+
+            # Convert the mask to 16-bit format (ensure values fit into 16-bit range)
+            mask_to_save = mask.astype(np.uint16)
+
+            # Save the post-processed mask as a .tif file using cv2
+            mask_output_path = os.path.join(self.output_path, f"{img_names[i]}_mask.tiff")
+            cv2.imwrite(mask_output_path, mask_to_save)
+
+            print(f"Predicted mask saved at: {mask_output_path}")
+
+            self.display_image_and_mask(imgs[i].cpu().numpy(), mask)
+
+        print(f"Test predictions saved in {self.output_path}")
+
+    def preprocess_image(self, img):
+        """
+        Preprocess input image (numpy array) for compatibility with the model.
+        """
+        if isinstance(img, np.ndarray):  # Check if the input is a numpy array
+            if len(img.shape) == 2:  # Grayscale image (H, W)
+                img = np.repeat(img[:, :, np.newaxis], 3, axis=2)
+
+            elif img.shape[2] == 1:  # Single channel grayscale (H, W, 1)
+                img = np.repeat(img, 3, axis=2)  # Convert to 3-channel RGB
+
+            img_tensor = torch.tensor(img.astype(np.float32).transpose(2, 0, 1))  # Change shape to (C, H, W)
+        else:
+            img_tensor = img  # If it's already a tensor, assume it's in (C, H, W) format
+
+        return img_tensor.float()