|
| 1 | +# Copyright 2021-2025 MONAI Consortium |
| 2 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 3 | +# you may not use this file except in compliance with the License. |
| 4 | +# You may obtain a copy of the License at |
| 5 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 6 | +# Unless required by applicable law or agreed to in writing, software |
| 7 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 8 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 9 | +# See the License for the specific language governing permissions and |
| 10 | +# limitations under the License. |
| 11 | + |
| 12 | +import logging |
| 13 | +from pathlib import Path |
| 14 | +from typing import List |
| 15 | + |
| 16 | +import torch |
| 17 | +from numpy import float32, int16 |
| 18 | + |
| 19 | +# import custom transforms from post_transforms.py |
| 20 | +from post_transforms import CalculateVolumeFromMaskd, ExtractVolumeToTextd, LabelToContourd, OverlayImageLabeld |
| 21 | + |
| 22 | +import monai |
| 23 | +from monai.deploy.core import AppContext, Fragment, Model, Operator, OperatorSpec |
| 24 | +from monai.deploy.operators.monai_seg_inference_operator import InfererType, InMemImageReader, MonaiSegInferenceOperator |
| 25 | +from monai.transforms import ( |
| 26 | + Activationsd, |
| 27 | + AsDiscreted, |
| 28 | + CastToTyped, |
| 29 | + Compose, |
| 30 | + CropForegroundd, |
| 31 | + EnsureChannelFirstd, |
| 32 | + EnsureTyped, |
| 33 | + Invertd, |
| 34 | + LoadImaged, |
| 35 | + Orientationd, |
| 36 | + SaveImaged, |
| 37 | + ScaleIntensityRanged, |
| 38 | + Spacingd, |
| 39 | +) |
| 40 | + |
| 41 | + |
| 42 | +# this operator performs inference with the new version of the bundle |
| 43 | +class AbdomenSegOperator(Operator): |
| 44 | + """Performs segmentation inference with a custom model architecture.""" |
| 45 | + |
| 46 | + DEFAULT_OUTPUT_FOLDER = Path.cwd() / "output" |
| 47 | + |
| 48 | + def __init__( |
| 49 | + self, |
| 50 | + fragment: Fragment, |
| 51 | + *args, |
| 52 | + app_context: AppContext, |
| 53 | + model_path: Path, |
| 54 | + output_folder: Path = DEFAULT_OUTPUT_FOLDER, |
| 55 | + output_labels: List[int], |
| 56 | + **kwargs, |
| 57 | + ): |
| 58 | + |
| 59 | + self._logger = logging.getLogger(f"{__name__}.{type(self).__name__}") |
| 60 | + self._input_dataset_key = "image" |
| 61 | + self._pred_dataset_key = "pred" |
| 62 | + |
| 63 | + # self.model_path is compatible with TorchScript and PyTorch model workflows (pythonic and MAP) |
| 64 | + self.model_path = self._find_model_file_path(model_path) |
| 65 | + |
| 66 | + self.output_folder = output_folder |
| 67 | + self.output_folder.mkdir(parents=True, exist_ok=True) |
| 68 | + self.output_labels = output_labels |
| 69 | + self.app_context = app_context |
| 70 | + self.input_name_image = "image" |
| 71 | + self.output_name_seg = "seg_image" |
| 72 | + self.output_name_text_dicom_sr = "result_text_dicom_sr" |
| 73 | + self.output_name_text_mongodb = "result_text_mongodb" |
| 74 | + self.output_name_sc_path = "dicom_sc_dir" |
| 75 | + |
| 76 | + # the base class has an attribute called fragment to hold the reference to the fragment object |
| 77 | + super().__init__(fragment, *args, **kwargs) |
| 78 | + |
| 79 | + # find model path; supports TorchScript and PyTorch model workflows (pythonic and MAP) |
| 80 | + def _find_model_file_path(self, model_path: Path): |
| 81 | + # when executing pythonically, model_path is a file |
| 82 | + # when executing as MAP, model_path is a directory (/opt/holoscan/models) |
| 83 | + # torch.load() from PyTorch workflow needs file path; can't load model from directory |
| 84 | + # returns first found file in directory in this case |
| 85 | + if model_path: |
| 86 | + if model_path.is_file(): |
| 87 | + return model_path |
| 88 | + elif model_path.is_dir(): |
| 89 | + for file in model_path.rglob("*"): |
| 90 | + if file.is_file(): |
| 91 | + return file |
| 92 | + |
| 93 | + raise ValueError(f"Model file not found in the provided path: {model_path}") |
| 94 | + |
| 95 | + # load a PyTorch model and register it in app_context |
| 96 | + def _load_pytorch_model(self): |
| 97 | + |
| 98 | + _device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 99 | + _kernel_size: tuple = (3, 3, 3, 3, 3, 3) |
| 100 | + _strides: tuple = (1, 2, 2, 2, 2, (2, 2, 1)) |
| 101 | + _upsample_kernel_size: tuple = (2, 2, 2, 2, (2, 2, 1)) |
| 102 | + |
| 103 | + # create DynUNet model with the specified architecture parameters + move to computational device (GPU or CPU) |
| 104 | + # parameters pulled from inference.yaml file of the MONAI bundle |
| 105 | + model = monai.networks.nets.dynunet.DynUNet( |
| 106 | + spatial_dims=3, |
| 107 | + in_channels=1, |
| 108 | + out_channels=4, |
| 109 | + kernel_size=_kernel_size, |
| 110 | + strides=_strides, |
| 111 | + upsample_kernel_size=_upsample_kernel_size, |
| 112 | + norm_name="INSTANCE", |
| 113 | + deep_supervision=False, |
| 114 | + res_block=True, |
| 115 | + ).to(_device) |
| 116 | + |
| 117 | + # load model state dictionary (i.e. mapping param names to tensors) via torch.load |
| 118 | + # weights_only=True to avoid arbitrary code execution during unpickling |
| 119 | + state_dict = torch.load(self.model_path, weights_only=True) |
| 120 | + |
| 121 | + # assign loaded weights to model architecture via load_state_dict |
| 122 | + model.load_state_dict(state_dict) |
| 123 | + |
| 124 | + # set model in evaluation (inference) mode |
| 125 | + model.eval() |
| 126 | + |
| 127 | + # create a MONAI Model object to encapsulate the PyTorch model and metadata |
| 128 | + loaded_model = Model(self.model_path, name="ped_abd_ct_seg") |
| 129 | + |
| 130 | + # assign loaded PyTorch model as the predictor for the Model object |
| 131 | + loaded_model.predictor = model |
| 132 | + |
| 133 | + # register the loaded Model object in the application context so other operators can access it |
| 134 | + # MonaiSegInferenceOperator uses _get_model method to load models; looks at app_context.models first |
| 135 | + self.app_context.models = loaded_model |
| 136 | + |
| 137 | + def setup(self, spec: OperatorSpec): |
| 138 | + spec.input(self.input_name_image) |
| 139 | + |
| 140 | + # DICOM SEG |
| 141 | + spec.output(self.output_name_seg) |
| 142 | + |
| 143 | + # DICOM SR |
| 144 | + spec.output(self.output_name_text_dicom_sr) |
| 145 | + |
| 146 | + # MongoDB |
| 147 | + spec.output(self.output_name_text_mongodb) |
| 148 | + |
| 149 | + # DICOM SC |
| 150 | + spec.output(self.output_name_sc_path) |
| 151 | + |
| 152 | + def compute(self, op_input, op_output, context): |
| 153 | + input_image = op_input.receive(self.input_name_image) |
| 154 | + if not input_image: |
| 155 | + raise ValueError("Input image is not found.") |
| 156 | + |
| 157 | + # this operator gets an in-memory Image object, so a specialized ImageReader is needed. |
| 158 | + _reader = InMemImageReader(input_image) |
| 159 | + |
| 160 | + # preprocessing and postprocessing |
| 161 | + pre_transforms = self.pre_process(_reader) |
| 162 | + post_transforms = self.post_process(pre_transforms) |
| 163 | + |
| 164 | + # if PyTorch model |
| 165 | + if self.model_path.suffix.lower() == ".pt": |
| 166 | + # load the PyTorch model |
| 167 | + self._logger.info("PyTorch model detected") |
| 168 | + self._load_pytorch_model() |
| 169 | + # else, we have TorchScript model |
| 170 | + else: |
| 171 | + self._logger.info("TorchScript model detected") |
| 172 | + |
| 173 | + # delegates inference and saving output to the built-in operator. |
| 174 | + infer_operator = MonaiSegInferenceOperator( |
| 175 | + self.fragment, |
| 176 | + roi_size=(96, 96, 96), |
| 177 | + pre_transforms=pre_transforms, |
| 178 | + post_transforms=post_transforms, |
| 179 | + overlap=0.75, |
| 180 | + app_context=self.app_context, |
| 181 | + model_name="", |
| 182 | + inferer=InfererType.SLIDING_WINDOW, |
| 183 | + sw_batch_size=4, |
| 184 | + model_path=self.model_path, |
| 185 | + name="monai_seg_inference_op", |
| 186 | + ) |
| 187 | + |
| 188 | + # setting the keys used in the dictionary-based transforms |
| 189 | + infer_operator.input_dataset_key = self._input_dataset_key |
| 190 | + infer_operator.pred_dataset_key = self._pred_dataset_key |
| 191 | + |
| 192 | + seg_image = infer_operator.compute_impl(input_image, context) |
| 193 | + |
| 194 | + # DICOM SEG |
| 195 | + op_output.emit(seg_image, self.output_name_seg) |
| 196 | + |
| 197 | + # grab result_text_dicom_sr and result_text_mongodb from ExractVolumeToTextd transform |
| 198 | + result_text_dicom_sr, result_text_mongodb = self.get_result_text_from_transforms(post_transforms) |
| 199 | + if not result_text_dicom_sr or not result_text_mongodb: |
| 200 | + raise ValueError("Result text could not be generated.") |
| 201 | + |
| 202 | + # only log volumes for target organs so logs reflect MAP behavior |
| 203 | + self._logger.info(f"Calculated Organ Volumes: {result_text_dicom_sr}") |
| 204 | + |
| 205 | + # DICOM SR |
| 206 | + op_output.emit(result_text_dicom_sr, self.output_name_text_dicom_sr) |
| 207 | + |
| 208 | + # MongoDB |
| 209 | + op_output.emit(result_text_mongodb, self.output_name_text_mongodb) |
| 210 | + |
| 211 | + # DICOM SC |
| 212 | + # temporary DICOM SC (w/o source DICOM metadata) saved in output_folder / temp directory |
| 213 | + dicom_sc_dir = self.output_folder / "temp" |
| 214 | + |
| 215 | + self._logger.info(f"Temporary DICOM SC saved at: {dicom_sc_dir}") |
| 216 | + |
| 217 | + op_output.emit(dicom_sc_dir, self.output_name_sc_path) |
| 218 | + |
| 219 | + def pre_process(self, img_reader) -> Compose: |
| 220 | + """Composes transforms for preprocessing the input image before predicting on a model.""" |
| 221 | + |
| 222 | + my_key = self._input_dataset_key |
| 223 | + |
| 224 | + return Compose( |
| 225 | + [ |
| 226 | + # img_reader: specialized InMemImageReader, derived from MONAI ImageReader |
| 227 | + LoadImaged(keys=my_key, reader=img_reader), |
| 228 | + EnsureChannelFirstd(keys=my_key), |
| 229 | + Orientationd(keys=my_key, axcodes="RAS"), |
| 230 | + Spacingd(keys=my_key, pixdim=[1.5, 1.5, 3.0], mode=["bilinear"]), |
| 231 | + ScaleIntensityRanged(keys=my_key, a_min=-250, a_max=400, b_min=0.0, b_max=1.0, clip=True), |
| 232 | + CropForegroundd(keys=my_key, source_key=my_key, mode="minimum"), |
| 233 | + EnsureTyped(keys=my_key), |
| 234 | + CastToTyped(keys=my_key, dtype=float32), |
| 235 | + ] |
| 236 | + ) |
| 237 | + |
| 238 | + def post_process(self, pre_transforms: Compose) -> Compose: |
| 239 | + """Composes transforms for postprocessing the prediction results.""" |
| 240 | + |
| 241 | + pred_key = self._pred_dataset_key |
| 242 | + |
| 243 | + labels = {"background": 0, "liver": 1, "spleen": 2, "pancreas": 3} |
| 244 | + |
| 245 | + return Compose( |
| 246 | + [ |
| 247 | + Activationsd(keys=pred_key, softmax=True), |
| 248 | + Invertd( |
| 249 | + keys=[pred_key, self._input_dataset_key], |
| 250 | + transform=pre_transforms, |
| 251 | + orig_keys=[self._input_dataset_key, self._input_dataset_key], |
| 252 | + meta_key_postfix="meta_dict", |
| 253 | + nearest_interp=[False, False], |
| 254 | + to_tensor=True, |
| 255 | + ), |
| 256 | + AsDiscreted(keys=pred_key, argmax=True), |
| 257 | + # custom post-processing steps |
| 258 | + CalculateVolumeFromMaskd(keys=pred_key, label_names=labels), |
| 259 | + # optional code for saving segmentation masks as a NIfTI |
| 260 | + # SaveImaged( |
| 261 | + # keys=pred_key, |
| 262 | + # output_ext=".nii.gz", |
| 263 | + # output_dir=self.output_folder / "NIfTI", |
| 264 | + # meta_keys="pred_meta_dict", |
| 265 | + # separate_folder=False, |
| 266 | + # output_dtype=int16 |
| 267 | + # ), |
| 268 | + # volume data stored in dictionary under pred_key + '_volumes' key |
| 269 | + ExtractVolumeToTextd( |
| 270 | + keys=[pred_key + "_volumes"], label_names=labels, output_labels=self.output_labels |
| 271 | + ), |
| 272 | + # comment out LabelToContourd for seg masks instead of contours; organ filtering will be lost |
| 273 | + LabelToContourd(keys=pred_key, output_labels=self.output_labels), |
| 274 | + OverlayImageLabeld(image_key=self._input_dataset_key, label_key=pred_key, overlay_key="overlay"), |
| 275 | + SaveImaged( |
| 276 | + keys="overlay", |
| 277 | + output_ext=".dcm", |
| 278 | + # save temporary DICOM SC (w/o source DICOM metadata) in output_folder / temp directory |
| 279 | + output_dir=self.output_folder / "temp", |
| 280 | + separate_folder=False, |
| 281 | + output_dtype=int16, |
| 282 | + ), |
| 283 | + ] |
| 284 | + ) |
| 285 | + |
| 286 | + # grab volume data from ExtractVolumeToTextd transform |
| 287 | + def get_result_text_from_transforms(self, post_transforms: Compose): |
| 288 | + """Extracts the result_text variables from post-processing transforms output.""" |
| 289 | + |
| 290 | + # grab the result_text variables from ExractVolumeToTextd transfor |
| 291 | + for transform in post_transforms.transforms: |
| 292 | + if isinstance(transform, ExtractVolumeToTextd): |
| 293 | + return transform.result_text_dicom_sr, transform.result_text_mongodb |
| 294 | + return None |
0 commit comments