Preprocessing cleanup (#41)

setup.py

@@ -5,7 +5,7 @@
         long_description = fh.read()
 
     setuptools.setup(
-        version='0.8.1',
+        version='0.8.2',
         author_email='Joeran.Bosma@radboudumc.nl',
         long_description=long_description,
         long_description_content_type="text/markdown",

src/picai_baseline/nnunet/training_docker/Dockerfile

@@ -31,9 +31,9 @@ RUN SITE_PKG=`pip3 show nnunet | grep "Location:" | awk '{print $2}'` && \
 
 # Install PI-CAI repositories
 RUN pip3 install \
-    picai_eval>=1.4.4 \
-    picai_prep>=2.1.2 \
-    picai_baseline>=0.8.1 \
+    picai_eval==1.4.4 \
+    picai_prep==2.1.2 \
+    picai_baseline==0.8.2 \
     mlxtend==0.19.0
 
 # copy custom files to nnUNet folder

src/picai_baseline/nnunet/training_docker/build.sh

@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
 
 docker build . \
-  --tag joeranbosma/picai_nnunet:1.7.0-customized-v1.12 \
+  --tag joeranbosma/picai_nnunet:1.7.0-customized-v1.13 \
   --tag joeranbosma/picai_nnunet:latest

src/picai_baseline/nnunet/training_docker/nnunet_wrapper.py

@@ -204,6 +204,9 @@ def plan_train(argv):
     parser.add_argument('--dont_plan_3d', action='store_true', help='Disable planning of 3D experiments')
     parser.add_argument('--carbontracker', action='store_true', help='Enables tracking of energy consumption')
     parser.add_argument('--pretrained_weights', type=str, required=False, default=None)
+    parser.add_argument('--disable_validation_inference', required=False, action='store_true', 
+                        help="If set nnU-Net will not run inference on the validation set. This is useful if you are only interested in the test set results and want to save some disk space and time.")
+    parser.add_argument('--dont_copy_preprocessed_data', action='store_true', help="Don't copy preprocessed data to datadir")
     args = parser.parse_args(argv)
 
     # aid type hinting
@@ -277,7 +280,7 @@ def plan_train(argv):
                     pickle.dump(splits, fp)
                 shutil_sol.copyfile(args.custom_split, splits_file.with_suffix('.json'))
 
-            if (prepdir / args.task).absolute() != taskdir.absolute():
+            if (prepdir / args.task).absolute() != taskdir.absolute() and not args.dont_copy_preprocessed_data:
                 # Copy preprocessed data to storage server
                 print('[#] Copying plans and preprocessed data from compute node to storage server')
                 taskdir.parent.mkdir(parents=True, exist_ok=True)
@@ -316,6 +319,8 @@ def plan_train(argv):
             cmd.append('--npz')
         if args.kwargs is not None:
             cmd.extend(args.kwargs.split(" "))
+        if args.disable_validation_inference:
+            cmd.append('--disable_validation_inference')
         print(f'[#] Running {" ".join(cmd)}')
 
         subprocess.check_call(cmd)

src/picai_baseline/nnunet_semi_supervised/training_sagemaker/code/preprocess.py

@@ -55,12 +55,13 @@ def main(taskname="Task2203_picai_baseline"):
     labels_dir = Path(args.labelsdir)
     output_dir = Path(args.outputdir)
     splits_path = workdir / f"nnUNet_raw_data/{taskname}/splits.json"
+    nnUNet_prep_dir = output_dir / "nnUNet_preprocessed"
 
     workdir.mkdir(parents=True, exist_ok=True)
     output_dir.mkdir(parents=True, exist_ok=True)
 
     # set environment variables
-    os.environ["prepdir"] = str(workdir / "nnUNet_preprocessed")
+    os.environ["prepdir"] = str(nnUNet_prep_dir)
 
     # set nnU-Net's number of preprocessing threads
     os.environ["nnUNet_tf"] = str(args.nnUNet_tf)
@@ -83,23 +84,21 @@ def main(taskname="Task2203_picai_baseline"):
         imagesdir=images_dir,
         labelsdir=labels_dir,
         preprocessing_kwargs='{"physical_size": [81.0, 192.0, 192.0], "crop_only": true}',
-        splits="picai_pubpriv",
+        splits=args.splits,
     )
 
     # Preprocess data with nnU-Net
     print("Preprocessing data with nnU-Net...")
     cmd = [
         "nnunet", "plan_train", str(taskname), workdir.as_posix(),
         "--custom_split", str(splits_path),
-        "--plan_only",
+        "--plan_only", "--dont_copy_preprocessed_data",
     ]
     check_call(cmd)
 
     # Export preprocessed dataset
-    print("Exporting preprocessed dataset...")
-    dst = output_dir / f"nnUNet_preprocessed/{taskname}/"
-    dst.parent.mkdir(parents=True, exist_ok=True)
-    shutil.copytree(workdir / f"nnUNet_preprocessed/{taskname}/", dst)
+    # nnU-Net will create a folder with the preprocessed data in the output directory.
+    # Make sure you export any files that you want to save to the output directory!
 
 
 if __name__ == '__main__':