allow dicom zip files as input

README.md

@@ -41,7 +41,7 @@ For MR images:
 ```
 TotalSegmentator -i mri.nii.gz -o segmentations --task total_mr
 ```
-> Note: A Nifti file or a folder with all DICOM slices of one patient is allowed as input.
+> Note: A Nifti file or a folder (or zip file) with all DICOM slices of one patient is allowed as input.
 
 > Note: If you run on CPU use the option `--fast` or `--roi_subset` to greatly improve runtime.
 

totalsegmentator/bin/TotalSegmentator.py

@@ -13,7 +13,7 @@ def main():
                                      epilog="Written by Jakob Wasserthal. If you use this tool please cite https://pubs.rsna.org/doi/10.1148/ryai.230024")
 
     parser.add_argument("-i", metavar="filepath", dest="input",
-                        help="CT nifti image or folder of dicom slices",
+                        help="CT nifti image or folder of dicom slices or zip file of dicom slices.",
                         type=lambda p: Path(p).absolute(), required=True)
 
     parser.add_argument("-o", metavar="directory", dest="output",

totalsegmentator/dicom_io.py

@@ -111,13 +111,23 @@ def dcm_to_nifti_LEGACY(input_path, output_path, verbose=False):
     os.remove(str(output_path)[:-7] + ".json")
 
 
-def dcm_to_nifti(input_path, output_path, verbose=False):
+def dcm_to_nifti(input_path, output_path, tmp_dir=None, verbose=False):
     """
     Uses dicom2nifti package (also works on windows)
 
-    input_path: a directory of dicom slices
+    input_path: a directory of dicom slices or a zip file of dicom slices
     output_path: a nifti file path
+    tmp_dir: extract zip file to this directory, else to the same directory as the zip file
     """
+    # Check if input_path is a zip file and extract it
+    if zipfile.is_zipfile(input_path):
+        if verbose: print(f"Extracting zip file: {input_path}")
+        extract_dir = os.path.splitext(input_path)[0] if tmp_dir is None else tmp_dir / "extracted_dcm"
+        with zipfile.ZipFile(input_path, 'r') as zip_ref:
+            zip_ref.extractall(extract_dir)
+            input_path = extract_dir
+
+    # Convert to nifti
     dicom2nifti.dicom_series_to_nifti(input_path, output_path, reorient_nifti=True)
 
 

totalsegmentator/nnunet.py

@@ -295,7 +295,10 @@ def nnUNet_predict_image(file_in: Union[str, Path, Nifti1Image], file_out, task_
     """
     if not isinstance(file_in, Nifti1Image):
         file_in = Path(file_in)
-        img_type = "nifti" if str(file_in).endswith(".nii") or str(file_in).endswith(".nii.gz") else "dicom"
+        if str(file_in).endswith(".nii") or str(file_in).endswith(".nii.gz"):
+            img_type = "nifti"
+        else:
+            img_type = "dicom"
         if not file_in.exists():
             sys.exit("ERROR: The input file or directory does not exist.")
     else:
@@ -325,7 +328,7 @@ def nnUNet_predict_image(file_in: Union[str, Path, Nifti1Image], file_out, task_
         if img_type == "dicom":
             if not quiet: print("Converting dicom to nifti...")
             (tmp_dir / "dcm").mkdir()  # make subdir otherwise this file would be included by nnUNet_predict
-            dcm_to_nifti(file_in, tmp_dir / "dcm" / "converted_dcm.nii.gz", verbose=verbose)
+            dcm_to_nifti(file_in, tmp_dir / "dcm" / "converted_dcm.nii.gz", tmp_dir, verbose=verbose)
             file_in_dcm = file_in
             file_in = tmp_dir / "dcm" / "converted_dcm.nii.gz"