feat: Added benchmark setup script (WIP)

Author: Karol-G

scribblebench/setup_benchmark.py

@@ -7,6 +7,8 @@
 import argparse
 from git import Repo
 import os
+from utils.download_kits23 import download_dataset
+from natsort import natsorted
 
 
 def setup_word_dataset(dataset_dir):
@@ -146,7 +148,7 @@ def setup_mscmr_dataset(dataset_dir):
         zip_ref.extractall(mscmr_preprocessed_dir)
 
     ####################################################################################################################
-    #### Preprocess WORD dataset
+    #### Preprocess MSCMR dataset
     ####################################################################################################################
 
     print("Preprocessing MSCMR dataset...")
@@ -171,7 +173,7 @@ def setup_mscmr_dataset(dataset_dir):
 
     names = [path.name[:-7] for path in (mscmr_raw_dir / "TestSet" / "labels").rglob("*.nii.gz")]
     for name in names:
-        shutil.move(mscmr_raw_dir / "TestSet" / "labels" / f"{name}.nii.gz", mscmr_preprocessed_dir / "labelsTs" / f"{name}_0000.nii.gz")
+        shutil.move(mscmr_raw_dir / "TestSet" / "labels" / f"{name}.nii.gz", mscmr_preprocessed_dir / "labelsTs" / f"{name}.nii.gz")
 
     # These two images have no dense GT so it is not possible to generate scribbles for them
     os.remove(mscmr_preprocessed_dir / "imagesTr" / "subject2_DE_0000.nii.gz")
@@ -194,14 +196,90 @@ def setup_mscmr_dataset(dataset_dir):
     print("Finished setting up MSCMR dataset.")
 
 
+def setup_kits_dataset(dataset_dir):
+    dataset_dir = Path(dataset_dir) / "ScribbleBench"
+    raw_dir = dataset_dir / "raw"
+    kits_raw_dir = raw_dir / "KiTS2023" / "dataset"
+    preprocessed_dir = dataset_dir
+    kits_preprocessed_dir = preprocessed_dir / "KiTS2023"
+    preprocessed_dir.mkdir(parents=True, exist_ok=True)
+    kits_preprocessed_dir.mkdir(parents=True, exist_ok=True)
+
+    test_set = ['case_00007', 'case_00013', 'case_00003', 'case_00018', 'case_00019', 'case_00016', 'case_00026', 'case_00014', 
+                'case_00038', 'case_00000', 'case_00045', 'case_00050', 'case_00061', 'case_00070', 'case_00069', 'case_00074', 
+                'case_00087', 'case_00090', 'case_00086', 'case_00084', 'case_00059', 'case_00106', 'case_00105', 'case_00102', 
+                'case_00112', 'case_00117', 'case_00114', 'case_00121', 'case_00101', 'case_00096', 'case_00115', 'case_00126', 
+                'case_00127', 'case_00120', 'case_00124', 'case_00116', 'case_00133', 'case_00145', 'case_00147', 'case_00152', 
+                'case_00144', 'case_00118', 'case_00132', 'case_00135', 'case_00141', 'case_00146', 'case_00164', 'case_00167', 
+                'case_00172', 'case_00179', 'case_00181', 'case_00184', 'case_00192', 'case_00194', 'case_00195', 'case_00165', 
+                'case_00199', 'case_00210', 'case_00188', 'case_00191', 'case_00211', 'case_00212', 'case_00222', 'case_00217', 
+                'case_00221', 'case_00227', 'case_00232', 'case_00236', 'case_00231', 'case_00234', 'case_00214', 'case_00238', 
+                'case_00223', 'case_00237', 'case_00240', 'case_00216', 'case_00266', 'case_00269', 'case_00268', 'case_00275', 
+                'case_00279', 'case_00253', 'case_00273', 'case_00282', 'case_00287', 'case_00286', 'case_00281', 'case_00284', 
+                'case_00291', 'case_00283', 'case_00276', 'case_00404', 'case_00411', 'case_00416', 'case_00418', 'case_00422', 
+                'case_00423', 'case_00290', 'case_00424', 'case_00295', 'case_00426', 'case_00428', 'case_00293', 'case_00429', 
+                'case_00433', 'case_00441', 'case_00443', 'case_00437', 'case_00444', 'case_00452', 'case_00449', 'case_00453', 
+                'case_00463', 'case_00468', 'case_00476', 'case_00483', 'case_00485', 'case_00480', 'case_00491', 'case_00474', 
+                'case_00486', 'case_00496', 'case_00494', 'case_00492', 'case_00503', 'case_00442', 'case_00518', 'case_00521', 
+                'case_00522', 'case_00525', 'case_00515', 'case_00533', 'case_00532', 'case_00539', 'case_00517', 'case_00546', 
+                'case_00550', 'case_00554', 'case_00558', 'case_00557', 'case_00544', 'case_00567', 'case_00574', 'case_00555', 
+                'case_00576', 'case_00575', 'case_00564']
+
+    ####################################################################################################################
+    #### Download KiTS2023 dataset
+    ####################################################################################################################
+
+    print("Downloading KiTS2023 dataset...")
+    repo_url = "https://github.com/neheller/kits23.git"
+    Repo.clone_from(repo_url, str(kits_raw_dir.parent))
+    download_dataset(kits_raw_dir)
+
+    ####################################################################################################################
+    #### Preprocess KiTS2023 dataset
+    ####################################################################################################################
+
+    print("Preprocessing KiTS2023 dataset...")
+
+    (kits_preprocessed_dir / "imagesTr").mkdir(parents=True, exist_ok=True)
+    (kits_preprocessed_dir / "imagesTs").mkdir(parents=True, exist_ok=True)
+    (kits_preprocessed_dir / "labelsTr").mkdir(parents=True, exist_ok=True)
+    (kits_preprocessed_dir / "labelsTs").mkdir(parents=True, exist_ok=True)
+
+    names = [p.name for p in kits_raw_dir.iterdir() if p.is_dir()]
+    names = natsorted(names)
+    for name in names:
+        postfix = "Tr" if name not in test_set else "Ts"
+        shutil.move(kits_raw_dir / name / "imaging.nii.gz", kits_preprocessed_dir / f"images{postfix}" / f"{name}_0000.nii.gz")
+        shutil.move(kits_raw_dir / name / "segmentation.nii.gz", kits_preprocessed_dir / f"labels{postfix}" / f"{name}.nii.gz")
+
+    dataset_json_url = "https://syncandshare.desy.de/index.php/s/Cfpwyg5dmi9a2Df/download/dataset.json"
+    response = requests.get(dataset_json_url)
+    response.raise_for_status()  # Raise an error on bad status
+    with open(kits_preprocessed_dir / "dataset.json", "wb") as f:
+        f.write(response.content)
+
+    ####################################################################################################################
+    #### Delete raw dataset files
+    ####################################################################################################################
+
+    print("Deleting raw dataset files...")
+    shutil.rmtree(raw_dir, ignore_errors=True)
+
+    print("Finished setting up KiTS2023 dataset.")
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('-d', "--dataset_dir", required=True, type=str, help="Path to the dir used for setting up ScribbleBench.")
     parser.add_argument('--word', required=False, default=False, action="store_true", help="Download and preprocess the WORD dataset for ScribbleBench.")
     parser.add_argument('--mscmr', required=False, default=False, action="store_true", help="Download and preprocess the MSCMR dataset for ScribbleBench.")
