save segmentations arg

bin/C2C

@@ -168,6 +168,7 @@ def argument_parser():
     base_parser = argparse.ArgumentParser(add_help=False)
     base_parser.add_argument("--input_path", "-i", type=str, required=True)
     base_parser.add_argument("--output_path", "-o", type=str)
+    base_parser.add_argument("--save_segmentations", action="store_true")
     base_parser.add_argument("--overwrite_outputs", action="store_true")
 
     parser = argparse.ArgumentParser()

comp2comp/utils/process.py

@@ -8,6 +8,7 @@
 from datetime import datetime
 from pathlib import Path
 from time import time
+import shutil
 
 from comp2comp.io.io_utils import get_dicom_or_nifti_paths_and_num
 
@@ -97,11 +98,16 @@ def process_3d(args, pipeline_builder):
 
             pipeline(output_dir=output_dir, model_dir=model_dir)
 
+
+            if not args.save_segmentations:
+                # remove the segmentations folder
+                segmentations_dir = os.path.join(output_dir, "segmentations")
+                if os.path.exists(segmentations_dir):
+                    shutil.rmtree(segmentations_dir)
+
             print(f"Finished processing {path} in {time() - st:.1f} seconds\n")
 
         except Exception:
             print(f"ERROR PROCESSING {path}\n")
             traceback.print_exc()
-            # if os.path.exists(output_dir):
-            #     shutil.rmtree(output_dir)
             continue