run nnunet test with body task

wasserth · Nov 27, 2024 · 48d7818 · 48d7818
1 parent 1de4511
commit 48d7818
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/tests/tests_nnunet.py b/tests/tests_nnunet.py
@@ -19,7 +19,8 @@
 def run_tests_and_exit_on_failure():
 
     # Download weights
-    download_pretrained_weights(297)  # total 3mm
+    # download_pretrained_weights(297)  # total 3mm
+    download_pretrained_weights(300)  # body 6mm
 
     # Set nnUNet_results env var
     weights_dir = Path.home() / ".totalsegmentator" / "nnunet" / "results"
@@ -31,7 +32,10 @@ def run_tests_and_exit_on_failure():
     shutil.copy("tests/reference_files/example_ct_sm.nii.gz", "tests/nnunet_input_files/example_ct_sm_0000.nii.gz")
 
     # Run nnunet
-    subprocess.call(f"nnUNetv2_predict -i tests/nnunet_input_files -o tests/nnunet_input_files -d 297 -tr nnUNetTrainer_4000epochs_NoMirroring -c 3d_fullres -f 0 -device cpu", shell=True)
+    # Task 297
+    # subprocess.call(f"nnUNetv2_predict -i tests/nnunet_input_files -o tests/nnunet_input_files -d 297 -tr nnUNetTrainer_4000epochs_NoMirroring -c 3d_fullres -f 0 -device cpu", shell=True)
+    # Task 300
+    subprocess.call(f"nnUNetv2_predict -i tests/nnunet_input_files -o tests/nnunet_input_files -d 300 -tr nnUNetTrainer -c 3d_fullres -f 0 -device cpu", shell=True)
 
     # Check if output file exists
     assert os.path.exists("tests/nnunet_input_files/example_ct_sm.nii.gz"), "A nnunet output file was not generated."