Add modality_list parameter to nnUNetExecutor and related functions

Author: SimoneBendazzoli93

monai/nvflare/nnunet_executor.py

@@ -77,6 +77,8 @@ class nnUNetExecutor(Executor):
         Extra configurations for training.
     exclude_vars : list, optional
         List of variables to exclude.
+    modality_list : list, optional
+        List of modalities.
 
     Methods
     -------
@@ -119,6 +121,7 @@ def __init__(
         tracking_uri=None,
         mlflow_token=None,
         bundle_root=None,
+        modality_list=None,
         train_extra_configs=None,
         exclude_vars=None,
     ):
@@ -143,6 +146,7 @@ def __init__(
         self.prepare_bundle_name = prepare_bundle_name
         self.bundle_root = bundle_root
         self.train_extra_configs = train_extra_configs
+        self.modality_list = modality_list
 
     def handle_event(self, event_type: str, fl_ctx: FLContext):
         if event_type == EventType.START_RUN:
@@ -211,6 +215,7 @@ def prepare_dataset(self) -> Shareable:
             mlflow_token=self.mlflow_token,
             subfolder_suffix=self.subfolder_suffix,
             trainer_class_name=nnunet_trainer_name,
+            modality_list=self.modality_list,
         )
 
         outgoing_dxo = DXO(data_kind=DataKind.COLLECTION, data=data_list, meta={})

monai/nvflare/nvflare_generate_job_configs.py

@@ -135,6 +135,9 @@ def prepare_config(clients, experiment, root_dir, script_dir, nvflare_exec):
             ],
         }
 
+        if "modality_list" in experiment["modality_list"]:
+            client["executors"][0]["executor"]["args"]["modality_list"] = experiment["modality_list"]
+
         if "subfolder_suffix" in clients[client_id]:
             client["executors"][0]["executor"]["args"]["subfolder_suffix"] = clients[client_id]["subfolder_suffix"]
         if "mlflow_token" in experiment:

monai/nvflare/nvflare_nnunet.py

@@ -305,6 +305,7 @@ def prepare_data_folder(
     experiment_name,
     client_name,
     dataset_format,
+    modality_list = None,
     tracking_uri=None,
     mlflow_token=None,
     subfolder_suffix=None,
@@ -332,6 +333,8 @@ def prepare_data_folder(
         Format of the dataset. Supported formats are "subfolders", "decathlon", and "nnunet".
     tracking_uri : str, optional
         URI for MLflow tracking server.
+    modality_list : list, optional
+        List of modalities. Default is None.
     mlflow_token : str, optional
         Token for MLflow authentication.
     subfolder_suffix : str, optional
@@ -438,9 +441,12 @@ def prepare_data_folder(
 
     os.makedirs(nnunet_root_dir, exist_ok=True)
 
+    if modality_list is None:
+        modality_list = [k for k in modality_dict.keys() if k != "label"]
+    
     data_src_cfg = os.path.join(nnunet_root_dir, "data_src_cfg.yaml")
     data_src = {
-        "modality": [k for k in modality_dict.keys() if k != "label"],
+        "modality": modality_list,
         "dataset_name_or_id": dataset_name_or_id,
         "datalist": str(datalist_file),
         "dataroot": str(data_dir),