[tbx11k] Implements database, improves documentation

pyproject.toml

@@ -117,31 +117,31 @@ indian-f7 = "ptbench.data.indian.fold_7"
 indian-f8 = "ptbench.data.indian.fold_8"
 indian-f9 = "ptbench.data.indian.fold_9"
 
-# TBX11K simplified dataset split 1 (and cross-validation folds)
-tbx11k_simplified = "ptbench.data.tbx11k_simplified.default"
-tbx11k_simplified_f0 = "ptbench.data.tbx11k_simplified.fold_0"
-tbx11k_simplified_f1 = "ptbench.data.tbx11k_simplified.fold_1"
-tbx11k_simplified_f2 = "ptbench.data.tbx11k_simplified.fold_2"
-tbx11k_simplified_f3 = "ptbench.data.tbx11k_simplified.fold_3"
-tbx11k_simplified_f4 = "ptbench.data.tbx11k_simplified.fold_4"
-tbx11k_simplified_f5 = "ptbench.data.tbx11k_simplified.fold_5"
-tbx11k_simplified_f6 = "ptbench.data.tbx11k_simplified.fold_6"
-tbx11k_simplified_f7 = "ptbench.data.tbx11k_simplified.fold_7"
-tbx11k_simplified_f8 = "ptbench.data.tbx11k_simplified.fold_8"
-tbx11k_simplified_f9 = "ptbench.data.tbx11k_simplified.fold_9"
+# TBX11K dataset split 1: healthy vs active tb, and cross-validation folds
+tbx11k-v1-healthy-vs-atb = "ptbench.data.tbx11k.v1_healthy_vs_atb"
+tbx11k-v1-f0 = "ptbench.data.tbx11k.v1_fold_0"
+tbx11k-v1-f1 = "ptbench.data.tbx11k.v1_fold_1"
+tbx11k-v1-f2 = "ptbench.data.tbx11k.v1_fold_2"
+tbx11k-v1-f3 = "ptbench.data.tbx11k.v1_fold_3"
+tbx11k-v1-f4 = "ptbench.data.tbx11k.v1_fold_4"
+tbx11k-v1-f5 = "ptbench.data.tbx11k.v1_fold_5"
+tbx11k-v1-f6 = "ptbench.data.tbx11k.v1_fold_6"
+tbx11k-v1-f7 = "ptbench.data.tbx11k.v1_fold_7"
+tbx11k-v1-f8 = "ptbench.data.tbx11k.v1_fold_8"
+tbx11k-v1-f9 = "ptbench.data.tbx11k.v1_fold_9"
 
-# TBX11K simplified dataset split 2 (and cross-validation folds)
-tbx11k_simplified_v2 = "ptbench.data.tbx11k_simplified_v2.default"
-tbx11k_simplified_v2_f0 = "ptbench.data.tbx11k_simplified_v2.fold_0"
-tbx11k_simplified_v2_f1 = "ptbench.data.tbx11k_simplified_v2.fold_1"
-tbx11k_simplified_v2_f2 = "ptbench.data.tbx11k_simplified_v2.fold_2"
-tbx11k_simplified_v2_f3 = "ptbench.data.tbx11k_simplified_v2.fold_3"
-tbx11k_simplified_v2_f4 = "ptbench.data.tbx11k_simplified_v2.fold_4"
-tbx11k_simplified_v2_f5 = "ptbench.data.tbx11k_simplified_v2.fold_5"
-tbx11k_simplified_v2_f6 = "ptbench.data.tbx11k_simplified_v2.fold_6"
-tbx11k_simplified_v2_f7 = "ptbench.data.tbx11k_simplified_v2.fold_7"
-tbx11k_simplified_v2_f8 = "ptbench.data.tbx11k_simplified_v2.fold_8"
-tbx11k_simplified_v2_f9 = "ptbench.data.tbx11k_simplified_v2.fold_9"
+# TBX11K dataset split 2: others vs active tb, and cross-validation folds
+tbx11k-v2-others-vs-atb = "ptbench.data.tbx11k.v2_others_vs_atb"
+tbx11k-v2-f0 = "ptbench.data.tbx11k.v2_fold_0"
+tbx11k-v2-f1 = "ptbench.data.tbx11k.v2_fold_1"
+tbx11k-v2-f2 = "ptbench.data.tbx11k.v2_fold_2"
+tbx11k-v2-f3 = "ptbench.data.tbx11k.v2_fold_3"
+tbx11k-v2-f4 = "ptbench.data.tbx11k.v2_fold_4"
+tbx11k-v2-f5 = "ptbench.data.tbx11k.v2_fold_5"
+tbx11k-v2-f6 = "ptbench.data.tbx11k.v2_fold_6"
+tbx11k-v2-f7 = "ptbench.data.tbx11k.v2_fold_7"
+tbx11k-v2-f8 = "ptbench.data.tbx11k.v2_fold_8"
+tbx11k-v2-f9 = "ptbench.data.tbx11k.v2_fold_9"
 
 # montgomery-shenzhen aggregated dataset
 montgomery-shenzhen = "ptbench.data.montgomery_shenzhen.default"

