Added performance optimizations and bumped version

README.md

@@ -41,6 +41,9 @@ We provide three Jupyter notebooks (in the `templates` folder), each with its ow
    - Lets you interactively explore the Nimbus Gold Standard labeled dataset.  
    - [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://drive.google.com/file/d/1LW0vHC3sKKA3TyvW_9FeIaHj3PonzhGS/view?usp=sharing)
 
+- **[4_Generic_Cell_Clustering.ipynb](https://github.com/angelolab/Nimbus-Inference/blob/main/templates/4_Generic_Cell_Clustering.ipynb)**  
+   - Performs cell phenotyping using SOM (Self-Organizing Map) and consensus clustering on Nimbus predictions. Takes the cell table output from `1_Nimbus_Predict.ipynb` and clusters cells based on marker expression, enabling downstream analysis and visualization of cell populations.
+
 Each notebook loads an example dataset directly from the Hugging Face Hub, so you can get hands-on with Nimbus right away.
 
 ---
@@ -66,7 +69,7 @@ Nimbus is a deep learning model designed to make **human-like, visual classifica
 - Generalizes across many tissue types, imaging platforms, and markers — without retraining.
 - Can be integrated into downstream clustering or phenotyping pipelines to improve accuracy.
 
-For more details, please see our [preprint](https://pmc.ncbi.nlm.nih.gov/articles/PMC11185540/).
+For more details, please see our [publication](https://www.nature.com/articles/s41592-025-02826-9).
 
 ---
 

pyproject.toml

@@ -4,7 +4,7 @@ requires = ["hatchling"]
 
 [project]
 name = "Nimbus-Inference"
-version = "0.0.4"
+version = "0.0.5"
 description = "A model for classification of cells into marker positive / negative"
 readme = "README.md"
 requires-python = ">=3.9"

src/nimbus_inference/nimbus.py

@@ -82,11 +82,14 @@ class Nimbus(nn.Module):
                 , with "cpu" as a fallback), "cpu", "cuda", or "mps". Defaults to "auto".
             checkpoint: which checkpoint to use for the model, either "latest" or one of the local
                 checkpoints.
+            compile_model (bool): Whether to use torch.compile() for faster inference.
+            mixed_precision (bool): Whether to use mixed precision (FP16) inference.
     """
     def __init__(
         self, dataset: MultiplexDataset, output_dir: str, save_predictions: bool=True,
         batch_size: int=4, test_time_aug: bool=True, model_magnification: int=10,
-        input_shape: list=[1024, 1024], device: str="auto",  checkpoint="latest"
+        input_shape: list=[1024, 1024], device: str="auto", checkpoint="latest",
+        compile_model: bool=False, mixed_precision: bool=False
     ):
         super(Nimbus, self).__init__()
         self.dataset = dataset
@@ -98,6 +101,8 @@ def __init__(
         self.test_time_aug = test_time_aug
         self.input_shape = input_shape
         self.checkpoint = checkpoint
+        self.compile_model = compile_model
+        self.mixed_precision = mixed_precision
         if self.output_dir != "":
             os.makedirs(self.output_dir, exist_ok=True)
 
@@ -159,6 +164,12 @@ def load_local_checkpoint(self, checkpoint, padding="reflect"):
         model.load_state_dict(torch.load(self.checkpoint_path))
         print(f"Loaded weights from {self.checkpoint_path}")
         self.model = model.to(self.device).eval()
+
+        # Apply torch.compile() if enabled
+        if self.compile_model:
+            print("Compiling model with torch.compile()...")
+            self.model = torch.compile(self.model, mode="reduce-overhead")
+            print("Model compiled successfully.")
 
     def load_latest_checkpoint(self, padding="reflect"):
         """Initializes the model and loads the latest weights from Hugging Face Hub if newer
@@ -226,6 +237,12 @@ def load_latest_checkpoint(self, padding="reflect"):
         model.load_state_dict(torch.load(self.checkpoint_path))
         print(f"Loaded weights from {self.checkpoint_path}")
         self.model = model.to(self.device).eval()
+
+        # Apply torch.compile() if enabled
+        if self.compile_model:
+            print("Compiling model with torch.compile()...")
+            self.model = torch.compile(self.model, mode="reduce-overhead")
+            print("Model compiled successfully.")
 
     def load_checkpoint(self, padding="reflect"):
         """Loads the model checkpoint.
@@ -276,7 +293,12 @@ def predict_segmentation(self, input_data):
                 if not isinstance(input_data, torch.Tensor):
                     input_data = torch.tensor(input_data).float()
                 input_data = input_data.to(self.device)
-                prediction = self.model(input_data)
+                # Use mixed precision (FP16) on CUDA if enabled
+                if self.mixed_precision:
+                    with torch.autocast(device_type=self.device.type, dtype=torch.float16):
+                        prediction = self.model(input_data)
+                else:
+                    prediction = self.model(input_data)
                 prediction = prediction.cpu()
         else:
             if not hasattr(self, "model") or self.model.padding != "valid":
@@ -316,7 +338,12 @@ def _tile_and_stitch(self, input_data):
             batch = torch.from_numpy(tiled_input[i : i + self.batch_size]).float()
             batch = batch.to(self.device)
             with torch.no_grad():
-                pred = self.model(batch).cpu().numpy()
+                # Use mixed precision (FP16) on CUDA if enabled
+                if self.mixed_precision:
+                    with torch.autocast(device_type=self.device.type, dtype=torch.float16):
+                        pred = self.model(batch).cpu().numpy()
+                else:
+                    pred = self.model(batch).cpu().numpy()
                 # crop pred
                 if self.crop_by.any():
                     pred = pred[

src/nimbus_inference/utils.py

@@ -449,8 +449,10 @@ def get_channel_stack(self, fov: str, channel: str):
         """
         idx = self.fovs.index(fov)
         fov_path = self.fov_paths[idx]
-        self.img_reader = LazyOMETIFFReader(fov_path)
-        return np.squeeze(self.img_reader.get_channel(channel))
+        # Cache the reader to avoid recreating it for each channel access
+        if not hasattr(self, '_reader_cache') or self._reader_cache[0] != fov_path:
+            self._reader_cache = (fov_path, LazyOMETIFFReader(fov_path))
+        return np.squeeze(self._reader_cache[1].get_channel(channel))
 
     def get_segmentation(self, fov: str):
         """Get the instance mask for a fov
@@ -500,18 +502,39 @@ def prepare_normalization_dict(
                 n_jobs
             )
 
-def prepare_input_data(mplex_img, instance_mask):
-    """Prepares the input data for the segmentation model
+def prepare_binary_mask(instance_mask):
+    """Prepares the binary mask from the instance mask. Can be cached per FOV.
 
     Args:
-        mplex_img (np.array): multiplex image
         instance_mask (np.array): instance mask
     Returns:
-        np.array: input data for segmentation model
+        np.array: binary mask (1 where cells are, 0 at boundaries and background)
     """
-    mplex_img = mplex_img.astype(np.float32)
     edge = find_boundaries(instance_mask, mode="inner").astype(np.uint8)
     binary_mask = np.logical_and(edge == 0, instance_mask > 0).astype(np.float32)
+    return binary_mask
+
+
+def prepare_input_data(mplex_img, instance_mask=None, binary_mask=None):
+    """Prepares the input data for the segmentation model.
+
+    Either instance_mask or binary_mask must be provided. If binary_mask is provided,
+    it will be used directly (faster for processing multiple channels from the same FOV).
+    If only instance_mask is provided, the binary_mask will be computed from it.
+
+    Args:
+        mplex_img (np.array): multiplex image
+        instance_mask (np.array, optional): instance mask (used to compute binary_mask if 
+            binary_mask is not provided)
+        binary_mask (np.array, optional): pre-computed binary mask from prepare_binary_mask()
+    Returns:
+        np.array: input data for segmentation model
+    """
+    if binary_mask is None:
+        if instance_mask is None:
+            raise ValueError("Either instance_mask or binary_mask must be provided")
+        binary_mask = prepare_binary_mask(instance_mask)
+    mplex_img = mplex_img.astype(np.float32)
     input_data = np.stack([mplex_img, binary_mask], axis=0)[np.newaxis,...] # bhwc
     return input_data
 
@@ -606,18 +629,24 @@ def predict_fovs(
         )
         df_fov = pd.DataFrame()
         instance_mask = dataset.get_segmentation(fov)
+        # Cache binary mask for this FOV (avoids recomputing find_boundaries for each channel)
+        binary_mask = prepare_binary_mask(instance_mask)
+        # Pre-compute scaled binary mask if magnification differs
+        if dataset.magnification != nimbus.model_magnification:
+            scale = nimbus.model_magnification / dataset.magnification
+            h, w = binary_mask.shape
+            scaled_h, scaled_w = int(h * scale), int(w * scale)
+            binary_mask_scaled = cv2.resize(binary_mask, [scaled_w, scaled_h], interpolation=0)
+        else:
+            binary_mask_scaled = None
         for channel_name in tqdm(dataset.channels):
             mplex_img = dataset.get_channel_normalized(fov, channel_name)
-            input_data = prepare_input_data(mplex_img, instance_mask)
             if dataset.magnification != nimbus.model_magnification:
-                scale = nimbus.model_magnification / dataset.magnification
-                input_data = np.squeeze(input_data)
-                _, h,w = input_data.shape
-                img = cv2.resize(input_data[0], [int(w*scale), int(h*scale)])
-                binary_mask = cv2.resize(
-                    input_data[1], [int(w*scale), int(h*scale)], interpolation=0
-                )
-                input_data = np.stack([img, binary_mask], axis=0)[np.newaxis,...]
+                h, w = mplex_img.shape
+                img_scaled = cv2.resize(mplex_img, [scaled_w, scaled_h])
+                input_data = prepare_input_data(img_scaled, binary_mask=binary_mask_scaled)
+            else:
+                input_data = prepare_input_data(mplex_img, binary_mask=binary_mask)
             if test_time_augmentation:
                 prediction = test_time_aug(
                     input_data, channel_name, nimbus, dataset.normalization_dict,

templates/1_Nimbus_Predict.ipynb

@@ -48,7 +48,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "id": "974f8dda",
    "metadata": {},
    "outputs": [],
@@ -67,7 +67,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "37733de5",
    "metadata": {},
    "outputs": [],
@@ -87,7 +87,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "id": "292e4524",
    "metadata": {},
    "outputs": [],
@@ -115,7 +115,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "id": "65a319c9",
    "metadata": {},
    "outputs": [],
@@ -223,6 +223,8 @@
     "    input_shape=[1024,1024],\n",
     "    device=\"auto\",\n",
     "    output_dir=nimbus_output_dir,\n",
+    "    compile_model=False,\n",
+    "    mixed_precision=False,\n",
     ")\n",
     "\n",
     "# check if all inputs are valid\n",
@@ -308,7 +310,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Nimbus",
+   "display_name": "nimbus",
    "language": "python",
    "name": "python3"
   },
@@ -322,7 +324,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.0"
+   "version": "3.11.11"
   }
  },
  "nbformat": 4,