Add original_path and model_name parameters to cross-site validation functions

Author: SimoneBendazzoli93

monai/nvflare/nnunet_executor.py

@@ -130,7 +130,8 @@ def __init__(
         train_extra_configs=None,
         exclude_vars=None,
         continue_training=False,
-        label_dict = None
+        label_dict = None,
+        original_path = None,
     ):
         super().__init__()
 
@@ -159,6 +160,7 @@ def __init__(
         self.cross_site_validation_task_name = cross_site_validation_task_name
         self.monai_deploy_config = monai_deploy_config
         self.label_dict = label_dict
+        self.original_path = original_path
 
     def handle_event(self, event_type: str, fl_ctx: FLContext):
         if event_type == EventType.START_RUN:
@@ -173,19 +175,20 @@ def initialize(self, fl_ctx: FLContext):
 
         with open("init_logfile_out.log", "w") as f_o:
             with open("init_logfile_err.log", "w") as f_e:
-                subprocess.call(
-                    [
-                        sys.executable,
-                        "-m",
-                        "pip",
-                        "install",
-                        "--user",
-                        "-r",
-                        str(Path(self.custom_app_dir).joinpath("requirements.txt")),
-                    ],
-                    stdout=f_o,
-                    stderr=f_e,
-                )
+                ...
+                #subprocess.call(
+                #    [
+                #        sys.executable,
+                #        "-m",
+                #        "pip",
+                #        "install",
+                #        "--user",
+                #        "-r",
+                #        str(Path(self.custom_app_dir).joinpath("requirements.txt")),
+                #    ],
+                #    stdout=f_o,
+                #    stderr=f_e,
+                #)
 
     def execute(self, task_name: str, shareable: Shareable, fl_ctx: FLContext, abort_signal: Signal) -> Shareable:
         self.run_dir = fl_ctx.get_engine().get_workspace().get_run_dir(fl_ctx.get_job_id())
@@ -223,6 +226,10 @@ def prepare_dataset(self) -> Shareable:
         dataset_name_or_id = self.nnunet_config["dataset_name_or_id"]
         if isinstance(dataset_name_or_id, dict):
             dataset_name_or_id = dataset_name_or_id.get(self.client_name, None)
+            if isinstance(dataset_name_or_id, dict):
+                dataset_name = dataset_name_or_id.get("name", None)
+                dataset_name_or_id = dataset_name_or_id.get("id", None)
+                
         data_list = prepare_data_folder(
             data_dir=self.data_dir,
             nnunet_root_dir=self.nnunet_root_folder,
@@ -237,6 +244,7 @@ def prepare_dataset(self) -> Shareable:
             subfolder_suffix=self.subfolder_suffix,
             trainer_class_name=nnunet_trainer_name,
             modality_list=self.modality_list,
+            dataset_name=dataset_name
         )
 
         outgoing_dxo = DXO(data_kind=DataKind.COLLECTION, data=data_list, meta={})
@@ -274,6 +282,10 @@ def plan_and_preprocess(self):
         dataset_name_or_id = self.nnunet_config["dataset_name_or_id"]
         if isinstance(dataset_name_or_id, dict):
             dataset_name_or_id = dataset_name_or_id.get(self.client_name, None)
+            if isinstance(dataset_name_or_id, dict):
+                dataset_name = dataset_name_or_id.get("name", None)
+                dataset_name_or_id = dataset_name_or_id.get("id", None)
+                
         nnunet_plans = plan_and_preprocess(
             self.nnunet_root_folder,
             dataset_name_or_id,
@@ -282,6 +294,7 @@ def plan_and_preprocess(self):
             self.tracking_uri,
             nnunet_plans_name=nnunet_plans_name,
             trainer_class_name=nnunet_trainer_name,
+            dataset_name=dataset_name,
         )
 
         outgoing_dxo = DXO(data_kind=DataKind.COLLECTION, data=nnunet_plans, meta={})
@@ -301,6 +314,10 @@ def preprocess(self):
         dataset_name_or_id = self.nnunet_config["dataset_name_or_id"]
         if isinstance(dataset_name_or_id, dict):
             dataset_name_or_id = dataset_name_or_id.get(self.client_name, None)
+            if isinstance(dataset_name_or_id, dict):
+                dataset_name = dataset_name_or_id.get("name", None)
+                dataset_name_or_id = dataset_name_or_id.get("id", None)
+                
         nnunet_plans = preprocess(
             self.nnunet_root_folder,
             dataset_name_or_id,
@@ -324,6 +341,10 @@ def train(self):
         dataset_name_or_id = self.nnunet_config["dataset_name_or_id"]
         if isinstance(dataset_name_or_id, dict):
             dataset_name_or_id = dataset_name_or_id.get(self.client_name, None)
+            if isinstance(dataset_name_or_id, dict):
+                dataset_name = dataset_name_or_id.get("name", None)
+                dataset_name_or_id = dataset_name_or_id.get("id", None)
+                
         validation_summary = train(
             self.nnunet_root_folder,
             trainer_class_name=nnunet_trainer_name,
@@ -354,6 +375,10 @@ def prepare_bundle(self):
         dataset_name_or_id = self.nnunet_config["dataset_name_or_id"]
         if isinstance(dataset_name_or_id, dict):
             dataset_name_or_id = dataset_name_or_id.get(self.client_name, None)
+            if isinstance(dataset_name_or_id, dict):
+                dataset_name = dataset_name_or_id.get("name", None)
+                dataset_name_or_id = dataset_name_or_id.get("id", None)
+                
         bundle_config = {
             "bundle_root": self.bundle_root,
             "tracking_uri": self.tracking_uri,
@@ -363,6 +388,7 @@ def prepare_bundle(self):
             "nnunet_trainer_class_name": nnunet_trainer_name,
             "dataset_name_or_id": dataset_name_or_id,
             "label_dict": self.label_dict,
+            "dataset_name": dataset_name,
         }
 
         bundle_config = prepare_bundle(bundle_config, self.train_extra_configs)
@@ -386,6 +412,10 @@ def finalize_bundle(self):
         dataset_name_or_id = self.nnunet_config["dataset_name_or_id"]
         if isinstance(dataset_name_or_id, dict):
             dataset_name_or_id = dataset_name_or_id.get(self.client_name, None)
+            if isinstance(dataset_name_or_id, dict):
+                dataset_name = dataset_name_or_id.get("name", None)
+                dataset_name_or_id = dataset_name_or_id.get("id", None)
+                
         validation_summary = finalize_bundle(
             self.bundle_root,
             self.nnunet_root_folder,
@@ -395,7 +425,8 @@ def finalize_bundle(self):
             client_name=self.client_name,
             tracking_uri=self.tracking_uri,
             nnunet_plans_name=nnunet_plans_name,
-            dataset_name_or_id=dataset_name_or_id
+            dataset_name_or_id=dataset_name_or_id,
+            dataset_name=dataset_name,
         )
         outgoing_dxo = DXO(data_kind=DataKind.COLLECTION, data=validation_summary, meta={})
         return outgoing_dxo.to_shareable()
@@ -414,18 +445,25 @@ def run_cross_site_validation(self):
         dataset_name_or_id = self.nnunet_config["dataset_name_or_id"]
         if isinstance(dataset_name_or_id, dict):
             dataset_name_or_id = dataset_name_or_id.get(self.client_name, None)
+            if isinstance(dataset_name_or_id, dict):
+                dataset_name = dataset_name_or_id.get("name", None)
+                dataset_name_or_id = dataset_name_or_id.get("id", None)
+                
         validation_summary = run_cross_site_validation(
             self.nnunet_root_folder,
             dataset_name_or_id,
             self.monai_deploy_config["app_path"],
             self.monai_deploy_config["app_model_path"],
             self.monai_deploy_config["app_output_path"],
+            self.monai_deploy_config["model_name"],
             trainer_class_name=nnunet_trainer_name,
             fold=0,
             experiment_name=self.nnunet_config["experiment_name"],
             client_name=self.client_name,
             tracking_uri=self.tracking_uri,
-            nnunet_plans_name=nnunet_plans_name
+            nnunet_plans_name=nnunet_plans_name,
+            dataset_name=dataset_name,
+            original_path=self.original_path,
         )
 
         outgoing_dxo = DXO(data_kind=DataKind.COLLECTION, data=validation_summary, meta={})

monai/nvflare/nvflare_generate_job_configs.py

@@ -1322,6 +1322,7 @@ def cross_site_validation_config(clients, experiment, root_dir, script_dir, nvfl
                             "monai_deploy_config": {
                                 "app_path": clients[client_id]["app_path"],
                                 "app_model_path": clients[client_id]["app_model_path"],
+                                "model_name": clients[client_id]["model_name"],
                                 "app_output_path": clients[client_id]["app_output_path"],
                             },
                             "client_name": clients[client_id]["client_name"],
@@ -1341,6 +1342,11 @@ def cross_site_validation_config(clients, experiment, root_dir, script_dir, nvfl
         if "bundle_root" in clients[client_id]:
             client["executors"][0]["executor"]["args"]["bundle_root"] = clients[client_id]["bundle_root"]
 
+        if "original_path" in clients[client_id]:
+            client["executors"][0]["executor"]["args"]["original_path"] = clients[client_id][
+                "original_path"
+            ]
+
         Path(root_dir).joinpath(task_name).joinpath(f"{task_name}-client-{client_id}").mkdir(
             parents=True, exist_ok=True
         )

monai/nvflare/nvflare_nnunet.py

@@ -541,7 +541,7 @@ def prepare_bundle(bundle_config, train_extra_configs=None):
     None
     """
 
-    prepare_bundle_api(bundle_config, train_extra_configs=train_extra_configs, is_federated=True)
+    return prepare_bundle_api(bundle_config, train_extra_configs=train_extra_configs, is_federated=True)
 
 
 
@@ -644,9 +644,9 @@ def finalize_bundle(bundle_root, nnunet_root_dir=None, validate_with_nnunet=True
         return validation_summary_dict
 
 
-def run_cross_site_validation(nnunet_root_dir, dataset_name_or_id, app_path, app_model_path, app_output_path, trainer_class_name="nnUNetTrainer", fold=0,
+def run_cross_site_validation(nnunet_root_dir, dataset_name_or_id, app_path, app_model_path, app_output_path, model_name, trainer_class_name="nnUNetTrainer", fold=0,
                     experiment_name=None, client_name=None, tracking_uri=None,
-                    nnunet_plans_name="nnUNetPlans", mlflow_token=None, skip_prediction=False, dataset_name=None):
+                    nnunet_plans_name="nnUNetPlans", mlflow_token=None, skip_prediction=False, dataset_name=None, original_path=None):
 
     validation_summary_dict, labels = cross_site_evaluation_api(
         nnunet_root_dir,
@@ -658,6 +658,8 @@ def run_cross_site_validation(nnunet_root_dir, dataset_name_or_id, app_path, app
         fold=fold,
         nnunet_plans_name=nnunet_plans_name,
         skip_prediction=skip_prediction,
+        original_path = original_path,
+
     )
     if mlflow_token is not None:
         os.environ["MLFLOW_TRACKING_TOKEN"] = mlflow_token
@@ -670,28 +672,28 @@ def run_cross_site_validation(nnunet_root_dir, dataset_name_or_id, app_path, app
         print(e)
         mlflow.set_experiment(experiment_id=(mlflow.get_experiment_by_name(experiment_name).experiment_id))
 
-    run_name = f"run_cross_site_validation_{client_name}"
+    run_name = f"run_cross_site_validation_{client_name}_Model_{model_name}"
 
     runs = mlflow.search_runs(
     experiment_names=[experiment_name],
     filter_string=f"tags.mlflow.runName = '{run_name}'",
     order_by=["start_time DESC"]
     )
-    tags = {"client": client_name}
+    tags = {"client": client_name,"model": model_name}
     if dataset_name is not None:
         tags["dataset_name"] = dataset_name
 
 
     if len(runs) == 0:
-        with mlflow.start_run(run_name=f"run_{client_name}", tags={"client": client_name}):
+        with mlflow.start_run(run_name=run_name, tags=tags):
             mlflow.log_dict(validation_summary_dict, "validation_summary.json")
             for label in validation_summary_dict["mean"]:
                 for metric in validation_summary_dict["mean"][label]:
                     label_name = labels[label]
                     mlflow.log_metric(f"{label_name}_{metric}", float(validation_summary_dict["mean"][label][metric]))
 
     else:
-        with mlflow.start_run(run_id=runs.iloc[0].run_id, tags={"client": client_name}):
+        with mlflow.start_run(run_id=runs.iloc[0].run_id, tags=tags):
             mlflow.log_dict(validation_summary_dict, "validation_summary.json")
             for label in validation_summary_dict["mean"]:
                 for metric in validation_summary_dict["mean"][label]:

monai/nvflare/utils.py

@@ -256,7 +256,7 @@ def prepare_data_folder_api(data_dir,
     return data_list
 
 
-def cross_site_evaluation_api(nnunet_root_dir, dataset_name_or_id, app_path, app_model_path, app_output_path, fold=0, trainer_class_name="nnUNetTrainer", nnunet_plans_name="nnUNetPlans", skip_prediction=False):
+def cross_site_evaluation_api(nnunet_root_dir, dataset_name_or_id, app_path, app_model_path, app_output_path, fold=0, trainer_class_name="nnUNetTrainer", nnunet_plans_name="nnUNetPlans", skip_prediction=False, original_path=None):
     data_src_cfg = os.path.join(nnunet_root_dir, f"Task{dataset_name_or_id}_data_src_cfg.yaml")
 
     runner = nnUNetV2Runner(input_config=data_src_cfg, trainer_class_name=trainer_class_name, work_dir=nnunet_root_dir)
@@ -285,14 +285,28 @@ def cross_site_evaluation_api(nnunet_root_dir, dataset_name_or_id, app_path, app
         filename = Path(app_input_path).name
 
         new_id = None  
+        updated_image_path = False
         for case in nnunet_datalist["training"]:
             if case["image"].endswith(filename):
                 new_id = case["new_name"]
                 break
+            if filename.startswith(case["new_name"]+"_"):
+                new_id = case["new_name"]
+                data["image"] = os.path.join(original_path, Path(case["image"]).name)
+                app_input_path = data["image"]
+                updated_image_path = True
+                break
         if new_id in nnunet_splits[fold]["val"]:
-            id_mapping[Path(data["image"]).name.split("_")[0].split(".")[0]] = new_id
+            if updated_image_path:
+                id_mapping[Path(data["image"]).name.split(".")[0]] = new_id
+            else:
+                id_mapping[Path(data["image"]).name.split("_")[0].split(".")[0]] = new_id
             if skip_prediction:
                 continue
+            print(f"Processing case: {new_id}")
+            print(f"App input path: {app_input_path}")
+            mapped_filename = Path(data["image"]).name.split(".")[0]
+            print(f"Mapping: {mapped_filename} -> {new_id}")
             subprocess.run(
                 [
                 "python",