src/ptbench/data/tbx11k/datamodule.py

+# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import importlib.resources
+import os
+
+import PIL.Image
+
+from torchvision.transforms.functional import to_tensor
+
+from ...utils.rc import load_rc
+from ..datamodule import CachingDataModule
+from ..split import JSONDatabaseSplit
+from ..typing import DatabaseSplit
+from ..typing import RawDataLoader as _BaseRawDataLoader
+from ..typing import Sample
+
+DatabaseSample = (
+    tuple[str, int] | tuple[str, int, list[tuple[int, int, int, int, int]]]
+)
+"""Type of objects in our JSON representation for this database.
+
+For healthy/sick (no TB)/latent TB cases, each sample is represented by a
+filename, followed by the number 0 (negative class).
+
+For active TB cases, each sample is represented by a filename, followed by the
+number 1, and then 1 or more 5-tuples with radiological finding locations.  The
+radiological findings are defined as such:
+
+* 0/1: Signs of latent TB (0), or active TB (1)
+* xmin
+* ymin
+* width
+* height
+"""
+
+
+class RawDataLoader(_BaseRawDataLoader):
+    """A specialized raw-data-loader for the TBX11k dataset.
+
+    Attributes
+    ----------
+
+    datadir
+        This variable contains the base directory where the database raw data
+        is stored.
+    """
+
+    datadir: str
+
+    def __init__(self):
+        self.datadir = load_rc().get(
+            "datadir.tbx11k", os.path.realpath(os.curdir)
+        )
+
+    def sample(self, sample: DatabaseSample) -> Sample:
+        """Loads a single image sample from the disk.
+
+        Parameters
+        ----------
+
+        sample:
+            A tuple containing the path suffix, within the dataset root folder,
+            where to find the image to be loaded, an integer, representing the
+            sample label, and possible radiological findings represented by
+            bounding boxes.
+
+
+        Returns
+        -------
+
+        sample
+            The sample representation
+        """
+        image = PIL.Image.open(os.path.join(self.datadir, sample[0]))
+        tensor = to_tensor(image)
+
+        # use the code below to view generated images
+        # from torchvision.transforms.functional import to_pil_image
+        # to_pil_image(tensor).show()
+        # __import__("pdb").set_trace()
+
+        return tensor, dict(label=sample[1], name=sample[0])  # type: ignore[arg-type]
+
+    def label(self, sample: DatabaseSample) -> int:
+        """Loads a single image sample label from the disk.
+
+        Parameters
+        ----------
+
+        sample:
+            A tuple containing the path suffix, within the dataset root folder,
+            where to find the image to be loaded, an integer, representing the
+            sample label, and possible radiological findings represented by
+            bounding boxes.
+
+
+        Returns
+        -------
+
+        label
+            The integer label associated with the sample
+        """
+        return sample[1]
+
+
+def make_split(basename: str) -> DatabaseSplit:
+    """Returns a database split for the Montgomery database."""
+
+    return JSONDatabaseSplit(
+        importlib.resources.files(__name__.rsplit(".", 1)[0]).joinpath(basename)
+    )
+
+
+class DataModule(CachingDataModule):
+    """TBX11k dataset for TB detection.
+
+    * Database reference: [TBX11K-2020]_
+    * The original dataset contains samples of healthy, sick (no TB), active
+      and latent TB cases.  There is a total of 11702 samples in the database.
+      Healthy and sick individuals are kept in separate folders.  Latent and
+      active TB cases are merged in the same directory.  One must check the
+      radiological annotations to understand if samples contain either, or both
+      (latent and active TB) signs.
+    * There is one case of patient (file ``imgs/tb/tb1199.png``), that is
+      inside the ``tb`` folder, but contains no annotations.  This sample was
+      **excluded** from our splits.
+    * There are 30 cases of patients that have both active and latent TB
+      radiological signs, over the entire database.  Those samples were also
+      **excluded** from our splits:
+
+      - imgs/tb/tb0135.png
+      - imgs/tb/tb0142.png
+      - imgs/tb/tb0154.png
+      - imgs/tb/tb0167.png
+      - imgs/tb/tb0190.png
+      - imgs/tb/tb0246.png
+      - imgs/tb/tb0255.png
+      - imgs/tb/tb0279.png
+      - imgs/tb/tb0284.png
+      - imgs/tb/tb0350.png
+      - imgs/tb/tb0378.png
+      - imgs/tb/tb0392.png
+      - imgs/tb/tb0395.png
+      - imgs/tb/tb0501.png
+      - imgs/tb/tb0506.png
+      - imgs/tb/tb0526.png
+      - imgs/tb/tb0543.png
+      - imgs/tb/tb0639.png
+      - imgs/tb/tb0640.png
+      - imgs/tb/tb0667.png
+      - imgs/tb/tb0676.png
+      - imgs/tb/tb0713.png
+      - imgs/tb/tb0786.png
+      - imgs/tb/tb0870.png
+      - imgs/tb/tb0875.png
+      - imgs/tb/tb0945.png
+      - imgs/tb/tb0949.png
+      - imgs/tb/tb0968.png
+      - imgs/tb/tb1104.png
+      - imgs/tb/tb1143.png
+
+      - Original train dataset samples:
+          - Healthy: 3000
+          - Sick (but no TB): 3000
+          - Active TB only: 473
+          - Latent TB only: 103
+          - Both active and latent TB: 23
+          - Unknown: 1
+          - Total: 6600
+
+      - Original validation dataset samples:
+          - Healthy: 800
+          - Sick (but no TB): 800
+          - Latent TB only: 36
+          - Active TB only: 157
+          - Both active and latent TB: 7
+          - Total: 1800
+
+      - Original test dataset samples:
+          - Unknown: 3302
+          - Total: 3302
+
+    * Because the test set does not have annotations, we generate train,
+      validation and test datasets as such:
+
+      - The original validation dataset becomes our test set.
+      - The original train dataset is split into new train and validation
+        datasets (validation ratio = 0.203 w.r.t. original train dataset size).
+        The selection of samples is stratified (see comments through our split
+        code, which is shipped alongside this file.)
+
+    * Split v1 contains healthy subjects against active TB cases (total samples
+      = 4430):
+
+      - ``train`` dataset samples:
+          - Healthy: 2390
+          - Active TB only: 377
+          - Total: 2767
+
+      - ``validation`` dataset samples:
+          - Healthy: 610
+          - Active TB only: 96
+          - Total: 706
+
+      - ``test`` dataset samples:
+          - Healthy: 800
+          - Active TB only: 157
+          - Total: 957
+
+    * Split v2 contains healthy, sick (no TB), and latent TB subjects against
+      active TB cases (total samples = 8369):
+
+      - ``train`` dataset samples:
+          - Healthy: 4864
+          - Active TB only: 377
+          - Total: 5241
+
+      - ``validation`` dataset samples:
+          - Active TB only: 96
+          - Healthy: 1239
+          - Total: 1335
+
+      - ``test`` dataset samples:
+          - Healthy: 1636
+          - Active TB only: 157
+          - Total: 1793
+
+
+    Data specifications:
+
+    * Raw data input (on disk):
+
+        * PNG images 8 bits RGB
+        * Resolution: 512x512 pixels
+
+    * Output image:
+
+        * Transforms:
+
+            * Load raw PNG with :py:mod:`PIL`
+
+        * Final specifications:
+
+            * RGB, encoded as a 3-plane image, 8 bits
+            * Square (512x512 px)
+    """
+
+    def __init__(self, split_filename: str):
+        super().__init__(
+            database_split=make_split(split_filename),
+            raw_data_loader=RawDataLoader(),
+        )

