From ce6eb7771a8b1e4eb915224861c4cf2a92e7149f Mon Sep 17 00:00:00 2001
From: erasdna <solheim@simula.no>
Date: Fri, 13 Mar 2026 10:44:57 +0100
Subject: [PATCH 1/6] add segmentation class, rework statistics dataframe

---
 src/mritk/data/base.py                 |  18 ++
 src/mritk/data/io.py                   |  62 +++---
 src/mritk/segmentation/__init__.py     |   4 +-
 src/mritk/segmentation/lookup_table.py |  32 ++-
 src/mritk/segmentation/segmentation.py |  43 ++++
 src/mritk/statistics/__init__.py       |   4 +-
 src/mritk/statistics/cli.py            | 133 ++++++------
 src/mritk/statistics/compute_stats.py  | 155 ++++----------
 src/mritk/statistics/metadata.py       |  45 ++++
 src/mritk/statistics/stat_functions.py |  53 +++++
 src/mritk/statistics/utils.py          |   5 +-
 test/conftest.py                       |  28 +++
 test/test_metadata.py                  |  27 +++
 test/test_mri_io.py                    |  45 +++-
 test/test_mri_stats.py                 | 271 ++++++++++++++-----------
 15 files changed, 575 insertions(+), 350 deletions(-)
 create mode 100644 src/mritk/segmentation/segmentation.py
 create mode 100644 src/mritk/statistics/metadata.py
 create mode 100644 src/mritk/statistics/stat_functions.py
 create mode 100644 test/test_metadata.py

diff --git a/src/mritk/data/base.py b/src/mritk/data/base.py
index 7880e7d..08e8b53 100644
--- a/src/mritk/data/base.py
+++ b/src/mritk/data/base.py
@@ -6,12 +6,23 @@
 
 
 import numpy as np
+from pathlib import Path
+from .io import load_mri_data, save_mri_data
 
 
 class MRIData:
     def __init__(self, data: np.ndarray, affine: np.ndarray):
         self.data = data
         self.affine = affine
+        self.dtype = data.dtype
+
+    @classmethod
+    def from_file(cls, filepath: Path):
+        data, affine = load_mri_data(filepath, np.float64)
+        return cls(data=data, affine=affine)
+
+    def save_mri_data(self, save_path: Path, intent_code: int | None = None):
+        save_mri_data(self.data, self.affine, save_path, self.dtype, intent_code)
 
     def get_data(self):
         return self.data
@@ -22,3 +33,10 @@ def get_metadata(self):
     @property
     def shape(self) -> tuple[int, ...]:
         return self.data.shape
+
+    @property
+    def voxel_ml_volume(self) -> float:
+        # Calculate the volume of a single voxel in milliliters
+        voxel_volume_mm3 = abs(np.linalg.det(self.affine[:3, :3]))
+        voxel_volume_ml = voxel_volume_mm3 / 1000.0  # Convert from mm^3 to ml
+        return voxel_volume_ml
diff --git a/src/mritk/data/io.py b/src/mritk/data/io.py
index ba3547c..601ac42 100644
--- a/src/mritk/data/io.py
+++ b/src/mritk/data/io.py
@@ -9,53 +9,47 @@
 import nibabel
 import numpy as np
 import numpy.typing as npt
-import re
 from typing import Optional
 
-from .base import MRIData
-from .orientation import data_reorientation
-
-
-def load_mri_data(
-    path: Path | str,
-    dtype: type = np.float64,
-    orient: bool = True,
-) -> MRIData:
-    suffix_regex = re.compile(r".+(?P<suffix>(\.nii(\.gz|)|\.mg(z|h)))")
-    m = suffix_regex.match(Path(path).name)
-    if (m is not None) and (m.groupdict()["suffix"] in (".nii", ".nii.gz")):
-        mri = nibabel.nifti1.load(path)
-    elif (m is not None) and (m.groupdict()["suffix"] in (".mgz", ".mgh")):
-        mri = nibabel.freesurfer.mghformat.load(path)
+
+def check_suffix(filepath: Path):
+    suffix = filepath.suffix
+    if suffix == ".gz":
+        suffixes = filepath.suffixes
+        if len(suffixes) >= 2 and suffixes[-2] == ".nii":
+            return ".nii.gz"
+    return suffix
+
+
+def load_mri_data(filepath: Path, dtype: type = np.float64) -> tuple[np.ndarray, np.ndarray]:
+    suffix = check_suffix(filepath)
+    if suffix in (".nii", ".nii.gz"):
+        mri = nibabel.nifti1.load(filepath)
+    elif suffix in (".mgz", ".mgh"):
+        mri = nibabel.freesurfer.mghformat.load(filepath)
     else:
-        raise ValueError(f"Invalid suffix {path}, should be either '.nii', or '.mgz'")
+        raise ValueError(f"Invalid suffix {filepath}, should be either '.nii', or '.mgz'")
 
     affine = mri.affine
     if affine is None:
-        raise RuntimeError("MRI do not contain affine")
+        raise RuntimeError("MRI does not contain an affine")
 
     data = np.asarray(mri.get_fdata("unchanged"), dtype=dtype)
-    mri = MRIData(data=data, affine=affine)
 
-    if orient:
-        return data_reorientation(mri)
-    else:
-        return mri
+    return data, affine
 
 
-def save_mri_data(mri: MRIData, path: Path, dtype: npt.DTypeLike, intent_code: Optional[int] = None):
-    # TODO : Choose other way to check extension than regex ?
-    suffix_regex = re.compile(r".+(?P<suffix>(\.nii(\.gz|)|\.mg(z|h)))")
-    m = suffix_regex.match(Path(path).name)
-    if (m is not None) and (m.groupdict()["suffix"] in (".nii", ".nii.gz")):
-        nii = nibabel.nifti1.Nifti1Image(mri.data.astype(dtype), mri.affine)
+def save_mri_data(data: np.ndarray, affine: np.ndarray, save_path: Path, dtype: npt.DTypeLike, intent_code: Optional[int] = None):
+    suffix = check_suffix(save_path)
+    if suffix in (".nii", ".nii.gz"):
+        nii = nibabel.nifti1.Nifti1Image(data.astype(dtype), affine)
         if intent_code is not None:
             nii.header.set_intent(intent_code)
-        nibabel.nifti1.save(nii, path)
-    elif (m is not None) and (m.groupdict()["suffix"] in (".mgz", ".mgh")):
-        mgh = nibabel.freesurfer.mghformat.MGHImage(mri.data.astype(dtype), mri.affine)
+        nibabel.nifti1.save(nii, save_path)
+    elif suffix in (".mgz", ".mgh"):
+        mgh = nibabel.freesurfer.mghformat.MGHImage(data.astype(dtype), affine)
         if intent_code is not None:
             mgh.header.set_intent(intent_code)
-        nibabel.freesurfer.mghformat.save(mgh, path)
+        nibabel.freesurfer.mghformat.save(mgh, save_path)
     else:
-        raise ValueError(f"Invalid suffix {path}, should be either '.nii', or '.mgz'")
+        raise ValueError(f"Invalid suffix {save_path}, should be either '.nii', or '.mgz'")
diff --git a/src/mritk/segmentation/__init__.py b/src/mritk/segmentation/__init__.py
index 67400a7..689410d 100644
--- a/src/mritk/segmentation/__init__.py
+++ b/src/mritk/segmentation/__init__.py
@@ -3,6 +3,6 @@
 # Copyright (C) 2026   Simula Research Laboratory
 
 
-from . import groups, lookup_table
+from . import groups, lookup_table, segmentation
 
-__all__ = ["groups", "lookup_table"]
+__all__ = ["groups", "lookup_table", "segmentation"]
diff --git a/src/mritk/segmentation/lookup_table.py b/src/mritk/segmentation/lookup_table.py
index 9d55348..dabe313 100644
--- a/src/mritk/segmentation/lookup_table.py
+++ b/src/mritk/segmentation/lookup_table.py
@@ -5,7 +5,6 @@
 # Copyright (C) 2026   Simula Research Laboratory
 
 
-import re
 import os
 from pathlib import Path
 import pandas as pd
@@ -23,12 +22,21 @@ def read_lut(filename: Path | str | None) -> pd.DataFrame:
         if not filename.exists():
             url = "https://github.com/freesurfer/freesurfer/raw/dev/distribution/FreeSurferColorLUT.txt"
             urlretrieve(url, filename)
-    lut_regex = re.compile(
-        r"^(?P<label>\d+)\s+(?P<description>[_\da-zA-Z-]+)\s+(?P<R>\d+)\s+(?P<G>\d+)\s+(?P<B>\d+)\s+(?P<A>\d+)"
-    )
+    # lut_regex = re.compile(
+    #     r"^(?P<label>\d+)\s+(?P<description>[_\da-zA-Z-]+)\s+(?P<R>\d+)\s+(?P<G>\d+)\s+(?P<B>\d+)\s+(?P<A>\d+)"
+    # )
+    # with open(filename, "r") as f:
+    #     records = [lut_record(m) for m in map(lut_regex.match, f) if m is not None]
+    # return pd.DataFrame.from_records(records)
     with open(filename, "r") as f:
-        records = [lut_record(m) for m in map(lut_regex.match, f) if m is not None]
-    return pd.DataFrame.from_records(records)
+        return pd.read_csv(
+            f,
+            sep=r"\s+",
+            header=None,
+            names=["ID", "Label", "R", "G", "B", "A"],
+            index_col="ID",
+            comment="#",  # Skips comment lines if any exist
+        )
 
 
 def write_lut(filename: Path | str, table: pd.DataFrame):
@@ -37,15 +45,3 @@ def write_lut(filename: Path | str, table: pd.DataFrame):
         newtable[col] = (newtable[col] * 255).astype(int)
     newtable["A"] = 255 - (newtable["A"] * 255).astype(int)
     newtable.to_csv(filename, sep="\t", index=False, header=False)
-
-
-def lut_record(match: re.Match) -> dict[str, str | float]:
-    groupdict = match.groupdict()
-    return {
-        "label": int(groupdict["label"]),
-        "description": groupdict["description"],
-        "R": float(groupdict["R"]) / 255,
-        "G": float(groupdict["G"]) / 255,
-        "B": float(groupdict["B"]) / 255,
-        "A": 1.0 - float(groupdict["A"]) / 255,
-    }
diff --git a/src/mritk/segmentation/segmentation.py b/src/mritk/segmentation/segmentation.py
new file mode 100644
index 0000000..74a6272
--- /dev/null
+++ b/src/mritk/segmentation/segmentation.py
@@ -0,0 +1,43 @@
+from ..data.base import MRIData
+import numpy as np
+import pandas as pd
+from .lookup_table import read_lut
+from pathlib import Path
+from ..data.io import load_mri_data
+
+
+class Segmentation(MRIData):
+    def __init__(self, data: np.ndarray, affine: np.ndarray):
+        super().__init__(data, affine)
+        self.data = self.data.astype(int)
+        self.rois = np.unique(self.data[self.data > 0])
+
+    @classmethod
+    def from_file(cls, filepath: Path):
+        data, affine = load_mri_data(filepath, dtype=int)
+        return cls(data=data, affine=affine)
+
+    @property
+    def num_rois(self) -> int:
+        return len(self.rois)
+
+    @property
+    def roi_labels(self) -> np.ndarray:
+        return self.rois
+
+
+class FreeSurferSegmentation(Segmentation):
+    def __init__(self, data: np.ndarray, affine: np.ndarray):
+        super().__init__(data, affine)
+
+        # TODO: verify that all labels in the segmentation are valid FS labels
+        # Retrieve the FreeSurfer LUT, which contains the labels associated with each ROI
+        self.freesurfer_lut = read_lut(None)  # Finds freesurfer LUT
+
+    def get_roi_labels(self, rois: list[int] | None = None) -> pd.DataFrame:
+        if not np.isin(rois, self.rois).all():
+            raise ValueError("Some of the provided ROIs are not present in the segmentation.")
+
+        if rois is None:
+            rois = self.rois
+        return self.freesurfer_lut.loc[self.freesurfer_lut.index.isin(rois), ["Label"]].rename_axis("ROI").reset_index()
diff --git a/src/mritk/statistics/__init__.py b/src/mritk/statistics/__init__.py
index 2c107e5..0e0f164 100644
--- a/src/mritk/statistics/__init__.py
+++ b/src/mritk/statistics/__init__.py
@@ -2,6 +2,6 @@
 # Copyright (C) 2026   Cécile Daversin-Catty (cecile@simula.no)
 # Copyright (C) 2026   Simula Research Laboratory
 
