[~] small refacto

Author: Kirscher

README.md

@@ -1,6 +1,6 @@
 # PSAT
 
-This repository contains the code and configuration files for PSAT (Pediatric Segmentation Approaches via Adult Augmentations and Transfer).
+This repository contains the code and configuration files for PSAT (Pediatric Segmentation Approaches via Adult Augmentations and Transfer Learning).
 
 ## Overview
 
@@ -12,6 +12,42 @@ PSAT addresses pediatric segmentation challenges by combining:
 
 <img src="resources/images/PSAT_overview.png" alt="PSAT Overview" style="width:80%; max-width:1000px; display:block; margin: 0 auto;">
 
+## Citation
+If you use this code, please cite our paper:
+
+```
+@article{kirscher2025psat,
+  title={PSAT: Pediatric Segmentation Approaches via Adult Augmentations and Transfer Learning},
+  author={T. Kirscher et al},
+  journal={MICCAI},
+  year={2025},
+  note={arXiv:xxxx.xxxxx}
+}
+```
+
+## Quickstart
+
+Install dependencies:
+```bash
+pip install -r requirements.txt
+```
+
+Run metrics evaluation (example):
+```bash
+python scripts/compute_metrics.py <ground_truth_dir> <predictions_dir>
+```
+Replace `<ground_truth_dir>` and `<predictions_dir>` with your folder paths containing NIfTI files.
+
+## Dependencies
+- nibabel
+- numpy
+- pandas
+- p_tqdm
+- scipy
+- surface-distance
+
+(See `requirements.txt` for full list.)
+
 ## Documentation
 
 - [nnUNet](nnUNet/nnUNet.md)
@@ -28,4 +64,4 @@ Install dependencies listed in `requirements.txt` and run:
 
 ```bash
 pytest -q
-```
+```

scripts/compute_metrics.py

@@ -7,6 +7,25 @@
   and Hausdorff distances (https://github.com/deepmind/surface-distance).
 - Dynamic retrieval of pixel spacing to support varying image resolutions.
 - Reorientation checks to ensure proper alignment and spacing of NIfTI images.
+
+Usage:
+    python scripts/compute_metrics.py <ground_truth_dir> <predictions_dir>
+
+Arguments:
+    ground_truth_dir: Directory containing ground truth NIfTI files (*.nii.gz)
+    predictions_dir: Directory containing predicted NIfTI files (*.nii.gz)
+
+Expected file format:
+    - Each subject should have a file named <subject>.nii.gz in both directories.
+    - Files must be 3D or 4D NIfTI images with integer labels.
+
+Dependencies:
+    - nibabel
+    - numpy
+    - pandas
+    - p_tqdm
+    - scipy
+    - surface-distance
 """
 
 import sys
@@ -34,6 +53,11 @@
 def dice_score(y_true, y_pred):
     """
     Binary Dice score. Same results as sklearn f1 binary.
+    Args:
+        y_true (np.ndarray): Binary ground truth mask.
+        y_pred (np.ndarray): Binary predicted mask.
+    Returns:
+        float: Dice coefficient.
     """
     intersect = np.sum(y_true * y_pred)
     denominator = np.sum(y_true) + np.sum(y_pred)
@@ -42,11 +66,12 @@ def dice_score(y_true, y_pred):
 
 
 def reorient_to_ras(img):
-    """Reorient the image to RAS (Right-Anterior-Superior) orientation."""
+    """Reorient the image to RAS (Right-Anterior-Superior) orientation using nibabel."""
     return nib.as_closest_canonical(img)
 
 
 def calc_metrics(subject, gt_dir=None, pred_dir=None, class_map=None):
+    # Load ground truth and prediction images for a subject
     try:
         gt_img = nib.load(gt_dir / f"{subject}.nii.gz")
         pred_img = nib.load(pred_dir / f"{subject}.nii.gz")
@@ -70,6 +95,7 @@ def calc_metrics(subject, gt_dir=None, pred_dir=None, class_map=None):
         gt = gt_all == idx
         pred = pred_all == idx
 
+        # Handle cases where ground truth or prediction is missing for a class
         if gt.max() > 0 and pred.max() == 0:
             r[f"dice-{roi_name}"] = 0
             r[f"hausdorff-{roi_name}"] = 0
@@ -103,9 +129,12 @@ def calculate_confidence_interval(data, confidence=0.95):
     """
     Calculate Dice score and Hausdorff distance for your nnU-Net predictions.
 
-    example usage:
-    python evaluate.py ground_truth_dir predictions_dir
+    Example usage:
+        python scripts/compute_metrics.py <ground_truth_dir> <predictions_dir>
+
+    See the top-level docstring for more details.
     """
+    # Parse input arguments
     gt_dir = Path(sys.argv[1])
     pred_dir = Path(sys.argv[2])
 

scripts/create_totalseg_subset.py

@@ -0,0 +1,74 @@
+import logging
+from pathlib import Path
+import shutil
+import pandas as pd
+import matplotlib.pyplot as plt
+from typing import List
+
+def plot_and_save_distribution(df: pd.DataFrame, title: str, filename: str) -> None:
+    """Plot age and gender distribution and save to file.
+
+    Parameters
+    ----------
+    df : pandas.DataFrame
+        DataFrame containing at least ``age`` and ``gender`` columns.
+    title : str
+        Title for the plot.
+    filename : str
+        Output image file path.
+    """
+    fig, ax = plt.subplots()
+    df.boxplot(column="age", by="gender", ax=ax)
+    ax.set_title(title)
+    ax.set_xlabel("gender")
+    ax.set_ylabel("age")
+    plt.suptitle("")
+    fig.tight_layout()
+    fig.savefig(filename)
+    plt.close(fig)
+
+def create_directory_structure(base_path: Path, subdirs: List[str]) -> None:
+    """Create the required directory structure for the new dataset."""
+    base_path.mkdir(parents=True, exist_ok=True)
+    for subdir in subdirs:
+        (base_path / subdir).mkdir(parents=True, exist_ok=True)
+
+def copy_selected_files(
+    original_base: Path, new_base: Path, subdirs: List[str], selected_image_ids: List[str]
+) -> None:
+    """Copy files matching selected image IDs from the original dataset to the new dataset."""
+    for subdir in subdirs:
+        original_dir = original_base / subdir
+        new_dir = new_base / subdir
+
+        if not original_dir.exists():
+            logging.warning(f"Directory {original_dir} does not exist. Skipping...")
+            continue
+
+        for file_name in original_dir.iterdir():
+            if any(file_name.name.startswith(image_id) for image_id in selected_image_ids):
+                shutil.copy(file_name, new_dir / file_name.name)
+
+def main() -> None:
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+
+    # Paths
+    base_path = Path("nnUNet_raw_data_base/Dataset297_TotalSegmentator")
+    new_base_path = Path("nnUNet_raw_data_base/Dataset797_TotalSegmentator_plus_TCIA")
+    subdirs = ["imagesTr", "imagesTs", "labelsTr", "labelsTs"]
+
+    # Load the selected subset
+    subset_csv = Path("resources/TotalSegmentator/dataset_subset.csv")
+    selected_subset = pd.read_csv(subset_csv)
+    selected_image_ids = selected_subset["image_id"].tolist()
+
+    # Create new dataset directory structure
+    create_directory_structure(new_base_path, subdirs)
+
+    # Copy relevant files
+    copy_selected_files(base_path, new_base_path, subdirs, selected_image_ids)
+
+    logging.info(f"Subset dataset created successfully at: {new_base_path}")
+
+if __name__ == "__main__":
+    main()