src/ptbench/data/tbx11k/make_splits_from_database.py

+#!/usr/bin/env python3
+"""Converts TBX11k JSON annotation files into simplified JSON datasets for
+ptbench.
+
+Requires ``datadir.tbx11k`` to be set on your configuration file, or that you
+are sitting at the root directory of the database.
+
+Because the test set does not have annotations, we generate train, validation
+and test datasets as such:
+
+1. The original validation set becomes the test set.
+2. The original training set is split into new training and validation
+   sets (validation ration = 0.203 by default).  The selection of samples is
+   stratified (respects class proportions in Özgür's way - see comments through
+   the code.)
+
+Our output format is the following:
+
+.. code:: json
+
+   {
+     "train": [
+       [
+         <filename-from-root>,
+         # label is one of:
+         # 0: healthy / 1: active-tb / 2: active-and-latent-tb
+         # 3: latent-tb / 4: sick (no tb)
+         <label>, [
+             # bounding-box annotations follow.  Box-labels are:
+             # 0: latent-tb sign / 1: active-tb sign
+             [<box-label>, <xmin>, <ymin>, <width>, <height>],
+             [0, <xmin>, <ymin>, <width>, <height>],
+             [1, <xmin>, <ymin>, <width>, <height>],
+             ...
+         ],
+       ],
+       ...
+     ],
+     "validation": [
+       # same format as for train
+       ...
+     ]
+     "test": [
+       # same format as for train
+       ...
+     ]
+"""
+
+import collections
+import json
+import os
+import pathlib
+import sys
+import typing
+
+from sklearn.model_selection import StratifiedKFold, train_test_split
+
+
+def reorder(data: dict) -> list:
+    """Reorders data from TBX11K into a sample-based organisation."""
+
+    categories = {k["id"]: k["name"] for k in data["categories"]}
+    assert len(set(categories.values())) == len(
+        categories
+    ), "Category ids are not unique"
+
+    # reset category values, so latent-tb = 0, and active-tb = 1
+    cat_translator = {
+        "ActiveTuberculosis": 1,
+        "ObsoletePulmonaryTuberculosis": 0,
+        "PulmonaryTuberculosis": 2,  # this should NOT exist anywhere!
+    }
+    categories = {k: cat_translator[v] for k, v in categories.items()}
+
+    images = {k["id"]: k["file_name"] for k in data["images"]}
+    assert len(set(images.values())) == len(images), "Image ids are not unique"
+
+    retval: dict[str, list[typing.Any]] = {
+        k["file_name"]: [-1] for k in data["images"]
+    }
+
+    # we now "consume" all annotations and assign each to an image
+    for annotation in data["annotations"]:
+        int_bbox: list[int] = [
+            categories[annotation["category_id"]],
+            *[round(k) for k in annotation["bbox"]],
+        ]
+        retval[images[annotation["image_id"]]].append(int_bbox)
+
+    return sorted([["imgs/" + k, *v] for k, v in retval.items()])
+
+
+def normalize_labels(data: list) -> list:
+    """Decides on the final labels for each sample.
+
+    Categories are decided on the following principles:
+
+    0: healthy, no other bounding box detected, comes from the imgs/health
+       subdir
+    1: active-tb, no latent tb, comes from the imgs/tb subdir, has one or more
+       bounding boxes with label 1, and no bounding box with label 0
+    2: active-tb and latent tb, comes from the imgs/tb subdir, has one or more
+       bounding boxes with label 1 and one or more with label 0
+    3: latent tb, comes from the imgs/tb subdir, has one or more
+       bounding boxes with label 0 and no bounding box with label 1
+    4: sick (but no tb), comes from the imgs/sick subdir, does not have any
+       annotated bounding box.
+    """
+
+    def _set_label(s: list) -> int:
+        if s[0].startswith("imgs/health"):
+            assert (
+                len(s) == 2
+            ), f"Image {s[0]} is healthy, but contains tb bbox annotations"
+            return 0  # patient is healthy
+
+        elif s[0].startswith("imgs/sick"):
+            assert (
+                len(s) == 2
+            ), f"Image {s[0]} is sick (no tb), but contains tb bbox annotations"
+            return 4  # patient is sick
+
+        elif s[0].startswith("imgs/tb"):
+            if len(s) == 2:
+                print(
+                    f"WARNING: Image {s[0]} is from the tb subdir, "
+                    f"but contains no tb bbox annotations"
+                )
+                return -1  # unknown diagnosis
+            bbx_labels: list[int] = [k[0] for k in s[2:]]
+            tb_counts = collections.Counter(bbx_labels)
+            assert 2 not in tb_counts, (
+                f"Label 2 (PulmonaryTuberculosis) was used in image {s[0]} "
+                f"- please check!"
+            )
+            if 0 in tb_counts:
+                if 1 not in tb_counts:
+                    return 3  # patient has latent tb
+                else:
+                    print(
+                        f"WARNING: Image {s[0]} has bboxes with both "
+                        f"active and latent tb."
+                    )
+                    return 2  # patient has active and latent tb
+            else:  # 1 in tb_counts:
+                assert 0 not in tb_counts  # cannot really happen, but check...
+                return 1  # patient has only active tb
+
+        else:
+            raise RuntimeError("Cannot happen - please check")
+
+    return [[k[0], _set_label(k), *k[2:]] for k in data]
+
+
+def print_statistics(d: dict):
+    """Print some statistics about the dataset."""
+
+    label_translations = {
+        -1: "Unknown",
+        0: "Healthy",
+        1: "Active TB only",
+        2: "Both active and latent TB",
+        3: "Latent TB only",
+        4: "Sick (but no TB)",
+    }
+
+    def _print_dataset(ds: list):
+        """Print stats only for the dataset."""
+        class_count = collections.Counter([k[1] for k in ds])
+        for k, v in class_count.items():
+            print(f"  - {label_translations[k]}: {v}")
+        print(f"  - Total: {len(ds)}")
+
+    print("Training set statistics:")
+    _print_dataset(d["train"])
+
+    print("\nValidation set statistics:")
+    _print_dataset(d["validation"])
+
+    print("\nTest set statistics:")
+    _print_dataset(d["test"])
+
+    total_samples = sum(len(ds) for ds in d.values())
+    print(f"\nTotal samples in database: {total_samples}")
+
+
+def create_v1_default_split(d: dict, seed: int, validation_size: float) -> dict:
+    """In the v1 split, we consider active-tb cases against healthy.
+
+    Because the test set is not annotated we do the following:
+
+    1. The original validation set becomes the test set.
+    2. The original training set is split into new training and validation
+       sets.  The selection of samples is stratified (respects class
+       proportions in Özgür's way - see comments)
+
+
+    Parameters
+    ----------
+
+    d
+        The original dataset that will be split
+
+    seed
+        The seed to use at the relevant RNG
+
+    validation_size
+        The proportion of data when we split the training set to make a
+        train and validation sets.
+    """
+
+    # filter cases (only interested in labels 0:healthy or 1:active-tb)
+    use_data = {
+        "train": [k for k in d["train"] if k[1] in (0, 1)],
+        "validation": [k for k in d["validation"] if k[1] in (0, 1)],
+    }
+
+    # Required to repeat Özgür's heuristic with labels that reverse somehow the
+    # sorting for "no_tb" (instead of 0), and "active_tb" (instead of 1).
+    # Reversing the labels used in the stratification process solves this
+    # issue.
+    targets = {0: 1, 1: 0}
+
+    train, val = train_test_split(
+        use_data["train"],
+        test_size=validation_size,
+        random_state=seed,
+        stratify=[targets[k[1]] for k in use_data["train"]],
+    )
+
+    return {
+        "train": train,
+        "validation": val,
+        "test": use_data["validation"],
+    }
+
+
+def create_v2_default_split(d: dict, seed: int, validation_size) -> dict:
+    """In the v2 split, we consider active-tb cases against healthy, sick and
+    latent-tb cases.
+
+    Because the test set is not annotated we do the following:
+
+    1. The original validation set becomes the test set.
+    2. The original training set is split into new training and validation
+       sets.  The selection of samples is stratified (respects class
+       proportions in Özgür's way - see comments)
+    """
+
+    # filter cases (only interested in labels 0:healthy or 1:active-tb)
+    use_data = {
+        "train": [k for k in d["train"] if k[1] in (0, 1, 3, 4)],
+        "validation": [k for k in d["validation"] if k[1] in (0, 1, 3, 4)],
+    }
+
+    # Required to repeat Özgür's heuristic with labels that reverse somehow the
+    # sorting for "no_tb" (instead of 0, 3 or 4), and "active_tb" (instead of
+    # 1). Reversing the labels used in the stratification process solves this
+    # issue.
+    targets = {0: 1, 1: 0, 3: 1, 4: 1}
+
+    train, val = train_test_split(
+        use_data["train"],
+        test_size=validation_size,
+        random_state=seed,
+        stratify=[targets[k[1]] for k in use_data["train"]],
+    )
+
+    # These are the targets that will show up in the split.  We make everything
+    # that is not active-tb to be label=0.
+    split_targets = {0: 0, 1: 1, 3: 0, 4: 0}
+    return {
+        "train": [[k[0], split_targets[k[1]], *k[2:]] for k in train],
+        "validation": [[k[0], split_targets[k[1]], *k[2:]] for k in val],
+        "test": [
+            [k[0], split_targets[k[1]], *k[2:]] for k in use_data["validation"]
+        ],
+    }
+
+
+def create_folds(
+    d: dict, n: int, seed: int, validation_size: float
+) -> list[dict]:
+    """Creates folds from existing splits.
+
+    Parameters
+    ----------
+
+    d
+        The original split to consider
+
+    n
+        The number of folds to produce
+
+
+    Returns
+    -------
+
+    folds
+        All the ``n`` folds
+    """
+
+    X = d["train"] + d["validation"] + d["test"]
+    y = [[k[1]] for k in X]
+
+    # Initializes a StratifiedKFold object with 10 folds
+    skf = StratifiedKFold(n_splits=n, shuffle=True, random_state=seed)
+
+    # Required to repeat Özgür's heuristic with labels that reverse somehow the
+    # sorting for "no_tb" (instead of 0), and "active_tb" (instead of 1).
+    # Reversing the labels used in the stratification process solves this
+    # issue.
+    targets = {0: 1, 1: 0}
+
+    # Loops over the 10 folds and split the data
+    retval = []
+    for train_idx, test_idx in skf.split(X, y):
+        # Get the training and test data for this fold
+        train_dataset = [X[k] for k in train_idx]
+        test_dataset = [X[k] for k in test_idx]
+
+        # Split the training data into training and validation sets
+        train_dataset, val_dataset = train_test_split(
+            train_dataset,
+            test_size=validation_size,
+            random_state=seed,
+            stratify=[targets[k[1]] for k in train_dataset],
+        )
+
+        retval.append(
+            {
+                "train": train_dataset,
+                "validation": val_dataset,
+                "test": test_dataset,
+            }
+        )
+
+    return retval
+
+
+def main():
+    if len(sys.argv) != 1:
+        print(__doc__)
+        print(f"Usage: python3 {sys.argv[0]} ")
+        sys.exit(0)
+
+    # program constants used by Özgür
+    seed = 42  # used to seed the relevant RNG
+    validation_size = 0.203  # proportion for test when splitting
+    n_folds = 10  # number of folds to create
+
+    from clapper.rc import UserDefaults
+
+    datadir = pathlib.Path(
+        UserDefaults("ptbench.toml").get(
+            "datadir.tbx11k", os.path.realpath(os.curdir)
+        )
+    )
+    train_filename = datadir / "annotations" / "json" / "TBX11K_train.json"
+    val_filename = datadir / "annotations" / "json" / "TBX11K_val.json"
+    test_filename = datadir / "annotations" / "json" / "all_test.json"
+
+    with open(train_filename) as f:
+        print(f"Loading {str(train_filename)}...")
+        data = json.load(f)
+        train_data = normalize_labels(reorder(data))
+
+    with open(val_filename) as f:
+        print(f"Loading {str(val_filename)}...")
+        data = json.load(f)
+        val_data = normalize_labels(reorder(data))
+
+    with open(test_filename) as f:
+        print(f"Loading {str(test_filename)}...")
+        data = json.load(f)
+        test_data = reorder(data)
+
+    final_data = {
+        "train": train_data,
+        "validation": val_data,
+        "test": test_data,
+    }
+    print_statistics(final_data)
+
+    # No need to record the re-processed data.
+    # with open(sys.argv[4], "w") as fout:
+    #     json.dump(final_data, fout, indent=2)
+
+    print("\nGenerating v1 split...")
+    v1_split = create_v1_default_split(
+        final_data, seed=seed, validation_size=validation_size
+    )
+    print_statistics(v1_split)
+    with open("v1-healthy-vs-atb.json", "w") as v1def:
+        json.dump(v1_split, v1def, indent=2)
+
+    # folds for the v1 split
+    print(f"\nGenerating {n_folds} v1 split folds...")
+    v1_folds = create_folds(
+        v1_split, n=n_folds, seed=seed, validation_size=validation_size
+    )
+    for i, k in enumerate(v1_folds):
+        with open(f"v1-fold-{i}.json", "w") as v1fold:
+            json.dump(k, v1fold, indent=2)
+
+    print("\nGenerating v2 split...")
+    v2_split = create_v2_default_split(
+        final_data, seed=seed, validation_size=validation_size
+    )
+    print_statistics(v2_split)
+    with open("v2-others-vs-atb.json", "w") as v2def:
+        json.dump(v2_split, v2def, indent=2)
+
+    # folds for the v2 split
+    print(f"\nGenerating {n_folds} v2 split folds...")
+    v2_folds = create_folds(
+        v2_split, n=n_folds, seed=seed, validation_size=validation_size
+    )
+    for i, k in enumerate(v2_folds):
+        with open(f"v2-fold-{i}.json", "w") as v2fold:
+            json.dump(k, v2fold, indent=2)
+
+
+if __name__ == "__main__":
+    main()

src/ptbench/data/tbx11k/v1_fold_0.py

+# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from .datamodule import DataModule
+
+datamodule = DataModule("v1-fold-0.json")

src/ptbench/data/tbx11k/v1_fold_1.py

+# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from .datamodule import DataModule
+
+datamodule = DataModule("v1-fold-1.json")

src/ptbench/data/tbx11k/v1_fold_2.py

+# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from .datamodule import DataModule
+
+datamodule = DataModule("v1-fold-2.json")

src/ptbench/data/tbx11k/v1_fold_3.py

+# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from .datamodule import DataModule
+
+datamodule = DataModule("v1-fold-3.json")

src/ptbench/data/tbx11k/v1_fold_4.py

+# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from .datamodule import DataModule
+
+datamodule = DataModule("v1-fold-4.json")