-from . import utils, compute_stats, cli
+from . import utils, compute_stats, cli, stat_functions
 
-__all__ = ["utils", "compute_stats", "cli"]
+__all__ = ["utils", "compute_stats", "cli", "stat_functions"]
diff --git a/src/mritk/statistics/cli.py b/src/mritk/statistics/cli.py
index 491f439..85371a6 100644
--- a/src/mritk/statistics/cli.py
+++ b/src/mritk/statistics/cli.py
@@ -3,20 +3,19 @@
 from pathlib import Path
 import pandas as pd
 
-from ..segmentation.groups import default_segmentation_groups
 from .compute_stats import generate_stats_dataframe
+from ..data.base import MRIData
+from ..segmentation.segmentation import Segmentation
+from .metadata import extract_metadata_from_bids
 
 
 def compute_mri_stats(
     segmentation: Path,
     mri: list[Path],
     output: Path,
-    timetable: Path | None = None,
-    timelabel: str | None = None,
-    seg_regex: str | None = None,
-    mri_regex: str | None = None,
     lut: Path | None = None,
     info: str | None = None,
+    use_bids_metadata: bool = False,
     **kwargs,
 ):
     import sys
@@ -40,6 +39,7 @@ def compute_mri_stats(
         console.print(f"[bold red]Error:[/bold red] Missing segmentation file: {segmentation}")
         sys.exit(1)
 
+    seg = Segmentation.from_file(segmentation)
     # Validate all MRI paths before starting
     for path in mri:
         if not path.exists():
@@ -51,19 +51,23 @@ def compute_mri_stats(
     # Loop through MRI paths
     console.print("[bold green]Processing MRIs...[/bold green]")
     for i, path in enumerate(mri):
-        # console.print(f"[blue]Processing MRI {i + 1}/{len(mri)}:[/blue] {path.name}")
+        if use_bids_metadata:
+            try:
+                bids_metadata = extract_metadata_from_bids(segmentation, path)
+            except Exception as e:
+                console.print(f"[bold red]Error extracting BIDS metadata:[/bold red] {e}")
+                sys.exit(1)
 
+            info_dict = (info_dict if info_dict else {}) | bids_metadata
+
+        mri_object = MRIData.from_file(path)  # Load MRI data
         try:
             # Call the logic function
+            # TODO: Add option to specify statistics to compute
             df = generate_stats_dataframe(
-                seg_path=segmentation,
-                mri_path=path,
-                timestamp_path=timetable,
-                timestamp_sequence=timelabel,
-                seg_pattern=seg_regex,
-                mri_data_pattern=mri_regex,
-                lut_path=lut,
-                info_dict=info_dict,
+                seg=seg,
+                mri=mri_object,
+                metadata=info_dict,
             )
             dataframes.append(df)
         except Exception as e:
@@ -84,7 +88,7 @@ def compute_mri_stats(
         console.print("[yellow]No dataframes generated.[/yellow]")
 
 
-def get_stats_value(stats_file: Path, region: str, info: str, **kwargs):
+def get_stats_value(stats_file: Path, ROI: int, statistic: str, **kwargs):
     """
     Replaces the @click.command('get') decorated function.
     """
@@ -94,51 +98,40 @@ def get_stats_value(stats_file: Path, region: str, info: str, **kwargs):
     # Setup Rich
     console = Console()
 
-    # Validate inputs
-    valid_regions = default_segmentation_groups().keys()
-    if region not in valid_regions:
-        console.print(f"[bold red]Error:[/bold red] Region '{region}' not found in default segmentation groups.")
-        sys.exit(1)
-
-    valid_infos = [
-        "sum",
-        "mean",
-        "median",
-        "std",
-        "min",
-        "max",
-        "PC1",
-        "PC5",
-        "PC25",
-        "PC75",
-        "PC90",
-        "PC95",
-        "PC99",
-    ]
-    if info not in valid_infos:
-        console.print(f"[bold red]Error:[/bold red] Info '{info}' is invalid. Choose from: {', '.join(valid_infos)}")
-        sys.exit(1)
-
+    # Verify that csv exists
     if not stats_file.exists():
         console.print(f"[bold red]Error:[/bold red] Stats file not found: {stats_file}")
         sys.exit(1)
 
     # Process
     try:
+        # Read csv
         df = pd.read_csv(stats_file, sep=";")
-        region_row = df.loc[df["description"] == region]
 
-        if region_row.empty:
-            console.print(f"[red]Region '{region}' not found in the stats file.[/red]")
+        # Verify that the requested statistic exists in the dataframe
+        valid_statistics = set(df["statistic"])
+        if statistic not in valid_statistics:
+            console.print(
+                f"[bold red]Error:[/bold red] Statistic '{statistic}' is invalid. Choose from: {', '.join(valid_statistics)}"
+            )
             sys.exit(1)
 
-        info_value = region_row[info].values[0]
+        # Verify that the requested ROI exists in the dataframe
+        valid_rois = set(df["ROI"])
+        if ROI not in valid_rois:
+            console.print(
+                f"[bold red]Error:[/bold red] ROI '{ROI}' not found in stats file. Valid ROIs: {', '.join(map(str, valid_rois))}"
+            )
+            sys.exit(1)
+
+        statistic_value = df.loc[(df["ROI"] == ROI) & (df["statistic"] == statistic), "value"]
 
         # Output
         console.print(
-            f"[bold cyan]{info}[/bold cyan] for [bold green]{region}[/bold green] = [bold white]{info_value}[/bold white]"
+            f"[bold cyan]{statistic}[/bold cyan] for ROI \
+                [bold green]{ROI}[/bold green] = [bold white]{statistic_value.item()}[/bold white]"
         )
-        return info_value
+        return statistic_value.item()  # Return as scalar
 
     except Exception as e:
         console.print(f"[bold red]Error reading stats file:[/bold red] {e}")
@@ -150,28 +143,50 @@ def add_arguments(parser: argparse.ArgumentParser):
 
     # --- Compute Command ---
     parser_compute = subparsers.add_parser("compute", help="Compute MRI statistics", formatter_class=parser.formatter_class)
-    parser_compute.add_argument("--segmentation", "-s", type=Path, required=True, help="Path to segmentation file")
-    parser_compute.add_argument("--mri", "-m", type=Path, nargs="+", required=True, help="Path to MRI data file(s)")
+    parser_compute.add_argument(
+        "--segmentation",
+        "-s",
+        type=Path,
+        required=True,
+        help="Path to segmentation file",
+    )
+    parser_compute.add_argument(
+        "--mri",
+        "-m",
+        type=Path,
+        nargs="+",
+        required=True,
+        help="Path to MRI data file(s)",
+    )
     parser_compute.add_argument("--output", "-o", type=Path, required=True, help="Output CSV file path")
-    parser_compute.add_argument("--timetable", "-t", type=Path, help="Path to timetable file")
-    parser_compute.add_argument("--timelabel", "-l", dest="timelabel", type=str, help="Time label sequence")
+    parser_compute.add_argument("--lut", "-lt", dest="lut", type=Path, help="Path to Lookup Table")
     parser_compute.add_argument(
-        "--seg_regex",
-        "-sr",
-        dest="seg_regex",
+        "--info",
+        "-i",
         type=str,
-        help="Regex pattern for segmentation filename",
+        help="Info dictionary as JSON string. \
+            If using --use_bids_metadata, overlapping fields will be overwritten by BIDS metadata extraction.",
+    )
+    parser_compute.add_argument(
+        "--use_bids_metadata",
+        "-b",
+        action="store_true",
+        help="Assumes file naming follows BIDS convention and extracts metadata accordingly.\
+            Checks that subject IDs match between segmentation and MRI data.",
     )
-    parser_compute.add_argument("--mri_regex", "-mr", dest="mri_regex", type=str, help="Regex pattern for MRI filename")
-    parser_compute.add_argument("--lut", "-lt", dest="lut", type=Path, help="Path to Lookup Table")
-    parser_compute.add_argument("--info", "-i", type=str, help="Info dictionary as JSON string")
     parser_compute.set_defaults(func=compute_mri_stats)
 
     # --- Get Command ---
     parser_get = subparsers.add_parser("get", help="Get specific stats value", formatter_class=parser.formatter_class)
     parser_get.add_argument("--stats_file", "-f", type=Path, required=True, help="Path to stats CSV file")
-    parser_get.add_argument("--region", "-r", type=str, required=True, help="Region description")
-    parser_get.add_argument("--info", "-i", type=str, required=True, help="Statistic to retrieve (mean, std, etc.)")
+    parser_get.add_argument("--ROI", "-r", type=int, required=True, help="Region of interest to extract")
+    parser_get.add_argument(
+        "--statistic",
+        "-s",
+        type=str,
+        required=True,
+        help="Statistic to retrieve (mean, std, etc.)",
+    )
     parser_get.set_defaults(func=get_stats_value)
 
 
diff --git a/src/mritk/statistics/compute_stats.py b/src/mritk/statistics/compute_stats.py
index d54f981..0f5009d 100644
--- a/src/mritk/statistics/compute_stats.py
+++ b/src/mritk/statistics/compute_stats.py
@@ -4,129 +4,60 @@
 # Copyright (C) 2026   Cécile Daversin-Catty (cecile@simula.no)
 # Copyright (C) 2026   Simula Research Laboratory
 
-from pathlib import Path
 from typing import Optional
-import re
 import numpy as np
 import pandas as pd
 import tqdm.rich
 
-from ..data.io import load_mri_data
 from ..data.orientation import assert_same_space
-from ..segmentation.groups import default_segmentation_groups
-from ..segmentation.lookup_table import read_lut
-from .utils import voxel_count_to_ml_scale, find_timestamp, prepend_info
+from .utils import prepend_info
+from ..segmentation.segmentation import Segmentation
+from ..data.base import MRIData
+from .stat_functions import Statistic, Mean, Std, Median
 
 
 def generate_stats_dataframe(
-    seg_path: Path,
-    mri_path: Path,
-    timestamp_path: Optional[str | Path] = None,
-    timestamp_sequence: Optional[str | Path] = None,
-    seg_pattern: Optional[str | Path] = None,
-    mri_data_pattern: Optional[str | Path] = None,
-    lut_path: Optional[Path] = None,
-    info_dict: Optional[dict] = None,
+    seg: Segmentation,
+    mri: MRIData,
+    qois: list[Statistic] = [Mean, Std, Median],
+    metadata: Optional[dict] = None,
 ) -> pd.DataFrame:
-    # Load the data (mri and seg)
-    mri = load_mri_data(mri_path, dtype=np.single)
-    seg = load_mri_data(seg_path, dtype=np.int16)
+    # Verify that segmentation and MRI are in the same space
     assert_same_space(seg, mri)
-    # Load LUT
-    lut = read_lut(lut_path)
-    # Get LUT info
-    seg_labels = np.unique(seg.data[seg.data != 0])
-    lut_regions = lut.loc[lut.label.isin(seg_labels), ["label", "description"]].to_dict("records")
-    regions = {
-        **{d["description"]: sorted([d["label"]]) for d in lut_regions},
-        **default_segmentation_groups(),
-    }
-    # Get SEG info
-    seg_info = {}
-    if seg_pattern is not None:
-        seg_pattern = rf"{seg_pattern}"
-        if (m := re.match(seg_pattern, Path(seg_path).name)) is not None:
-            seg_info = m.groupdict()
-        else:
-            raise RuntimeError(f"Segmentation filename {seg_path.name} does not match the provided pattern.")
-    elif info_dict is not None:
-        seg_info["segmentation"] = info_dict["segmentation"] if "segmentation" in info_dict else None
-        seg_info["subject"] = info_dict["subject"] if "subject" in info_dict else None
-    else:
-        seg_info = {"segmentation": None, "subject": None}
-    # Get MRI info
-    mri_info = {}
-    if mri_data_pattern is not None:
-        mri_data_pattern = rf"{mri_data_pattern}"
-        if (m := re.match(mri_data_pattern, Path(mri_path).name)) is not None:
-            mri_info = m.groupdict()
-        else:
-            raise RuntimeError(f"MRI data filename {mri_path.name} does not match the provided pattern.")
-    elif info_dict is not None:
-        mri_info["mri_data"] = info_dict["mri_data"] if "mri_data" in info_dict else None
-        mri_info["subject"] = info_dict["subject"] if "subject" in info_dict else None
-        mri_info["session"] = info_dict["session"] if "session" in info_dict else None
-    else:
-        mri_info = {"mri_data": None, "subject": None, "session": None}
-    # Get timestamp
-    if timestamp_path is not None:
-        try:
-            mri_info["timestamp"] = find_timestamp(
-                Path(str(timestamp_path)),
-                str(timestamp_sequence),
-                str(mri_info["subject"]),
-                str(mri_info["session"]),
-            )
-        except (ValueError, RuntimeError, KeyError):
-            mri_info["timestamp"] = None
-    else:
-        mri_info["timestamp"] = None
 
-    info = seg_info | mri_info
+    qoi_records = []  # Collects records related to qois
+    roi_records = []  # Collects records related to ROIs,
 
-    records = []
+    # Mask infinite values
     finite_mask = np.isfinite(mri.data)
-    volscale = voxel_count_to_ml_scale(seg.affine)
-
-    for description, labels in tqdm.rich.tqdm(regions.items(), total=len(regions)):
-        region_mask = np.isin(seg.data, labels)
-        voxelcount = region_mask.sum()
-        record = {
-            "label": ",".join([str(x) for x in labels]),
-            "description": description,
-            "voxelcount": voxelcount,
-            "volume_ml": volscale * voxelcount,
-        }
-        if voxelcount == 0:
-            records.append(record)
-            continue
-
-        data_mask = region_mask * finite_mask
-        region_data = mri.data[data_mask]
-        num_nan = (~np.isfinite(region_data)).sum()
-        record["num_nan_values"] = num_nan
-        if num_nan == voxelcount:
-            records.append(record)
-            continue
-
-        stats = {
-            "sum": np.sum(region_data),
-            "mean": np.mean(region_data),
-            "median": np.median(region_data),
-            "std": np.std(region_data),
-            "min": np.min(region_data),
-            **{f"PC{pc}": np.quantile(region_data, pc / 100) for pc in [1, 5, 25, 75, 90, 95, 99]},
-            "max": np.max(region_data),
-        }
-        records.append({**record, **stats})
-
-    dframe = pd.DataFrame.from_records(records)
-    dframe = prepend_info(
-        dframe,
-        segmentation=info["segmentation"],
-        mri_data=info["mri_data"],
-        subject=info["subject"],
-        session=info["session"],
-        timestamp=info["timestamp"],
-    )
-    return dframe
+    for roi in tqdm.rich.tqdm(seg.roi_labels, total=len(seg.roi_labels)):
+        # Identify rois in segmentation
+        region_mask = (seg.data == roi) * finite_mask
+        # print(region_mask.shape)
+        region_data = mri.data[region_mask]
+        nb_nans = np.isnan(region_data).sum()
+
+        voxelcount = len(region_data)
+
+        roi_records.append(
+            {
+                "ROI": roi,
+                "voxel_count": voxelcount,
+                "volume_ml": seg.voxel_ml_volume * voxelcount,
+                "num_nan_values": nb_nans,
+            }
+        )
+        # Iterate qoi functions
+        for qoi in qois:
+            qoi_value = qoi(region_data)
+            # Store the qoi value in a dataframe, along with the roi label and description
+            qoi_records.append({"ROI": roi, "statistic": qoi.name, "value": qoi_value})
+
+    df = pd.DataFrame.from_records(qoi_records)
+    df_roi = pd.DataFrame.from_records(roi_records)
+    df = df.merge(df_roi, on="ROI", how="left")
+
+    # Add some metadata to each row
+    if metadata is not None:
+        df = prepend_info(df, **(metadata))
+    return df
diff --git a/src/mritk/statistics/metadata.py b/src/mritk/statistics/metadata.py
new file mode 100644
index 0000000..63d3708
--- /dev/null
+++ b/src/mritk/statistics/metadata.py
@@ -0,0 +1,45 @@
+import re
+from pathlib import Path
+
+
+def extract_pattern_from_path(pattern, path: Path):
+    if (m := re.match(pattern, Path(path).name)) is not None:
+        info = m.groupdict()
+    else:
+        raise RuntimeError(f"Filename {path.name} does not match the provided pattern.")
+
+    return info
+
+
+def extract_metadata_from_bids(
+    segmentation_path: Path,
+    mri_data_path: Path,
+) -> dict:
+    """Extract subject, session, mri data type and segmentation name from filepath. \
+        Assumes that naming follows the BIDS convention
+
+    Args:
+        segmentation_path (Path): Path so segmentation file
+        mri_data_path (Path): Path to mri data file
+
+    Raises:
+        RuntimeError: If subject ID in the segmentation filename does not match the subject ID in the mri data filename
+
+    Returns:
+        dict: Combined subject, session, mri data type and segmentation name
+    """
+
+    seg_pattern = r"sub-(?P<subject>[^\.]+)_seg-(?P<segmentation>[^\.]+)"
+    # Identify subject and segmentation from segmentation filename
+    seg_info = extract_pattern_from_path(pattern=seg_pattern, path=segmentation_path)
+
+    mri_data_pattern = r"sub-(?P<subject>[^\.]+)_(?P<session>ses-\d{2})_(?P<mri_data>[^\.]+)"
+    # Identify subject, session and mri data type from mri data filename
+    mri_info = extract_pattern_from_path(pattern=mri_data_pattern, path=mri_data_path)
+
+    if mri_info["subject"] != seg_info["subject"]:
+        raise RuntimeError(
+            f"Subject ID mismatch between segmentation and MRI data: {seg_info['subject']} vs {mri_info['subject']}"
+        )
+
+    return seg_info | mri_info
diff --git a/src/mritk/statistics/stat_functions.py b/src/mritk/statistics/stat_functions.py
new file mode 100644
index 0000000..b0648f3
--- /dev/null
+++ b/src/mritk/statistics/stat_functions.py
@@ -0,0 +1,53 @@
+from dataclasses import dataclass
+import numpy as np
+from typing import Callable, Any
+
+
+@dataclass
+class Statistic:
+    name: str
+    func: Callable
+
+    def __call__(self, data) -> Any:
+        return self.func(data)
+
+
+Median = Statistic("median", lambda x: np.median(x))
+Mean = Statistic("mean", lambda x: np.mean(x))
+Std = Statistic("std", lambda x: np.std(x))
+Sum = Statistic("sum", lambda x: np.sum(x))
+Min = Statistic("min", lambda x: np.min(x))
+Max = Statistic("max", lambda x: np.max(x))
+
+
+@dataclass
+class PCx(Statistic):
+    percentile: int
+
+    def __init__(self, percentile) -> None:
+        super().__init__(f"PC{percentile}", lambda x: np.percentile(x, percentile))
+        self.percentile = percentile
+
+
+# Etc
+PC1 = PCx(1)
+PC5 = PCx(5)
+PC25 = PCx(25)
+PC75 = PCx(75)
+PC95 = PCx(95)
+PC99 = PCx(99)
+
+
+@dataclass
+class StableStatistic(Statistic):
+    low: int
+    high: int
+
+    def __call__(self, data) -> Any:
+        low_value = np.percentile(data, self.low)
+        high_value = np.percentile(data, self.high)
+        return super().__call__(data[(data > low_value) & (data < high_value)])
+
+
+StableMean = StableStatistic("stable_mean", lambda x: np.mean(x), 5, 95)
+StableStd = StableStatistic("stable_std", lambda x: np.std(x), 5, 95)
diff --git a/src/mritk/statistics/utils.py b/src/mritk/statistics/utils.py
index d598373..e6fc31b 100644
--- a/src/mritk/statistics/utils.py
+++ b/src/mritk/statistics/utils.py
@@ -5,13 +5,12 @@
 # Copyright (C) 2026   Simula Research Laboratory
 
 
-import numpy as np
 from pathlib import Path
 import pandas as pd
 
 
-def voxel_count_to_ml_scale(affine: np.ndarray):
-    return 1e-3 * np.linalg.det(affine[:3, :3])
+# def voxel_count_to_ml_scale(affine: np.ndarray):
+#     return 1e-3 * np.linalg.det(affine[:3, :3])
 
 
 def find_timestamp(
diff --git a/test/conftest.py b/test/conftest.py
index dbfe221..17c3db8 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -1,8 +1,36 @@
 from pathlib import Path
 import os
 import pytest
+import numpy as np
+from mritk.data.base import MRIData
+from mritk.segmentation.segmentation import Segmentation
 
 
 @pytest.fixture(scope="session")
 def mri_data_dir() -> Path:
     return Path(os.getenv("MRITK_TEST_DATA_FOLDER", "test_data"))
+
+
+@pytest.fixture
+def example_segmentation() -> Segmentation:
+    """Example segmentation"""
+    base = np.array([0, 1, 2, 3], dtype=float)
+    seg = np.tile(base, (100, 1))
+
+    return Segmentation(seg, affine=np.eye(4))
+
+
+@pytest.fixture
+def example_values() -> MRIData:
+    """Example values for testing qoi computations"""
+    np.random.seed(0)
+
+    data = np.array(
+        [
+            np.random.normal(0.0, size=100),
+            np.random.normal(1.0, size=100),
+            np.random.normal(2.0, size=100),
+            np.random.normal(3.0, size=100),
+        ]
+    ).T
+    return MRIData(data, affine=np.eye(4))
diff --git a/test/test_metadata.py b/test/test_metadata.py
new file mode 100644
index 0000000..466bca1
--- /dev/null
+++ b/test/test_metadata.py
@@ -0,0 +1,27 @@
+from mritk.statistics.metadata import (
+    extract_metadata_from_bids,
+    extract_pattern_from_path,
+)
+from pathlib import Path
+
+
+def test_path_extraction():
+    pattern = r"sub-(?P<subject>[^\.]+)_(?P<session>ses-\d{2})_(?P<mri_data>[^\.]+)"
+
+    path = Path("sub-01_ses-01_concentration.nii.gz")
+    info = extract_pattern_from_path(pattern, path)
+    assert info["subject"] == "01"
+    assert info["session"] == "ses-01"
+    assert info["mri_data"] == "concentration"
+
+
+def test_bids_extraction():
+    seg_path = Path("sub-01_seg-aparc+aseg_refined.nii.gz")
+    mri_path = Path("sub-01_ses-01_concentration.nii.gz")
+
+    metadata = extract_metadata_from_bids(segmentation_path=seg_path, mri_data_path=mri_path)
+
+    assert metadata["subject"] == "01"
+    assert metadata["session"] == "ses-01"
+    assert metadata["mri_data"] == "concentration"
+    assert metadata["segmentation"] == "aparc+aseg_refined"
diff --git a/test/test_mri_io.py b/test/test_mri_io.py
index 34ffeeb..74743fe 100644
--- a/test/test_mri_io.py
+++ b/test/test_mri_io.py
@@ -7,12 +7,51 @@
 
 import numpy as np
 from mritk.data.io import load_mri_data, save_mri_data
+from mritk.data.base import MRIData
+from mritk.segmentation.segmentation import Segmentation
 
 
 def test_mri_io_nifti(tmp_path, mri_data_dir):
-    input_file = mri_data_dir / "mri-processed/mri_dataset/derivatives/sub-01/ses-01/sub-01_ses-01_acq-mixed_T1map.nii.gz"
+    input_file = mri_data_dir / "mri-dataset/mri_dataset/derivatives/sub-01/ses-01/sub-01_ses-01_acq-mixed_T1map.nii.gz"
 
     output_file = tmp_path / "output_nifti.nii.gz"
 
-    mri = load_mri_data(input_file, dtype=np.single, orient=False)  ## TODO : Test orient=True case
-    save_mri_data(mri, output_file, dtype=np.single)
+    data, affine = load_mri_data(input_file, dtype=np.single)  ## TODO : Test orient=True case
+    save_mri_data(data, affine, output_file, dtype=np.single)
+
+
+def test_MRIData_io(tmp_path, mri_data_dir):
+    input_file = mri_data_dir / "mri-dataset/mri_dataset/derivatives/sub-01/ses-01/sub-01_ses-01_acq-mixed_T1map.nii.gz"
+
+    output_file = tmp_path / "output_mridata.nii.gz"
+
+    mri_data = MRIData.from_file(input_file, dtype=np.single)
+    mri_data.save_mri_data(output_file, intent_code=1006)
+
+
+def test_MRIData_io_invalid_suffix(tmp_path, mri_data_dir):
+    input_file = mri_data_dir / "mri-dataset/mri_dataset/derivatives/sub-01/ses-01/sub-01_ses-01_acq-mixed_T1map.nii.gz"
+
+    output_file = tmp_path / "output_mridata.invalid"
+
+    mri_data = MRIData.from_file(input_file, dtype=np.single)
+    try:
+        mri_data.save_mri_data(output_file, intent_code=1006)
+        assert False, "Expected ValueError for invalid suffix"
+    except ValueError as e:
+        assert str(e) == f"Invalid suffix {output_file}, should be either '.nii', or '.mgz'"
+
+
+def test_load_mri_data_invalid_suffix(mri_data_dir):
+    input_file = mri_data_dir / "mri-dataset/mri_dataset/derivatives/sub-01/ses-01/sub-01_ses-01_acq-mixed_T1map.invalid"
+    try:
+        load_mri_data(input_file, dtype=np.single)
+        assert False, "Expected ValueError for invalid suffix"
+    except ValueError as e:
+        assert str(e) == f"Invalid suffix {input_file}, should be either '.nii', or '.mgz'"
+
+
+def test_load_Segmentation(mri_data_dir):
+    input_file = mri_data_dir / "mri-processed/mri_processed_data/sub-01/segmentations/sub-01_seg-aparc+aseg_refined.nii.gz"
+    seg = Segmentation.from_file(input_file)
+    assert seg.data.dtype == int
diff --git a/test/test_mri_stats.py b/test/test_mri_stats.py
index a259e86..872290b 100644
--- a/test/test_mri_stats.py
+++ b/test/test_mri_stats.py
@@ -5,19 +5,28 @@
 Copyright (C) 2026   Simula Research Laboratory
 """
 
-from pathlib import Path
+from mritk.statistics.compute_stats import (
+    generate_stats_dataframe,
+)
 
+from mritk.data.base import MRIData
+from mritk.segmentation.segmentation import Segmentation
 
-from mritk.statistics.compute_stats import generate_stats_dataframe  # , compute_mri_stats
-import mritk.cli as cli
-
-
-def test_compute_stats_default(mri_data_dir: Path):
-    seg_path = mri_data_dir / "mri-processed/mri_processed_data/sub-01" / "segmentations/sub-01_seg-aparc+aseg_refined.nii.gz"
-    mri_path = mri_data_dir / "mri-processed/mri_processed_data/sub-01" / "concentrations/sub-01_ses-01_concentration.nii.gz"
-
-    dataframe = generate_stats_dataframe(seg_path, mri_path)
 
+def test_compute_stats_default(example_segmentation: Segmentation, example_values: MRIData):
+    print(example_values.data.shape)
+    dataframe = generate_stats_dataframe(
+        example_segmentation,
+        example_values,
+        metadata={
+            "segmentation": "segmentation",
+            "mri_data": "mri_data",
+            "subject": "subject",
+            "session": "session",
+            "timestamp": "timestamp",
+        },
+    )
+    print(dataframe.columns)
     assert not dataframe.empty
     assert set(dataframe.columns) == {
         "segmentation",
@@ -25,114 +34,142 @@ def test_compute_stats_default(mri_data_dir: Path):
         "subject",
         "session",
         "timestamp",
-        "label",
-        "description",
-        "voxelcount",
+        "ROI",
+        "voxel_count",
         "volume_ml",
         "num_nan_values",
-        "sum",
-        "mean",
-        "median",
-        "std",
-        "min",
-        "PC1",
-        "PC5",
-        "PC25",
-        "PC75",
-        "PC90",
-        "PC95",
-        "PC99",
-        "max",
+        "statistic",
+        "value",
     }
 
-
-def test_compute_stats_patterns(mri_data_dir: Path):
-    seg_path = mri_data_dir / "mri-processed/mri_processed_data/sub-01" / "segmentations/sub-01_seg-aparc+aseg_refined.nii.gz"
-    mri_path = mri_data_dir / "mri-processed/mri_processed_data/sub-01" / "concentrations/sub-01_ses-01_concentration.nii.gz"
-    seg_pattern = "(?P<subject>sub-(control|patient)*\\d{2})_seg-(?P<segmentation>[^\\.]+)"
-    mri_data_pattern = "(?P<subject>sub-(control|patient)*\\d{2})_(?P<session>ses-\\d{2})_(?P<mri_data>[^\\.]+)"
-
-    dataframe = generate_stats_dataframe(
-        seg_path,
-        mri_path,
-        seg_pattern=seg_pattern,
-        mri_data_pattern=mri_data_pattern,
-    )
-
-    assert not dataframe.empty
-    assert dataframe["subject"].iloc[0] == "sub-01"
-    assert dataframe["segmentation"].iloc[0] == "aparc+aseg_refined"
-    assert dataframe["mri_data"].iloc[0] == "concentration"
-    assert dataframe["session"].iloc[0] == "ses-01"
-
-
-def test_compute_stats_timestamp(mri_data_dir: Path):
-    seg_path = mri_data_dir / "mri-processed/mri_processed_data/sub-01" / "segmentations/sub-01_seg-aparc+aseg_refined.nii.gz"
-    mri_path = mri_data_dir / "mri-processed/mri_processed_data/sub-01" / "concentrations/sub-01_ses-01_concentration.nii.gz"
-    seg_pattern = "(?P<subject>sub-(control|patient)*\\d{2})_seg-(?P<segmentation>[^\\.]+)"
-    mri_data_pattern = "(?P<subject>sub-(control|patient)*\\d{2})_(?P<session>ses-\\d{2})_(?P<mri_data>[^\\.]+)"
-    timetable = mri_data_dir / "timetable/timetable.tsv"
-    timetable_sequence = "mixed"
-
-    dataframe = generate_stats_dataframe(
-        seg_path,
-        mri_path,
-        seg_pattern=seg_pattern,
-        mri_data_pattern=mri_data_pattern,
-        timestamp_path=timetable,
-        timestamp_sequence=timetable_sequence,
-    )
-
-    assert dataframe["timestamp"].iloc[0] == -6414.9
-
-
-def test_compute_stats_info(mri_data_dir: Path):
-    seg_path = mri_data_dir / "mri-processed/mri_processed_data/sub-01" / "segmentations/sub-01_seg-aparc+aseg_refined.nii.gz"
-    mri_path = mri_data_dir / "mri-processed/mri_processed_data/sub-01" / "concentrations/sub-01_ses-01_concentration.nii.gz"
-    info = {
-        "mri_data": "concentration",
-        "subject": "sub-01",
-        "session": "ses-01",
-        "segmentation": "aparc+aseg_refined",
-    }
-
-    dataframe = generate_stats_dataframe(seg_path, mri_path, info_dict=info)
-
-    assert not dataframe.empty
-    assert dataframe["subject"].iloc[0] == "sub-01"
-    assert dataframe["segmentation"].iloc[0] == "aparc+aseg_refined"
-    assert dataframe["mri_data"].iloc[0] == "concentration"
-    assert dataframe["session"].iloc[0] == "ses-01"
-
-
-def test_compute_mri_stats_cli(capsys, tmp_path: Path, mri_data_dir: Path):
-    seg_path = mri_data_dir / "mri-processed/mri_processed_data/sub-01" / "segmentations/sub-01_seg-aparc+aseg_refined.nii.gz"
-    mri_path = mri_data_dir / "mri-processed/mri_processed_data/sub-01" / "concentrations/sub-01_ses-01_concentration.nii.gz"
-    seg_pattern = "(?P<subject>sub-(control|patient)*\\d{2})_seg-(?P<segmentation>[^\\.]+)"
-    mri_data_pattern = "(?P<subject>sub-(control|patient)*\\d{2})_(?P<session>ses-\\d{2})_(?P<mri_data>[^\\.]+)"
-    timetable = mri_data_dir / "timetable/timetable.tsv"
-    timetable_sequence = "mixed"
-
-    args = [
-        "--segmentation",
-        str(seg_path),
-        "--mri",
-        str(mri_path),
-        "--output",
-        str(tmp_path / "mri_stats_output.csv"),
-        "--timetable",
-        str(timetable),
-        "--timelabel",
-        timetable_sequence,
-        "--seg_regex",
-        seg_pattern,
-        "--mri_regex",
-        mri_data_pattern,
-    ]
-
-    ret = cli.main(["stats", "compute"] + args)
-    assert ret == 0
-    captured = capsys.readouterr()
-    assert "Processing MRIs..." in captured.out
-    assert "Stats successfully saved to" in captured.out
-    assert (tmp_path / "mri_stats_output.csv").exists()
+    assert set(dataframe["statistic"].unique()) == {"mean", "median", "std"}
+
+
+# def test_compute_stats_patterns(mri_data_dir: Path):
+#     seg_path = (
+#         mri_data_dir
+#         / "mri-processed/mri_processed_data/sub-01"
+#         / "segmentations/sub-01_seg-aparc+aseg_refined.nii.gz"
+#     )
+#     mri_path = (
+#         mri_data_dir
+#         / "mri-processed/mri_processed_data/sub-01"
+#         / "concentrations/sub-01_ses-01_concentration.nii.gz"
+#     )
+#     seg_pattern = (
+#         "(?P<subject>sub-(control|patient)*\\d{2})_seg-(?P<segmentation>[^\\.]+)"
+#     )
+#     mri_data_pattern = "(?P<subject>sub-(control|patient)*\\d{2})_(?P<session>ses-\\d{2})_(?P<mri_data>[^\\.]+)"
+
+#     dataframe = generate_stats_dataframe(
+#         seg_path,
+#         mri_path,
+#         seg_pattern=seg_pattern,
+#         mri_data_pattern=mri_data_pattern,
+#     )
+
+#     assert not dataframe.empty
+#     assert dataframe["subject"].iloc[0] == "sub-01"
+#     assert dataframe["segmentation"].iloc[0] == "aparc+aseg_refined"
+#     assert dataframe["mri_data"].iloc[0] == "concentration"
+#     assert dataframe["session"].iloc[0] == "ses-01"
+
+
+# def test_compute_stats_timestamp(mri_data_dir: Path):
+#     seg_path = (
+#         mri_data_dir
+#         / "mri-processed/mri_processed_data/sub-01"
+#         / "segmentations/sub-01_seg-aparc+aseg_refined.nii.gz"
+#     )
+#     mri_path = (
+#         mri_data_dir
+#         / "mri-processed/mri_processed_data/sub-01"
+#         / "concentrations/sub-01_ses-01_concentration.nii.gz"
+#     )
+#     seg_pattern = (
+#         "(?P<subject>sub-(control|patient)*\\d{2})_seg-(?P<segmentation>[^\\.]+)"
+#     )
+#     mri_data_pattern = "(?P<subject>sub-(control|patient)*\\d{2})_(?P<session>ses-\\d{2})_(?P<mri_data>[^\\.]+)"
+#     timetable = mri_data_dir / "timetable/timetable.tsv"
+#     timetable_sequence = "mixed"
+
+#     dataframe = generate_stats_dataframe(
+#         seg_path,
+#         mri_path,
+#         seg_pattern=seg_pattern,
+#         mri_data_pattern=mri_data_pattern,
+#         timestamp_path=timetable,
+#         timestamp_sequence=timetable_sequence,
+#     )
+
+#     assert dataframe["timestamp"].iloc[0] == -6414.9
+
+
+# def test_compute_stats_info(mri_data_dir: Path):
+#     seg_path = (
+#         mri_data_dir
+#         / "mri-processed/mri_processed_data/sub-01"
+#         / "segmentations/sub-01_seg-aparc+aseg_refined.nii.gz"
+#     )
+#     mri_path = (
+#         mri_data_dir
+#         / "mri-processed/mri_processed_data/sub-01"
+#         / "concentrations/sub-01_ses-01_concentration.nii.gz"
+#     )
+#     info = {
+#         "mri_data": "concentration",
+#         "subject": "sub-01",
+#         "session": "ses-01",
+#         "segmentation": "aparc+aseg_refined",
+#     }
+
+#     dataframe = generate_stats_dataframe(seg_path, mri_path, info_dict=info)
+
+#     assert not dataframe.empty
+#     assert dataframe["subject"].iloc[0] == "sub-01"
+#     assert dataframe["segmentation"].iloc[0] == "aparc+aseg_refined"
+#     assert dataframe["mri_data"].iloc[0] == "concentration"
+#     assert dataframe["session"].iloc[0] == "ses-01"
+
+
+# def test_compute_mri_stats_cli(capsys, tmp_path: Path, mri_data_dir: Path):
+#     seg_path = (
+#         mri_data_dir
+#         / "mri-processed/mri_processed_data/sub-01"
+#         / "segmentations/sub-01_seg-aparc+aseg_refined.nii.gz"
+#     )
+#     mri_path = (
+#         mri_data_dir
+#         / "mri-processed/mri_processed_data/sub-01"
+#         / "concentrations/sub-01_ses-01_concentration.nii.gz"
+#     )
+#     seg_pattern = (
+#         "(?P<subject>sub-(control|patient)*\\d{2})_seg-(?P<segmentation>[^\\.]+)"
+#     )
+#     mri_data_pattern = "(?P<subject>sub-(control|patient)*\\d{2})_(?P<session>ses-\\d{2})_(?P<mri_data>[^\\.]+)"
+#     timetable = mri_data_dir / "timetable/timetable.tsv"
+#     timetable_sequence = "mixed"
+
+#     args = [
+#         "--segmentation",
+#         str(seg_path),
+#         "--mri",
+#         str(mri_path),
+#         "--output",
+#         str(tmp_path / "mri_stats_output.csv"),
+#         "--timetable",
+#         str(timetable),
+#         "--timelabel",
+#         timetable_sequence,
+#         "--seg_regex",
+#         seg_pattern,
+#         "--mri_regex",
+#         mri_data_pattern,
+#     ]
+
+#     ret = cli.main(["stats", "compute"] + args)
+#     assert ret == 0
+#     captured = capsys.readouterr()
+#     assert "Processing MRIs..." in captured.out
+#     assert "Stats successfully saved to" in captured.out
+#     assert (tmp_path / "mri_stats_output.csv").exists()

From 80d21d539cd0f7eff0d22ddb1112a58596a110a6 Mon Sep 17 00:00:00 2001
From: erasdna <solheim@simula.no>
Date: Fri, 13 Mar 2026 12:49:22 +0100
Subject: [PATCH 2/6] passing tests

---
 src/mritk/concentration/concentration.py | 18 ++++-----
 src/mritk/data/base.py                   |  4 +-
 src/mritk/masking/masks.py               | 14 +++----
 src/mritk/segmentation/segmentation.py   | 16 ++++++--
 src/mritk/t1_maps/t1_maps.py             | 31 +++++++++-------
 test/test_concentration.py               |  9 +++--
 test/test_mri_io.py                      | 14 +++----
 test/test_mri_t1_maps.py                 | 47 ++++++++++++++++++++----
 8 files changed, 100 insertions(+), 53 deletions(-)

diff --git a/src/mritk/concentration/concentration.py b/src/mritk/concentration/concentration.py
index 08c0a67..d4c82b9 100644
--- a/src/mritk/concentration/concentration.py
+++ b/src/mritk/concentration/concentration.py
@@ -10,7 +10,6 @@
 
 import numpy as np
 from ..data.base import MRIData
-from ..data.io import load_mri_data, save_mri_data
 from ..data.orientation import assert_same_space
 
 
@@ -25,20 +24,19 @@ def concentration_from_R1(R1: np.ndarray, R1_0: np.ndarray, r1: float) -> np.nda
 
 
 def concentration(
-    input: Path,
-    reference: Path,
+    T1_mri: MRIData,
+    T10_mri: MRIData,
     output: Optional[Path] = None,
     r1: float = 0.0045,
-    mask: Optional[Path] = None,
+    mask: Optional[MRIData] = None,
 ) -> MRIData:
-    T1_mri = load_mri_data(input, np.single)
-    T10_mri = load_mri_data(reference, np.single)
+    # T1_mri = load_mri_data(input, np.single)
+    # T10_mri = load_mri_data(reference, np.single)
     assert_same_space(T1_mri, T10_mri)
 
     if mask is not None:
-        mask_mri = load_mri_data(mask, bool)
-        assert_same_space(mask_mri, T10_mri)
-        mask_data = mask_mri.data * (T10_mri.data > 1e-10) * (T1_mri.data > 1e-10)
+        assert_same_space(mask, T10_mri)
+        mask_data = mask.data * (T10_mri.data > 1e-10) * (T1_mri.data > 1e-10)
         T1_mri.data *= mask_data
         T10_mri.data *= mask_data
     else:
@@ -50,5 +48,5 @@ def concentration(
     concentrations[mask_data] = concentration_from_T1(T1=T1_mri.data[mask_data], T1_0=T10_mri.data[mask_data], r1=r1)
     mri_data = MRIData(data=concentrations, affine=T10_mri.affine)
     if output is not None:
-        save_mri_data(mri_data, output, np.single)
+        mri_data.save_mri_data(output)
     return mri_data
diff --git a/src/mritk/data/base.py b/src/mritk/data/base.py
index 08e8b53..6afa2e9 100644
--- a/src/mritk/data/base.py
+++ b/src/mritk/data/base.py
@@ -17,8 +17,8 @@ def __init__(self, data: np.ndarray, affine: np.ndarray):
         self.dtype = data.dtype
 
     @classmethod
-    def from_file(cls, filepath: Path):
-        data, affine = load_mri_data(filepath, np.float64)
+    def from_file(cls, filepath: Path, dtype: type = float):
+        data, affine = load_mri_data(filepath, dtype)
         return cls(data=data, affine=affine)
 
     def save_mri_data(self, save_path: Path, intent_code: int | None = None):
diff --git a/src/mritk/masking/masks.py b/src/mritk/masking/masks.py
index 1ead7ae..9c77802 100644
--- a/src/mritk/masking/masks.py
+++ b/src/mritk/masking/masks.py
@@ -40,12 +40,12 @@ def csf_mask(
     use_li: bool = False,
     output: Path | None = None,
 ) -> MRIData:
-    input_vol = load_mri_data(input, dtype=np.single)
-    mask = create_csf_mask(input_vol.data, connectivity, use_li)
+    data, affine = load_mri_data(input, dtype=np.single)
+    mask = create_csf_mask(data, connectivity, use_li)
     assert np.max(mask) > 0, "Masking failed, no voxels in mask"
-    mri_data = MRIData(data=mask, affine=input_vol.affine)
+    mri_data = MRIData(data=mask, affine=affine)
     if output is not None:
-        save_mri_data(mri_data, output, dtype=np.uint8)
+        save_mri_data(mri_data, output, output, dtype=np.uint8)
     return mri_data
 
 
@@ -63,10 +63,10 @@ def intracranial_mask(
     segmentation: Path,
     output: Optional[Path] = None,
 ) -> MRIData:
-    input_csf_mask = load_mri_data(csf_mask, dtype=bool)
-    segmentation_data = load_mri_data(segmentation, dtype=bool)
+    input_csf_mask = MRIData.from_file(csf_mask, dtype=bool)
+    segmentation_data = MRIData.from_file(segmentation, dtype=bool)
     mask_data = create_intracranial_mask(input_csf_mask, segmentation_data)
     mri_data = MRIData(data=mask_data, affine=segmentation_data.affine)
     if output is not None:
-        save_mri_data(mri_data, output, dtype=np.uint8)
+        mri_data.save_mri_data(output)
     return mri_data
diff --git a/src/mritk/segmentation/segmentation.py b/src/mritk/segmentation/segmentation.py
index 74a6272..8b9c941 100644
--- a/src/mritk/segmentation/segmentation.py
+++ b/src/mritk/segmentation/segmentation.py
@@ -13,8 +13,8 @@ def __init__(self, data: np.ndarray, affine: np.ndarray):
         self.rois = np.unique(self.data[self.data > 0])
 
     @classmethod
-    def from_file(cls, filepath: Path):
-        data, affine = load_mri_data(filepath, dtype=int)
+    def from_file(cls, filepath: Path, dtype: type = int):
+        data, affine = load_mri_data(filepath, dtype=dtype)
         return cls(data=data, affine=affine)
 
     @property
@@ -34,10 +34,20 @@ def __init__(self, data: np.ndarray, affine: np.ndarray):
         # Retrieve the FreeSurfer LUT, which contains the labels associated with each ROI
         self.freesurfer_lut = read_lut(None)  # Finds freesurfer LUT
 
-    def get_roi_labels(self, rois: list[int] | None = None) -> pd.DataFrame:
+    def get_roi_labels(self, rois: np.ndarray[int] | None = None) -> pd.DataFrame:
         if not np.isin(rois, self.rois).all():
             raise ValueError("Some of the provided ROIs are not present in the segmentation.")
 
         if rois is None:
             rois = self.rois
         return self.freesurfer_lut.loc[self.freesurfer_lut.index.isin(rois), ["Label"]].rename_axis("ROI").reset_index()
+
+
+class ExtendedFreeSurferSegmentation(FreeSurferSegmentation):
+    def __init__(self, data: np.ndarray, affine: np.ndarray):
+        super().__init__(data, affine)
+
+    def get_roi_labels(self, rois: np.ndarray[int] | None = None) -> pd.DataFrame:
+        rois = self.rois if rois is None else rois
+
+        return super().get_roi_labels(rois % 10000)
diff --git a/src/mritk/t1_maps/t1_maps.py b/src/mritk/t1_maps/t1_maps.py
index a6e4184..7a96203 100644
--- a/src/mritk/t1_maps/t1_maps.py
+++ b/src/mritk/t1_maps/t1_maps.py
@@ -16,7 +16,6 @@
 import scipy
 
 from ..data.base import MRIData
-from ..data.io import load_mri_data, save_mri_data
 from ..masking.masks import create_csf_mask
 from .utils import (
     mri_facemask,
@@ -26,8 +25,7 @@
 )
 
 
-def looklocker_t1map(looklocker_input: Path, timestamps: Path, output: Path | None = None) -> MRIData:
-    LL_mri = load_mri_data(looklocker_input, dtype=np.single)
+def looklocker_t1map(LL_mri: MRIData, timestamps: Path, output: Path | None = None) -> MRIData:
     D = LL_mri.data
     affine = LL_mri.affine
     t_data = np.loadtxt(timestamps) / 1000
@@ -59,13 +57,13 @@ def looklocker_t1map(looklocker_input: Path, timestamps: Path, output: Path | No
     T1map = np.minimum(T1map, T1_ROOF)
     T1map_mri = MRIData(T1map.astype(np.single), affine)
     if output is not None:
-        save_mri_data(T1map_mri, output, dtype=np.single)
+        T1map_mri.save_mri_data(output)
 
     return T1map_mri
 
 
 def looklocker_t1map_postprocessing(
-    T1map: Path,
+    T1map_mri: MRIData,
     T1_low: float,
     T1_high: float,
     radius: int = 10,
@@ -73,7 +71,6 @@ def looklocker_t1map_postprocessing(
     mask: Optional[np.ndarray] = None,
     output: Path | None = None,
 ) -> MRIData:
-    T1map_mri = load_mri_data(T1map, dtype=np.single)
     T1map_data = T1map_mri.data.copy()
     if mask is None:
         # Create mask for largest island.
@@ -106,16 +103,19 @@ def looklocker_t1map_postprocessing(
 
     processed_T1map = MRIData(T1map_data, T1map_mri.affine)
     if output is not None:
-        save_mri_data(processed_T1map, output, dtype=np.single)
+        processed_T1map.save_mri_data(output)
 
     return processed_T1map
 
 
 def mixed_t1map(
-    SE_nii_path: Path, IR_nii_path: Path, meta_path: Path, T1_low: float, T1_high: float, output: Path | None = None
+    SE: MRIData,
+    IR: MRIData,
+    meta_path: Path,
+    T1_low: float,
+    T1_high: float,
+    output: Path | None = None,
 ) -> nibabel.nifti1.Nifti1Image:
-    SE = load_mri_data(SE_nii_path, dtype=np.single)
-    IR = load_mri_data(IR_nii_path, dtype=np.single)
     with open(meta_path, "r") as f:
         meta = json.load(f)
 
@@ -137,10 +137,9 @@ def mixed_t1map(
     return nii
 
 
-def mixed_t1map_postprocessing(SE_nii_path: Path, T1_path: Path, output: Path | None = None) -> nibabel.nifti1.Nifti1Image:
+def mixed_t1map_postprocessing(SE_mri: MRIData, T1_path: Path, output: Path | None = None) -> nibabel.nifti1.Nifti1Image:
     T1map_nii = nibabel.nifti1.load(T1_path)
 
-    SE_mri = load_mri_data(SE_nii_path, np.single)
     mask = create_csf_mask(SE_mri.data, use_li=True)
     mask = skimage.morphology.erosion(mask)
 
@@ -154,8 +153,14 @@ def mixed_t1map_postprocessing(SE_nii_path: Path, T1_path: Path, output: Path |
     return masked_T1map_nii
 
 
+# TODO: Migrate this to MRIData class
 def hybrid_t1map(
-    LL_path: Path, mixed_path: Path, csf_mask_path: Path, threshold: float, erode: int = 0, output: Path | None = None
+    LL_path: Path,
+    mixed_path: Path,
+    csf_mask_path: Path,
+    threshold: float,
+    erode: int = 0,
+    output: Path | None = None,
 ) -> nibabel.nifti1.Nifti1Image:
     mixed_mri = nibabel.nifti1.load(mixed_path)
     mixed = mixed_mri.get_fdata()
diff --git a/test/test_concentration.py b/test/test_concentration.py
index 73936d8..b4a802e 100644
--- a/test/test_concentration.py
+++ b/test/test_concentration.py
@@ -6,9 +6,9 @@
 """
 
 from pathlib import Path
-
+import numpy as np
 from mritk.concentration.concentration import concentration
-
+from mritk.data.base import MRIData
 from mritk.t1_maps.utils import compare_nifti_images
 
 
@@ -29,5 +29,8 @@ def test_intracranial_concentration(tmp_path, mri_data_dir: Path):
     test_outputs = [tmp_path / f"output_ses-0{i}_concentration.nii.gz" for i in sessions]
 
     for i, s in enumerate(sessions):
-        concentration(input=images_path[i], reference=baseline_path, output=test_outputs[i], r1=r1, mask=mask_path)
+        T1_mri = MRIData.from_file(images_path[i], np.single)
+        T10_mri = MRIData.from_file(baseline_path, np.single)
+        mask = MRIData.from_file(mask_path, dtype=bool)
+        concentration(T1_mri=T1_mri, T10_mri=T10_mri, output=test_outputs[i], r1=r1, mask=mask)
         compare_nifti_images(test_outputs[i], ref_outputs[i], data_tolerance=1e-12)
diff --git a/test/test_mri_io.py b/test/test_mri_io.py
index 74743fe..01fd37d 100644
--- a/test/test_mri_io.py
+++ b/test/test_mri_io.py
@@ -12,7 +12,7 @@
 
 
 def test_mri_io_nifti(tmp_path, mri_data_dir):
-    input_file = mri_data_dir / "mri-dataset/mri_dataset/derivatives/sub-01/ses-01/sub-01_ses-01_acq-mixed_T1map.nii.gz"
+    input_file = mri_data_dir / "mri-processed/mri_processed_data/sub-01/T1maps/sub-01_ses-02_T1map_hybrid.nii.gz"
 
     output_file = tmp_path / "output_nifti.nii.gz"
 
@@ -21,20 +21,20 @@ def test_mri_io_nifti(tmp_path, mri_data_dir):
 
 
 def test_MRIData_io(tmp_path, mri_data_dir):
-    input_file = mri_data_dir / "mri-dataset/mri_dataset/derivatives/sub-01/ses-01/sub-01_ses-01_acq-mixed_T1map.nii.gz"
+    input_file = mri_data_dir / "mri-processed/mri_processed_data/sub-01/T1maps/sub-01_ses-02_T1map_hybrid.nii.gz"
 
     output_file = tmp_path / "output_mridata.nii.gz"
 
-    mri_data = MRIData.from_file(input_file, dtype=np.single)
+    mri_data = MRIData.from_file(input_file)
     mri_data.save_mri_data(output_file, intent_code=1006)
 
 
 def test_MRIData_io_invalid_suffix(tmp_path, mri_data_dir):
-    input_file = mri_data_dir / "mri-dataset/mri_dataset/derivatives/sub-01/ses-01/sub-01_ses-01_acq-mixed_T1map.nii.gz"
+    input_file = mri_data_dir / "mri-processed/mri_processed_data/sub-01/T1maps/sub-01_ses-02_T1map_hybrid.nii.gz"
 
     output_file = tmp_path / "output_mridata.invalid"
 
-    mri_data = MRIData.from_file(input_file, dtype=np.single)
+    mri_data = MRIData.from_file(input_file)
     try:
         mri_data.save_mri_data(output_file, intent_code=1006)
         assert False, "Expected ValueError for invalid suffix"
@@ -43,9 +43,9 @@ def test_MRIData_io_invalid_suffix(tmp_path, mri_data_dir):
 
 
 def test_load_mri_data_invalid_suffix(mri_data_dir):
-    input_file = mri_data_dir / "mri-dataset/mri_dataset/derivatives/sub-01/ses-01/sub-01_ses-01_acq-mixed_T1map.invalid"
+    input_file = mri_data_dir / "mri-processed/mri_processed_data/sub-01/T1maps/sub-01_ses-02_T1map_hybrid.invalid"
     try:
-        load_mri_data(input_file, dtype=np.single)
+        load_mri_data(input_file)
         assert False, "Expected ValueError for invalid suffix"
     except ValueError as e:
         assert str(e) == f"Invalid suffix {input_file}, should be either '.nii', or '.mgz'"
diff --git a/test/test_mri_t1_maps.py b/test/test_mri_t1_maps.py
index d40f934..9df9ae4 100644
--- a/test/test_mri_t1_maps.py
+++ b/test/test_mri_t1_maps.py
@@ -6,7 +6,9 @@
 """
 
 from pathlib import Path
+from mritk.data.base import MRIData
 import pytest
+import numpy as np
 
 from mritk.t1_maps.t1_maps import (
     looklocker_t1map,
@@ -23,7 +25,9 @@
 def test_looklocker_t1map(tmp_path, mri_data_dir: Path):
     LL_path = mri_data_dir / "mri-dataset/mri_dataset/sub-01" / "ses-01/anat/sub-01_ses-01_acq-looklocker_IRT1.nii.gz"
     timestamps = (
-        mri_data_dir / "mri-dataset/mri_dataset/sub-01" / "ses-01/anat/sub-01_ses-01_acq-looklocker_IRT1_trigger_times.txt"
+        mri_data_dir
+        / "mri-dataset/mri_dataset/derivatives/sub-01"
+        / "ses-01/anat/sub-01_ses-01_acq-looklocker_IRT1_trigger_times.txt"
     )
     T1_low = 100
     T1_high = 6000
@@ -32,9 +36,22 @@ def test_looklocker_t1map(tmp_path, mri_data_dir: Path):
     test_output_raw = tmp_path / "output_acq-looklocker_T1map_raw.nii.gz"
     test_output = tmp_path / "output_acq-looklocker_T1map.nii.gz"
 
-    looklocker_t1map(looklocker_input=LL_path, timestamps=timestamps, output=test_output_raw)
-    looklocker_t1map_postprocessing(T1map=test_output_raw, T1_low=T1_low, T1_high=T1_high, output=test_output)
-    compare_nifti_images(test_output, ref_output, data_tolerance=1e-12)
+    looklocker_t1map(
+        LL_mri=MRIData.from_file(LL_path, np.single),
+        timestamps=timestamps,
+        output=test_output_raw,
+    )
+    looklocker_t1map_postprocessing(
+        T1map_mri=MRIData.from_file(test_output_raw, np.single),
+        T1_low=T1_low,
+        T1_high=T1_high,
+        output=test_output,
+    )
+    compare_nifti_images(
+        MRIData.from_file(ref_output, np.single),
+        MRIData.from_file(test_output, np.single),
+        data_tolerance=1e-12,
+    )
 
 
 def test_mixed_t1map(tmp_path, mri_data_dir: Path):
@@ -42,7 +59,7 @@ def test_mixed_t1map(tmp_path, mri_data_dir: Path):
     IR_path = mri_data_dir / "mri-dataset/mri_dataset/sub-01" / "ses-01/mixed/sub-01_ses-01_acq-mixed_IR-corrected-real.nii.gz"
     meta_path = mri_data_dir / "mri-dataset/mri_dataset/sub-01" / "ses-01/mixed/sub-01_ses-01_acq-mixed_meta.json"
 
-    ref_output = mri_data_dir / "mri-dataset/mri_dataset/derivatives/sub-01" / "ses-01/sub-01_ses-01_acq-mixed_T1map.nii.gz"
+    ref_output = mri_data_dir / "mri-dataset/mri_dataset/sub-01" / "ses-01/sub-01_ses-01_acq-mixed_T1map.nii.gz"
     test_output_raw = tmp_path / "output_acq-mixed_T1map_raw.nii.gz"
     test_output = tmp_path / "output_acq-mixed_T1map.nii.gz"
 
@@ -50,9 +67,18 @@ def test_mixed_t1map(tmp_path, mri_data_dir: Path):
     T1_high = 10000
 
     mixed_t1map(
-        SE_nii_path=SE_path, IR_nii_path=IR_path, meta_path=meta_path, T1_low=T1_low, T1_high=T1_high, output=test_output_raw
+        SE=MRIData.from_file(SE_path, np.single),
+        IR=MRIData.from_file(IR_path, np.single),
+        meta_path=meta_path,
+        T1_low=T1_low,
+        T1_high=T1_high,
+        output=test_output_raw,
+    )
+    mixed_t1map_postprocessing(
+        SE_mri=MRIData.from_file(SE_path, np.single),
+        T1_path=test_output_raw,
+        output=test_output,
     )
-    mixed_t1map_postprocessing(SE_nii_path=SE_path, T1_path=test_output_raw, output=test_output)
 
     compare_nifti_images(test_output, ref_output, data_tolerance=1e-12)
 
@@ -71,6 +97,11 @@ def test_hybrid_t1map(tmp_path, mri_data_dir: Path):
     erode = 1
 
     hybrid_t1map(
-        LL_path=LL_path, mixed_path=mixed_path, csf_mask_path=csf_mask_path, threshold=threshold, erode=erode, output=test_output
+        LL_path=LL_path,
+        mixed_path=mixed_path,
+        csf_mask_path=csf_mask_path,
+        threshold=threshold,
+        erode=erode,
+        output=test_output,
     )
     compare_nifti_images(test_output, ref_output, data_tolerance=1e-12)

From 7f448e071ee76c33b2657ece8001b164e27fafb7 Mon Sep 17 00:00:00 2001
From: erasdna <solheim@simula.no>
Date: Fri, 13 Mar 2026 12:54:17 +0100
Subject: [PATCH 3/6] Extended FreeSurfer

---
 src/mritk/segmentation/segmentation.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/mritk/segmentation/segmentation.py b/src/mritk/segmentation/segmentation.py
index 8b9c941..e98dff9 100644
--- a/src/mritk/segmentation/segmentation.py
+++ b/src/mritk/segmentation/segmentation.py
@@ -51,3 +51,4 @@ def get_roi_labels(self, rois: np.ndarray[int] | None = None) -> pd.DataFrame:
         rois = self.rois if rois is None else rois
 
         return super().get_roi_labels(rois % 10000)
+        # TODO: Add column to distinguish parenchyma, csf etc

From 606b1a0675a6a3a584bf8d9538c80f83c1f97c6e Mon Sep 17 00:00:00 2001
From: erasdna <solheim@simula.no>
Date: Fri, 13 Mar 2026 13:32:09 +0100
Subject: [PATCH 4/6] add more segmentation clasess

---
 src/mritk/segmentation/segmentation.py | 50 ++++++++++++++++++--------
 test/test_segmentation.py              | 39 ++++++++++++++++++++
 2 files changed, 75 insertions(+), 14 deletions(-)
 create mode 100644 test/test_segmentation.py

diff --git a/src/mritk/segmentation/segmentation.py b/src/mritk/segmentation/segmentation.py
index e98dff9..033ea3b 100644
--- a/src/mritk/segmentation/segmentation.py
+++ b/src/mritk/segmentation/segmentation.py
@@ -7,10 +7,14 @@
 
 
 class Segmentation(MRIData):
-    def __init__(self, data: np.ndarray, affine: np.ndarray):
+    def __init__(self, data: np.ndarray, affine: np.ndarray, lut: pd.DataFrame | None = None):
         super().__init__(data, affine)
         self.data = self.data.astype(int)
         self.rois = np.unique(self.data[self.data > 0])
+        if lut is not None:
+            self.lut = lut
+        else:
+            self.lut = pd.DataFrame({"Label": self.rois}, index=self.rois)
 
     @classmethod
     def from_file(cls, filepath: Path, dtype: type = int):
@@ -25,22 +29,23 @@ def num_rois(self) -> int:
     def roi_labels(self) -> np.ndarray:
         return self.rois
 
+    def get_roi_labels(self, rois: np.ndarray[int] | None = None) -> pd.DataFrame:
+        if rois is None:
+            rois = self.rois
+
+        if not np.isin(rois, self.rois).all():
+            raise ValueError("Some of the provided ROIs are not present in the segmentation.")
+
+        return self.lut.loc[self.lut.index.isin(rois), ["Label"]].rename_axis("ROI").reset_index()
+
 
 class FreeSurferSegmentation(Segmentation):
     def __init__(self, data: np.ndarray, affine: np.ndarray):
-        super().__init__(data, affine)
-
-        # TODO: verify that all labels in the segmentation are valid FS labels
         # Retrieve the FreeSurfer LUT, which contains the labels associated with each ROI
-        self.freesurfer_lut = read_lut(None)  # Finds freesurfer LUT
-
-    def get_roi_labels(self, rois: np.ndarray[int] | None = None) -> pd.DataFrame:
-        if not np.isin(rois, self.rois).all():
-            raise ValueError("Some of the provided ROIs are not present in the segmentation.")
+        lut = read_lut(None)  # Finds freesurfer LUT
+        super().__init__(data, affine, lut)
 
-        if rois is None:
-            rois = self.rois
-        return self.freesurfer_lut.loc[self.freesurfer_lut.index.isin(rois), ["Label"]].rename_axis("ROI").reset_index()
+        # TODO: verify that all labels in the segmentation are valid FS labels
 
 
 class ExtendedFreeSurferSegmentation(FreeSurferSegmentation):
@@ -50,5 +55,22 @@ def __init__(self, data: np.ndarray, affine: np.ndarray):
     def get_roi_labels(self, rois: np.ndarray[int] | None = None) -> pd.DataFrame:
         rois = self.rois if rois is None else rois
 
-        return super().get_roi_labels(rois % 10000)
-        # TODO: Add column to distinguish parenchyma, csf etc
+        freesurfer_labels = super().get_roi_labels(rois % 10000).rename(columns={"ROI": "FreeSurfer_ROI"})
+        tissue_type = self.get_tissue_type(rois)
+        return freesurfer_labels.merge(
+            tissue_type,
+            left_on="FreeSurfer_ROI",
+            right_on="FreeSurfer_ROI",
+            how="outer",
+        ).drop(columns=["FreeSurfer_ROI"])[["ROI", "Label", "tissue_type"]]
+
+    def get_tissue_type(self, rois: np.ndarray[int] | None = None) -> pd.DataFrame:
+        rois = self.rois if rois is None else rois
+        tissue_types = pd.Series(
+            data=np.where(rois < 10000, "Parenchyma", np.where(rois < 20000, "CSF", "Dura")),
+            index=rois,
+            name="tissue_type",
+        )
+        ret = pd.DataFrame(tissue_types, columns=["tissue_type"]).rename_axis("ROI").reset_index()
+        ret["FreeSurfer_ROI"] = ret["ROI"] % 10000
+        return ret
diff --git a/test/test_segmentation.py b/test/test_segmentation.py
new file mode 100644
index 0000000..09d2d9a
--- /dev/null
+++ b/test/test_segmentation.py
@@ -0,0 +1,39 @@
+from mritk.segmentation.segmentation import Segmentation, FreeSurferSegmentation, ExtendedFreeSurferSegmentation
+from pathlib import Path
+import numpy as np
+
+
+def test_segmentation_initialization(example_segmentation: Segmentation):
+    assert example_segmentation.data.shape == (100, 4)
+    assert example_segmentation.affine.shape == (4, 4)
+    assert example_segmentation.num_rois == 3
+    assert set(example_segmentation.roi_labels) == {1, 2, 3}
+    assert example_segmentation.lut.shape == (3, 1)
+    assert set(example_segmentation.lut.columns) == {"Label"}
+
+
+def test_freesurfer_segmentation_labels(mri_data_dir: Path):
+    fs_seg = FreeSurferSegmentation.from_file(
+        mri_data_dir
+        / "mri-processed/mri_processed_data/sub-01\
+            /segmentations/sub-01_seg-aparc+aseg_refined.nii.gz"
+    )
+
+    labels = fs_seg.get_roi_labels()
+    assert not labels.empty
+    assert set(labels["ROI"]) == set(fs_seg.roi_labels)
+
+
+def test_extended_freesurfer_segmentation_labels(example_segmentation: Segmentation, mri_data_dir: Path):
+    data = example_segmentation.data
+    data[0:2, 0:2] = 10001  # csf
+    data[3:5, 3:5] = 20001  # dura
+
+    ext_fs_seg = ExtendedFreeSurferSegmentation(data, affine=np.eye(4))
+    labels = ext_fs_seg.get_roi_labels()
+
+    assert set(labels["ROI"]) == set(ext_fs_seg.roi_labels)
+    assert labels.loc[labels["ROI"] == 10001, "tissue_type"].iloc[0] == "CSF"
+    assert labels.loc[labels["ROI"] == 20001, "tissue_type"].iloc[0] == "Dura"
+    assert labels.loc[labels["ROI"] == 10001, "Label"].iloc[0] == labels.loc[labels["ROI"] == 1, "Label"].iloc[0]
+    assert labels.loc[labels["ROI"] == 20001, "Label"].iloc[0] == labels.loc[labels["ROI"] == 1, "Label"].iloc[0]

From 4c7a1c24f1750d28e0260ed739f9829f9a0454b3 Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henriknf@simula.no>
Date: Sun, 15 Mar 2026 17:48:02 -0700
Subject: [PATCH 5/6] Remove stats.py

---
 src/mritk/stats.py | 423 ---------------------------------------------
 1 file changed, 423 deletions(-)
 delete mode 100644 src/mritk/stats.py

diff --git a/src/mritk/stats.py b/src/mritk/stats.py
deleted file mode 100644
index 7a28f8f..0000000
--- a/src/mritk/stats.py
+++ /dev/null
@@ -1,423 +0,0 @@
-# MRI Statistics Module
-
-# Copyright (C) 2026   Jørgen Riseth (jnriseth@gmail.com)
-# Copyright (C) 2026   Cécile Daversin-Catty (cecile@simula.no)
-# Copyright (C) 2026   Simula Research Laboratory
-
-import argparse
-import re
-import typing
-from pathlib import Path
-
-import numpy as np
-import pandas as pd
-import tqdm.rich
-
-from .data import MRIData
-from .segmentation import default_segmentation_groups, read_lut
-from .testing import assert_same_space
-
-
-def voxel_count_to_ml_scale(affine: np.ndarray):
-    return 1e-3 * np.linalg.det(affine[:3, :3])
-
-
-def find_timestamp(
-    timetable_path: Path,
-    timestamp_sequence: str,
-    subject: str,
-    session: str,
-) -> float:
-    """Find single session timestamp"""
-    try:
-        timetable = pd.read_csv(timetable_path, sep="\t")
-    except pd.errors.EmptyDataError:
-        raise RuntimeError(f"Timetable-file {timetable_path} is empty.")
-    try:
-        timestamp = timetable.loc[
-            (timetable["sequence_label"].str.lower() == timestamp_sequence)
-            & (timetable["subject"] == subject)
-            & (timetable["session"] == session)
-        ]["acquisition_relative_injection"]
-    except ValueError as e:
-        print(timetable)
-        print(timestamp_sequence, subject)
-        raise e
-    return timestamp.item()
-
-
-def prepend_info(df, **kwargs):
-    nargs = len(kwargs)
-    for key, val in kwargs.items():
-        assert key not in df.columns, f"Column {key} already exist in df."
-        df[key] = val
-    return df[[*df.columns[-nargs:], *df.columns[:-nargs]]]
-
-
-def extract_metadata(
-    file_path: Path,
-    pattern: str | None = None,
-    info_dict: dict[str, str] | None = None,
-    required_keys: list[str] | None = None,
-) -> dict:
-    """
-    Extracts metadata from a filename using a regex pattern, falling back to a dictionary.
-
-    Args:
-        file_path (Path): The path to the file.
-        pattern (str, optional): Regex pattern with named capture groups.
-        info_dict (dict, optional): Fallback dictionary if pattern is not provided.
-        required_keys (list[str], optional): Keys to initialize with None if neither match.
-
-    Returns:
-        dict: A dictionary of the extracted metadata.
-
-    Raises:
-        RuntimeError: If a pattern is provided but the filename does not match.
-    """
-    if pattern is not None:
-        if (m := re.match(rf"{pattern}", file_path.name)) is not None:
-            return m.groupdict()
-        else:
-            raise RuntimeError(f"Filename {file_path.name} does not match the provided pattern.")
-
-    required_keys = required_keys or []
-    if info_dict is not None:
-        return {k: info_dict.get(k) for k in required_keys}
-
-    return {k: None for k in required_keys}
-
-
-def get_regions_dictionary(seg_data: np.ndarray, lut_path: Path | None = None) -> dict[str, list[int]]:
-    """
-    Builds a dictionary mapping region descriptions to their corresponding segmentation labels.
-
-    Args:
-        seg_data (np.ndarray): The segmentation array.
-        lut_path (Path, optional): Path to the FreeSurfer Color Look-Up Table.
-
-    Returns:
-        dict[str, list[int]]: Mapping of region names to a list of label integers.
-    """
-    lut = read_lut(lut_path)
-    seg_labels = np.unique(seg_data[seg_data != 0])
-
-    lut_regions = lut.loc[lut.label.isin(seg_labels), ["label", "description"]].to_dict("records")
-
-    regions = {
-        **{d["description"]: sorted([d["label"]]) for d in lut_regions},
-        **default_segmentation_groups(),
-    }
-    return regions
-
-
-def compute_region_statistics(
-    region_data: np.ndarray,
-    labels: list[int],
-    description: str,
-    volscale: float,
-    voxelcount: int,
-) -> dict:
-    """
-    Computes statistical metrics (mean, std, percentiles, etc.) for a specific masked region.
-
-    Args:
-        region_data (np.ndarray): The raw MRI data values mapped to this region (includes NaNs).
-        labels (list[int]): The segmentation label indices representing this region.
-        description (str): Human-readable name of the region.
-        volscale (float): Multiplier to convert voxel counts to milliliters.
-        voxelcount (int): Total number of voxels in the region.
-
-    Returns:
-        dict: A dictionary containing the computed statistics.
-    """
-    record = {
-        "label": ",".join([str(x) for x in labels]),
-        "description": description,
-        "voxelcount": voxelcount,
-        "volume_ml": volscale * voxelcount,
-    }
-
-    if voxelcount == 0:
-        return record
-
-    num_nan = int((~np.isfinite(region_data)).sum())
-    record["num_nan_values"] = num_nan
-
-    if num_nan == voxelcount:
-        return record
-
-    # Filter out NaNs for the mathematical stats
-    valid_data = region_data[np.isfinite(region_data)]
-
-    stats = {
-        "sum": float(np.sum(valid_data)),
-        "mean": float(np.mean(valid_data)),
-        "median": float(np.median(valid_data)),
-        "std": float(np.std(valid_data)),
-        "min": float(np.min(valid_data)),
-        **{f"PC{pc}": float(np.quantile(valid_data, pc / 100)) for pc in [1, 5, 25, 75, 90, 95, 99]},
-        "max": float(np.max(valid_data)),
-    }
-
-    return {**record, **stats}
-
-
-def generate_stats_dataframe(
-    seg_path: Path,
-    mri_path: Path,
-    timestamp_path: str | Path | None = None,
-    timestamp_sequence: str | Path | None = None,
-    seg_pattern: str | None = None,
-    mri_data_pattern: str | None = None,
-    lut_path: Path | None = None,
-    info_dict: dict | None = None,
-) -> pd.DataFrame:
-    """
-    Generates a Pandas DataFrame containing descriptive statistics of MRI data grouped by segmentation regions.
-
-    Args:
-        seg_path (Path): Path to the segmentation NIfTI file.
-        mri_path (Path): Path to the underlying MRI data NIfTI file.
-        timestamp_path (str | Path, optional): Path to the timetable TSV file.
-        timestamp_sequence (str | Path, optional): Sequence label to query in the timetable.
-        seg_pattern (str, optional): Regex to extract metadata from the seg_path filename.
-        mri_data_pattern (str, optional): Regex to extract metadata from the mri_path filename.
-        lut_path (Path, optional): Path to the look-up table.
-        info_dict (dict, optional): Fallback dictionary for metadata.
-
-    Returns:
-        pd.DataFrame: A formatted DataFrame with statistics for all identified regions.
-    """
-    # Load and validate the data
-    mri = MRIData.from_file(mri_path, dtype=np.single)
-    seg = MRIData.from_file(seg_path, dtype=np.int16)
-    assert_same_space(seg, mri)
-
-    # Resolve metadata
-    seg_info = extract_metadata(seg_path, seg_pattern, info_dict, ["segmentation", "subject"])
-    mri_info = extract_metadata(mri_path, mri_data_pattern, info_dict, ["mri_data", "subject", "session"])
-    info = seg_info | mri_info
-
-    # Resolve timestamps
-    info["timestamp"] = None
-    if timestamp_path is not None:
-        try:
-            info["timestamp"] = find_timestamp(
-                Path(str(timestamp_path)),
-                str(timestamp_sequence),
-                str(info.get("subject")),
-                str(info.get("session")),
-            )
-        except (ValueError, RuntimeError, KeyError):
-            pass
-
-    regions = get_regions_dictionary(seg.data, lut_path)
-    volscale = voxel_count_to_ml_scale(seg.affine)
-    records = []
-
-    # Iterate over regions and compute stats
-    for description, labels in tqdm.rich.tqdm(regions.items(), total=len(regions)):
-        region_mask = np.isin(seg.data, labels)
-        voxelcount = region_mask.sum()
-
-        # Extract raw data for this region (including NaNs)
-        region_data = mri.data[region_mask]
-
-        record = compute_region_statistics(
-            region_data=region_data,
-            labels=labels,
-            description=description,
-            volscale=volscale,
-            voxelcount=voxelcount,
-        )
-        records.append(record)
-
-    # Format output
-    dframe = pd.DataFrame.from_records(records)
-    dframe = prepend_info(
-        dframe,
-        segmentation=info.get("segmentation"),
-        mri_data=info.get("mri_data"),
-        subject=info.get("subject"),
-        session=info.get("session"),
-        timestamp=info.get("timestamp"),
-    )
-    return dframe
-
-
-def compute_mri_stats(
-    segmentation: Path,
-    mri: list[Path],
-    output: Path,
-    timetable: Path | None = None,
-    timelabel: str | None = None,
-    seg_regex: str | None = None,
-    mri_regex: str | None = None,
-    lut: Path | None = None,
-    info: str | None = None,
-    **kwargs,
-):
-    import json
-    import sys
-
-    from rich.console import Console
-    from rich.panel import Panel
-
-    # Setup Rich
-    console = Console()
-
-    # Parse info dict from JSON string if provided
-    info_dict = None
-    if info:
-        try:
-            info_dict = json.loads(info)
-        except json.JSONDecodeError:
-            console.print("[bold red]Error:[/bold red] --info must be a valid JSON string.")
-            sys.exit(1)
-
-    if not segmentation.exists():
-        console.print(f"[bold red]Error:[/bold red] Missing segmentation file: {segmentation}")
-        sys.exit(1)
-
-    # Validate all MRI paths before starting
-    for path in mri:
-        if not path.exists():
-            console.print(f"[bold red]Error:[/bold red] Missing MRI file: {path}")
-            sys.exit(1)
-
-    dataframes = []
-
-    # Loop through MRI paths
-    console.print("[bold green]Processing MRIs...[/bold green]")
-    for i, path in enumerate(mri):
-        # console.print(f"[blue]Processing MRI {i + 1}/{len(mri)}:[/blue] {path.name}")
-
-        try:
-            # Call the logic function
-            df = generate_stats_dataframe(
-                seg_path=segmentation,
-                mri_path=path,
-                timestamp_path=timetable,
-                timestamp_sequence=timelabel,
-                seg_pattern=seg_regex,
-                mri_data_pattern=mri_regex,
-                lut_path=lut,
-                info_dict=info_dict,
-            )
-            dataframes.append(df)
-        except Exception as e:
-            console.print(f"[bold red]Failed to process {path.name}:[/bold red] {e}")
-            sys.exit(1)
-
-    if dataframes:
-        final_df = pd.concat(dataframes)
-        final_df.to_csv(output, sep=";", index=False)
-        console.print(
-            Panel(
-                f"Stats successfully saved to:\n[bold green]{output}[/bold green]",
-                title="Success",
-                expand=False,
-            )
-        )
-    else:
-        console.print("[yellow]No dataframes generated.[/yellow]")
-
-
-def get_stats_value(stats_file: Path, region: str, info: str, **kwargs):
-    import sys
-
-    from rich.console import Console
-
-    # Setup Rich
-    console = Console()
-
-    # Validate inputs
-    valid_regions = default_segmentation_groups().keys()
-    if region not in valid_regions:
-        console.print(f"[bold red]Error:[/bold red] Region '{region}' not found in default segmentation groups.")
-        sys.exit(1)
-
-    valid_infos = [
-        "sum",
-        "mean",
-        "median",
-        "std",
-        "min",
-        "max",
-        "PC1",
-        "PC5",
-        "PC25",
-        "PC75",
-        "PC90",
-        "PC95",
-        "PC99",
-    ]
-    if info not in valid_infos:
-        console.print(f"[bold red]Error:[/bold red] Info '{info}' is invalid. Choose from: {', '.join(valid_infos)}")
-        sys.exit(1)
-
-    if not stats_file.exists():
-        console.print(f"[bold red]Error:[/bold red] Stats file not found: {stats_file}")
-        sys.exit(1)
-
-    # Process
-    try:
-        df = pd.read_csv(stats_file, sep=";")
-        region_row = df.loc[df["description"] == region]
-
-        if region_row.empty:
-            console.print(f"[red]Region '{region}' not found in the stats file.[/red]")
-            sys.exit(1)
-
-        info_value = region_row[info].values[0]
-
-        # Output
-        console.print(
-            f"[bold cyan]{info}[/bold cyan] for [bold green]{region}[/bold green] = [bold white]{info_value}[/bold white]"
-        )
-        return info_value
-
-    except Exception as e:
-        console.print(f"[bold red]Error reading stats file:[/bold red] {e}")
-        sys.exit(1)
-
-
-def add_arguments(parser: argparse.ArgumentParser):
-    subparsers = parser.add_subparsers(dest="stats-command", help="Available commands")
-
-    # --- Compute Command ---
-    parser_compute = subparsers.add_parser("compute", help="Compute MRI statistics", formatter_class=parser.formatter_class)
-    parser_compute.add_argument("--segmentation", "-s", type=Path, required=True, help="Path to segmentation file")
-    parser_compute.add_argument("--mri", "-m", type=Path, nargs="+", required=True, help="Path to MRI data file(s)")
-    parser_compute.add_argument("--output", "-o", type=Path, required=True, help="Output CSV file path")
-    parser_compute.add_argument("--timetable", "-t", type=Path, help="Path to timetable file")
-    parser_compute.add_argument("--timelabel", "-l", dest="timelabel", type=str, help="Time label sequence")
-    parser_compute.add_argument(
-        "--seg_regex",
-        "-sr",
-        dest="seg_regex",
-        type=str,
-        help="Regex pattern for segmentation filename",
-    )
-    parser_compute.add_argument("--mri_regex", "-mr", dest="mri_regex", type=str, help="Regex pattern for MRI filename")
-    parser_compute.add_argument("--lut", "-lt", dest="lut", type=Path, help="Path to Lookup Table")
-    parser_compute.add_argument("--info", "-i", type=str, help="Info dictionary as JSON string")
-    parser_compute.set_defaults(func=compute_mri_stats)
-
-    # --- Get Command ---
-    parser_get = subparsers.add_parser("get", help="Get specific stats value", formatter_class=parser.formatter_class)
-    parser_get.add_argument("--stats_file", "-f", type=Path, required=True, help="Path to stats CSV file")
-    parser_get.add_argument("--region", "-r", type=str, required=True, help="Region description")
-    parser_get.add_argument("--info", "-i", type=str, required=True, help="Statistic to retrieve (mean, std, etc.)")
-    parser_get.set_defaults(func=get_stats_value)
-
-
-def dispatch(args: dict[str, typing.Any]):
-    command = args.pop("stats-command")
-    if command == "compute":
-        compute_mri_stats(**args)
-    elif command == "get":
-        get_stats_value(**args)
-    else:
-        raise ValueError(f"Unknown command: {command}")

From 8e5c107a1b28773b6fb1d845b73381f9af3412cb Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henriknf@simula.no>
Date: Sun, 15 Mar 2026 17:49:28 -0700
Subject: [PATCH 6/6] Fix import in __init__

---
 src/mritk/__init__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/mritk/__init__.py b/src/mritk/__init__.py
index 0e93ba9..45be4fb 100644
--- a/src/mritk/__init__.py
+++ b/src/mritk/__init__.py
@@ -13,7 +13,7 @@
     mixed,
     r1,
     segmentation,
-    stats,
+    statistics,
     utils,
 )
 
@@ -34,5 +34,5 @@
     "mixed",
     "hybrid",
     "r1",
-    "stats",
+    "statistics",
 ]