From 4e6e42053915af332f6aefc43dd09ddf0d501585 Mon Sep 17 00:00:00 2001 From: Lorenz Rumberger Date: Tue, 6 Jan 2026 11:34:39 +0100 Subject: [PATCH] Added performance optimizations and bumped version --- README.md | 5 ++- pyproject.toml | 2 +- src/nimbus_inference/nimbus.py | 33 +++++++++++++++-- src/nimbus_inference/utils.py | 61 +++++++++++++++++++++++--------- templates/1_Nimbus_Predict.ipynb | 14 ++++---- 5 files changed, 88 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index 248db6d..3f00029 100644 --- a/README.md +++ b/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). --- diff --git a/pyproject.toml b/pyproject.toml index 8755ced..945d907 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/nimbus_inference/nimbus.py b/src/nimbus_inference/nimbus.py index ae5464c..77c57ed 100644 --- a/src/nimbus_inference/nimbus.py +++ b/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[ diff --git a/src/nimbus_inference/utils.py b/src/nimbus_inference/utils.py index 4ad4b0c..3dd808f 100644 --- a/src/nimbus_inference/utils.py +++ b/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, diff --git a/templates/1_Nimbus_Predict.ipynb b/templates/1_Nimbus_Predict.ipynb index fa269f7..3c97955 100644 --- a/templates/1_Nimbus_Predict.ipynb +++ b/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,