+    parser.add_argument('--kits', required=False, default=False, action="store_true", help="Download and preprocess the KiTS2023 dataset for ScribbleBench.")
     args = parser.parse_args()
 
     if args.word:
         setup_word_dataset(args.dataset_dir)
     if args.word:
         setup_mscmr_dataset(args.dataset_dir)
+    if args.kits:
+        setup_kits_dataset(args.dataset_dir)
+        

scribblebench/utils/__init__.py

@@ -0,0 +1,78 @@
+"""A script to download the KiTS23 dataset into this repository"""
+import sys
+from tqdm import tqdm
+from pathlib import Path
+import urllib.request
+import shutil
+from time import sleep
+
+TRAINING_CASE_NUMBERS = list(range(300)) + list(range(400, 589))
+
+
+def get_destination(case_id: str, save_dir, create: bool = False):
+    destination = save_dir / case_id / "imaging.nii.gz"
+    if create:
+        destination.parent.mkdir(exist_ok=True)
+    return destination
+
+
+def cleanup(tmp_pth: Path, e: Exception):
+    if tmp_pth.exists():
+        tmp_pth.unlink()
+
+    if e is None:
+        print("\nInterrupted.\n")
+        sys.exit()
+    raise(e)
+
+
+def download_case(case_num: int, save_dir, pbar: tqdm, retry=True):
+    remote_name = f"master_{case_num:05d}.nii.gz"
+    url = f"https://kits19.sfo2.digitaloceanspaces.com/{remote_name}"
+    destination = get_destination(f"case_{case_num:05d}", save_dir, True)
+    tmp_pth = destination.parent / f".partial.{destination.name}"
+    try:
+        urllib.request.urlretrieve(url, str(tmp_pth))
+        shutil.move(str(tmp_pth), str(destination))
+    except KeyboardInterrupt as e:
+        pbar.close()
+        while True:
+            try:
+                sleep(0.1)
+                cleanup(tmp_pth, None)
+            except KeyboardInterrupt:
+                pass
+    except Exception as e:
+        if retry:
+            print(f"\nFailed to download case_{case_num:05d}. Retrying...")
+            sleep(5)
+            download_case(case_num, pbar, retry=False)
+        pbar.close()
+        while True:
+            try:
+                cleanup(tmp_pth, e)
+            except KeyboardInterrupt:
+                pass
+
+
+def download_dataset(save_dir):
+    save_dir = Path(save_dir)
+    save_dir.mkdir(parents=True, exist_ok=True)
+
+    # Determine which cases still need to be downloaded
+    left_to_download = []
+    for case_num in TRAINING_CASE_NUMBERS:
+        case_id = f"case_{case_num:05d}"
+        dst = get_destination(case_id, save_dir)
+        if not dst.exists():
+            left_to_download = left_to_download + [case_num]
+
+    # Show progressbar as cases are downloaded
+    print(f"\nFound {len(left_to_download)} cases to download\n")
+    for case_num in (pbar := tqdm(left_to_download)):
+        pbar.set_description(f"Dowloading case_{case_num:05d}...")
+        download_case(case_num, save_dir, pbar)
+
+
+if __name__ == "__main__":
+    download_dataset()

scribblebench/utils/download_kits23.py

@@ -36,6 +36,7 @@ install_requires =
     medvol
     gdown
     GitPython
+    natsort
 python_requires = >=3.8
 include_package_data = True
 ; package_dir =