diff --git a/README.md b/README.md index 7f85b67ad..1b4297723 100644 --- a/README.md +++ b/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. diff --git a/totalsegmentator/bin/TotalSegmentator.py b/totalsegmentator/bin/TotalSegmentator.py index 8cccd4c37..9aefe9797 100644 --- a/totalsegmentator/bin/TotalSegmentator.py +++ b/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", diff --git a/totalsegmentator/dicom_io.py b/totalsegmentator/dicom_io.py index 4f50477a4..bf352428b 100644 --- a/totalsegmentator/dicom_io.py +++ b/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) diff --git a/totalsegmentator/nnunet.py b/totalsegmentator/nnunet.py index a2325a3e4..b03157c18 100644 --- a/totalsegmentator/nnunet.py +++ b/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"