bruise-data 1.2.0__py3-none-any.whl

bruise_data/__init__.py

@@ -0,0 +1,16 @@
+"""bruise_data - one loader for NIJ, Fenwick, bruise mimics and extra bruise images.
+
+    from bruise_data import get_dataloaders
+    train_dl, val_dl, test_dl = get_dataloaders(["nij", "fenwick", "mimics"])
+
+See README.md next to this folder for the full guide.
+"""
+from .checks import check_splits
+from .core import (DEFAULT_ROOT, REQUIRED_COLUMNS, BruiseDataset, dataset_info, get_dataloaders, get_datasets,
+                   get_test_loaders, list_datasets, load_index, read_manifest, summary)
+from .transforms import BruiseTransform, denormalize
+
+__all__ = ["get_dataloaders", "get_datasets", "get_test_loaders", "load_index", "summary", "read_manifest",
+           "check_splits", "BruiseDataset", "BruiseTransform", "denormalize", "list_datasets", "dataset_info",
+           "REQUIRED_COLUMNS", "DEFAULT_ROOT"]
+__version__ = "1.2.0"

bruise_data/__main__.py

@@ -0,0 +1,92 @@
+"""Command line:
+
+    python -m bruise_data datasets                 # every dataset found (any folder with a manifest.csv)
+    python -m bruise_data summary --datasets nij fenwick mimics
+    python -m bruise_data check   [--deep 50]
+    python -m bruise_data preview --datasets nij mimics --out preview.jpg
+    python -m bruise_data classes
+"""
+from __future__ import annotations
+
+import argparse
+
+import numpy as np
+import pandas as pd
+from PIL import Image, ImageDraw
+
+from . import check_splits, dataset_info, get_datasets, list_datasets, load_index, read_manifest, summary
+from .core import DEFAULT_MERGE
+from .transforms import denormalize
+
+
+def main():
+    ap = argparse.ArgumentParser(prog="python -m bruise_data", description=__doc__,
+                                 formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("command", choices=["datasets", "summary", "check", "preview", "classes"])
+    ap.add_argument("--root", default=None, help="BRUISE_DATA folder (default: the folder holding bruise_data/)")
+    ap.add_argument("--datasets", nargs="+", default=None,
+                    help="dataset folder names (default: nij fenwick mimics; check: every dataset found)")
+    ap.add_argument("--task", default="segmentation", choices=["segmentation", "classification"])
+    ap.add_argument("--mimic-classes", nargs="+", default=["all"])
+    ap.add_argument("--annotator", default="primary")
+    ap.add_argument("--test-annotator", default="majority")
+    ap.add_argument("--val-fold", type=int, default=0)
+    ap.add_argument("--deep", type=int, default=0, help="check: open N random images per dataset to compare mask size")
+    ap.add_argument("--split", default="train", choices=["train", "val", "test"], help="preview: which split")
+    ap.add_argument("--n", type=int, default=16, help="preview: number of images")
+    ap.add_argument("--out", default="preview.jpg")
+    a = ap.parse_args()
+    mc = "all" if a.mimic_classes == ["all"] else a.mimic_classes
+    kw = dict(root=a.root, task=a.task, mimic_classes=mc, annotator=a.annotator,
+              test_annotator=a.test_annotator, val_fold=a.val_fold)
+    pd.set_option("display.width", 200)
+
+    chosen = a.datasets or ["nij", "fenwick", "mimics"]
+
+    if a.command == "datasets":
+        rows = []
+        for d in list_datasets(a.root):
+            m = read_manifest(d, a.root)
+            i = dataset_info(d, manifest=m)
+            kind = "annotated bruise" if i["annotators"] else "bruise, no annotators" if i["has_bruise"] else "mimic / negative"
+            rows.append(dict(dataset=d, kind=kind, images=m.image_id.nunique(), masks=i["has_masks"],
+                             primary=i["primary"] or "-", annotators=" ".join(i["annotators"]) or "-",
+                             majority_test=i["has_majority"], folds=i["has_folds"]))
+        print(pd.DataFrame(rows).to_string(index=False))
+    elif a.command == "summary":
+        df = load_index(chosen, **kw)
+        print(summary(df).to_string())
+        print("\nclasses (label index order):", df.attrs["class_names"])
+    elif a.command == "check":
+        raise SystemExit(1 if check_splits(a.root, a.datasets, deep=a.deep) else 0)
+    elif a.command == "classes":
+        m = read_manifest("mimics", a.root)
+        m["class_name"] = m.class_name.replace(DEFAULT_MERGE)
+        print(pd.crosstab([m.class_name, m.source_dataset], m.split).to_string())
+    elif a.command == "preview":
+        tr, va, te = get_datasets(chosen, image_size=256, augment=(a.split == "train"), **kw)
+        ds = {"train": tr, "val": va, "test": te}[a.split]
+        idx = np.random.default_rng(0).choice(len(ds), size=min(a.n, len(ds)), replace=False)
+        tiles = []
+        for i in idx:
+            it = ds[int(i)]
+            img = Image.fromarray((denormalize(it["image"]).permute(1, 2, 0).numpy() * 255).astype(np.uint8))
+            if "mask" in it and it["mask"].sum() > 0:
+                m = Image.fromarray((it["mask"][0].numpy() * 255).astype(np.uint8))
+                red = Image.new("RGB", img.size, (255, 0, 0))
+                img = Image.composite(Image.blend(img, red, 0.45), img, m)
+            d = ImageDraw.Draw(img)
+            d.rectangle([0, 0, 256, 14], fill=(0, 0, 0))
+            d.text((3, 1), f"{it['dataset']}/{it['class_name']}", fill=(255, 255, 255))
+            tiles.append(img)
+        cols = 4
+        rows = (len(tiles) + cols - 1) // cols
+        grid = Image.new("RGB", (cols * 256, rows * 256))
+        for k, t in enumerate(tiles):
+            grid.paste(t, ((k % cols) * 256, (k // cols) * 256))
+        grid.save(a.out, quality=90)
+        print(f"saved {a.out} ({len(tiles)} {a.split} images, red = bruise mask)")
+
+
+if __name__ == "__main__":
+    main()

bruise_data/checks.py

@@ -0,0 +1,94 @@
+"""Sanity checks: missing files, subject leakage between splits, mask/image size."""
+from __future__ import annotations
+
+import random
+from pathlib import Path
+
+import pandas as pd
+from PIL import Image
+
+from .core import DEFAULT_ROOT, dataset_info, list_datasets, read_manifest
+
+
+def check_splits(root: str | Path | None = None, datasets=None, deep: int = 0, verbose: bool = True) -> list[str]:
+    """Run all checks; returns a list of problems (empty list = everything passed).
+
+    datasets: default = every dataset found under root (any folder with a manifest.csv).
+    deep: also open this many random images per dataset and confirm the mask has the image's size.
+    """
+    root = Path(root or DEFAULT_ROOT)
+    datasets = list_datasets(root) if datasets is None else ([datasets] if isinstance(datasets, str) else list(datasets))
+    problems: list[str] = []
+    shared: dict[str, list] = {}  # source_dataset -> [(dataset, manifest)] of annotated sets, e.g. nij + nij_als
+
+    def report(name, bad):
+        if verbose:
+            print(f"  [{'PASS' if not bad else 'FAIL'}] {name}" + (f"  ({len(bad)} problems, e.g. {bad[:3]})" if bad else ""))
+        problems.extend(f"{name}: {b}" for b in bad)
+
+    for d in datasets:
+        p = root / d / "manifest.csv"
+        if not p.exists():
+            if verbose:
+                print(f"{d}: manifest missing - skipped")
+            continue
+        try:
+            m = read_manifest(d, root)
+            info = dataset_info(d, manifest=m)
+        except ValueError as e:
+            report(f"{d} manifest readable", [str(e)])
+            continue
+        if verbose:
+            print(f"{d}: {m.image_id.nunique()} images, {len(m)} manifest rows"
+                  + (f", annotators {list(info['annotators'])} (primary {info['primary']})" if info["annotators"] else ""))
+        report(f"{d} split values are train/val/test", sorted(set(m.split.astype(str)) - {"train", "val", "test"}))
+
+        files = list(m.image_path) + [x for x in m.mask_path if isinstance(x, str) and x]
+        report(f"{d} files exist", [f for f in dict.fromkeys(files) if not (root / f).exists()])
+
+        if info["annotators"]:
+            no_subj = m.subject_id.isna() | (m.subject_id.astype(str).str.strip() == "")
+            report(f"{d} every row has a subject_id (needed for subject-level splits)", list(m[no_subj].image_id.unique()))
+            test_subj = set(m[(m.split == "test") & ~no_subj].subject_id)
+            tv_subj = set(m[(m.split != "test") & ~no_subj].subject_id)
+            report(f"{d} test subjects never in train/val (any annotator)", sorted(test_subj & tv_subj))
+            bad = []
+            for a, g in m[m.split != "test"].groupby("annotator"):
+                multi = g.groupby("subject_id").fold.nunique()
+                bad += [f"{a}:{s}" for s in multi[multi > 1].index]
+            report(f"{d} each subject sits in exactly one fold (so val never shares subjects with train)", bad)
+            dup = m[m.split != "test"].duplicated(["image_id", "annotator"])
+            report(f"{d} no duplicate train/val rows", list(m[m.split != "test"][dup].image_id))
+            shared.setdefault(str(m.source_dataset.mode().iat[0]), []).append((d, m[~no_subj]))
+        else:
+            report(f"{d} image_id unique", list(m[m.image_id.duplicated()].image_id))
+            for col in ("subject_id", "near_dup_group"):
+                g = m[m[col].notna() & (m[col] != "")].groupby(col).split.nunique()
+                report(f"{d} each {col} in one split only", list(g[g > 1].index))
+
+        if deep:
+            rows = m[m.mask_path.notna() & (m.mask_path != "")]
+            idx = random.Random(0).sample(range(len(rows)), min(deep, len(rows))) if len(rows) else []
+            bad = []
+            for i in idx:
+                r = rows.iloc[i]
+                with Image.open(root / r.image_path) as a, Image.open(root / r.mask_path) as b:
+                    if a.size != b.size:
+                        bad.append(f"{r.image_id} image {a.size} mask {b.size}")
+            report(f"{d} mask size == image size ({len(idx)} sampled)", bad)
+
+    # datasets with the same subjects (same source_dataset, e.g. nij + nij_als) are often loaded together
+    for parts in shared.values():
+        if len(parts) < 2:
+            continue
+        names = " + ".join(d for d, _ in parts)
+        m = pd.concat([x for _, x in parts], ignore_index=True)
+        tv = m[m.split != "test"]
+        report(f"{names}: test subjects never in train/val of any of them",
+               sorted(set(m[m.split == "test"].subject_id) & set(tv.subject_id)))
+        multi = tv[tv.is_primary_annotator].groupby("subject_id").fold.nunique()
+        report(f"{names}: a subject has the same fold in all of them (primary annotators)", sorted(multi[multi > 1].index))
+
+    if verbose:
+        print("ALL CHECKS PASSED" if not problems else f"{len(problems)} PROBLEMS FOUND")
+    return problems

bruise_data/core.py

@@ -0,0 +1,422 @@
+"""Index loading, the PyTorch Dataset and the one-call ``get_dataloaders``."""
+from __future__ import annotations
+
+import os
+import warnings
+from pathlib import Path
+from typing import Callable, Iterable
+
+import numpy as np
+import pandas as pd
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler
+
+from .transforms import BruiseTransform
+
+DEFAULT_ROOT = Path(os.environ.get("BRUISE_DATA_ROOT", Path(__file__).resolve().parents[1]))
+SHIPPED = ("nij", "nij_als", "fenwick", "fenwick_als", "mimics", "extra_bruise")  # listed first, in this order
+DEFAULT_MERGE = {"petechiae": "purpura"}
+SPLITS = ("train", "val", "test")
+REQUIRED_COLUMNS = ("image_id", "split", "class_name", "is_bruise", "image_path")
+_NOT_A_PERSON = {"", "nan", "none", "majority"}
+_STR_COLS = {"subject_id": str, "lesion_id": str, "near_dup_group": str, "image_id": str,
+             "mask_path": str, "source_class": str, "ita_category": str, "body_site": str}
+
+
+# ---------------------------------------------------------------------------
+# Index
+# ---------------------------------------------------------------------------
+def list_datasets(root: str | Path | None = None) -> list[str]:
+    """Every dataset under ``root``: any sub-folder that holds a ``manifest.csv``."""
+    root = Path(root or DEFAULT_ROOT)
+    found = sorted(p.parent.name for p in root.glob("*/manifest.csv"))
+    return [d for d in SHIPPED if d in found] + [d for d in found if d not in SHIPPED]
+
+
+def read_manifest(dataset: str, root: str | Path | None = None) -> pd.DataFrame:
+    """Raw ``<root>/<dataset>/manifest.csv`` as a DataFrame (no filtering).
+
+    Only REQUIRED_COLUMNS must exist; any other column the loader uses is filled with a default.
+    """
+    root = Path(root or DEFAULT_ROOT)
+    p = root / dataset / "manifest.csv"
+    if not p.exists():
+        raise FileNotFoundError(f"{p} not found. Is root={root} the BRUISE_DATA folder? "
+                                f"Datasets found there: {list_datasets(root)}")
+    m = pd.read_csv(p, dtype=_STR_COLS, keep_default_na=True, low_memory=False)
+    missing = [c for c in REQUIRED_COLUMNS if c not in m.columns]
+    if missing:
+        raise ValueError(f"{p} is missing required column(s) {missing}")
+    defaults = {"dataset": dataset, "fold": -1, "annotator": "none", "is_primary_annotator": True, "mask_path": "",
+                "source_dataset": dataset, "subject_id": "", "near_dup_group": "", "commercial_use": "unknown"}
+    defaults.update({k: np.nan for k in BruiseDataset.META_FLOATS})
+    for c, v in defaults.items():
+        if c not in m.columns:
+            m[c] = v
+    m["fold"] = pd.to_numeric(m.fold, errors="coerce").fillna(-1).astype(int)
+    m["is_bruise"] = _as_bool(m.is_bruise)
+    m["is_primary_annotator"] = _as_bool(m.is_primary_annotator)
+    return m
+
+
+def _as_bool(s: pd.Series) -> pd.Series:
+    return s.astype(str).str.strip().str.lower().isin(("true", "1", "yes"))
+
+
+def dataset_info(dataset: str, root: str | Path | None = None, manifest: pd.DataFrame | None = None) -> dict:
+    """What a dataset holds, read from its manifest (nothing about a dataset is hard-coded).
+
+    annotators    people with their own train/val masks, primary first (empty tuple if none)
+    primary       the annotator flagged ``is_primary_annotator`` (None if there are no annotators)
+    has_masks     at least one row has a mask_path
+    has_bruise    at least one bruise image (False = a mimic / negative dataset)
+    has_majority  test rows with annotator == "majority" exist
+    has_folds     train/val rows carry a fold number (0..4), so ``val_fold`` can re-split them
+    """
+    m = read_manifest(dataset, root) if manifest is None else manifest
+    tv = m[m.split != "test"]
+    ann = tv.annotator.astype(str).str.strip()
+    names = [a for a in dict.fromkeys(ann) if a.lower() not in _NOT_A_PERSON]
+    primary = [a for a in dict.fromkeys(ann[tv.is_primary_annotator]) if a in names]
+    if names and not primary:
+        if len(names) > 1:
+            raise ValueError(f"{dataset}: annotators {names} but none has is_primary_annotator=True")
+        primary = names
+    if len(primary) > 1:
+        raise ValueError(f"{dataset}: more than one primary annotator {primary}")
+    return dict(name=dataset,
+                annotators=tuple(primary + [a for a in names if a not in primary]),
+                primary=primary[0] if primary else None,
+                has_masks=bool((m.mask_path.fillna("").astype(str).str.len() > 0).any()),
+                has_bruise=bool(m.is_bruise.any()),
+                has_majority=bool((m[m.split == "test"].annotator.astype(str) == "majority").any()),
+                has_folds=bool((tv.fold >= 0).any()))
+
+
+def _pick(value, info: dict, kind: str, known: set) -> str:
+    """Resolve an annotator setting (str / dict / 'primary' / 'majority') for one dataset."""
+    d = info["name"]
+    v = value.get(d, "primary") if isinstance(value, dict) else value
+    if v == "primary":
+        return info["primary"]
+    if v == "majority" and kind == "test":
+        if info["has_majority"]:
+            return "majority"
+        warnings.warn(f"{d} has no majority-vote test masks; scoring against its primary annotator {info['primary']}")
+        return info["primary"]
+    if v not in info["annotators"]:
+        if isinstance(value, dict) or v not in known:
+            raise ValueError(f"{kind} annotator '{v}' does not exist for {d}; choose from {info['annotators']}")
+        return info["primary"]  # a single name only applies to the dataset that has it
+    return v
+
+
+def load_index(
+    datasets: Iterable[str] = ("nij", "fenwick"),
+    root: str | Path | None = None,
+    task: str = "segmentation",
+    mimic_classes: str | Iterable[str] = "all",
+    mimic_sources: Iterable[str] | None = None,
+    exclude_mimic_sources: Iterable[str] | None = None,
+    merge: dict | None = None,
+    min_class_images: int = 50,
+    annotator: str | dict = "primary",
+    test_annotator: str | dict = "majority",
+    val_fold: int = 0,
+    commercial_only: bool = False,
+    binary: bool = False,
+) -> pd.DataFrame:
+    """Build the filtered image table that the loaders use.
+
+    Any folder under ``root`` with a ``manifest.csv`` is a dataset. How it is handled depends only on
+    what its manifest holds (see ``dataset_info``):
+      - with annotators (nij, fenwick, ...): one annotator's train/val masks, val re-split by fold,
+        test masks from ``test_annotator``;
+      - bruise images without annotators (extra_bruise, ...): used as the manifest lists them;
+      - no bruise images at all (mimics, ...): the mimic_* / merge filters apply to all of them together.
+
+    Returns one row per image with a final ``split`` (train/val/test), ``class_name`` and
+    integer ``label``. The class list is in ``df.attrs["class_names"]``.
+    See README.md for what every argument does.
+    """
+    root = Path(root or DEFAULT_ROOT)
+    datasets = [datasets] if isinstance(datasets, str) else list(dict.fromkeys(datasets))
+    available = list_datasets(root)
+    bad = [d for d in datasets if d not in available]
+    if bad:
+        raise ValueError(f"unknown dataset(s) {bad}; datasets in {root}: {available}")
+    if task not in ("segmentation", "classification"):
+        raise ValueError("task must be 'segmentation' or 'classification'")
+    if not 0 <= int(val_fold) <= 4:
+        raise ValueError("val_fold must be 0..4")
+    merge = DEFAULT_MERGE if merge is None else merge
+
+    manifests = {d: read_manifest(d, root) for d in datasets}
+    infos = {d: dataset_info(d, manifest=m) for d, m in manifests.items()}
+    for d, info in infos.items():
+        if task == "segmentation" and info["has_bruise"] and not info["has_masks"]:
+            raise ValueError(f"'{d}' has no masks, so it can only be used with task='classification'")
+    known = {a for info in infos.values() for a in info["annotators"]}
+
+    parts, negatives = [], []
+    for d in datasets:
+        m, info = manifests[d], infos[d]
+        if info["annotators"]:
+            ann = _pick(annotator, info, "train", known)
+            tst = _pick(test_annotator, info, "test", known)
+            tv = m[(m.split != "test") & (m.annotator == ann)].copy()
+            if info["has_folds"]:
+                tv["split"] = np.where(tv.fold == int(val_fold), "val", "train")
+            te = m[(m.split == "test") & (m.annotator == tst)].copy()
+            parts.append(pd.concat([tv, te], ignore_index=True))
+        elif info["has_bruise"]:
+            parts.append(m)
+        else:
+            negatives.append(m)
+
+    if negatives:
+        m = pd.concat(negatives, ignore_index=True)
+        m["class_name"] = m.class_name.replace(merge)
+        if mimic_sources is not None:
+            m = m[m.source_dataset.isin(list(mimic_sources))]
+        if exclude_mimic_sources is not None:
+            m = m[~m.source_dataset.isin(list(exclude_mimic_sources))]
+        available = sorted(m.class_name.unique())
+        if mimic_classes == "all":
+            counts = m.class_name.value_counts()
+            small = sorted(counts[counts < min_class_images].index)
+            if small:
+                warnings.warn(f"mimic classes with < {min_class_images} images left out: {small} "
+                              f"(ask for them by name, or set min_class_images=0)")
+            m = m[~m.class_name.isin(small)]
+        else:
+            want = [mimic_classes] if isinstance(mimic_classes, str) else list(mimic_classes)
+            want = [merge.get(c, c) for c in want]
+            missing = [c for c in want if c not in available]
+            if missing:
+                raise ValueError(f"unknown mimic class(es) {missing}; available: {available}")
+            m = m[m.class_name.isin(want)]
+        parts.append(m)
+
+    df = pd.concat(parts, ignore_index=True)
+    if commercial_only:
+        df = df[df.commercial_use.astype(str).str.lower() == "yes"]
+
+    if binary:  # every non-bruise class becomes one "non_bruise" class; bruise is the positive class (label 1)
+        df["class_name"] = np.where(df.is_bruise, "bruise", "non_bruise")
+        class_names = ["non_bruise", "bruise"]
+    else:
+        mimic_names = sorted(c for c in df.class_name.unique() if c != "bruise")
+        class_names = (["bruise"] if df.is_bruise.any() or task == "segmentation" else []) + mimic_names
+    df["label"] = df.class_name.map({c: i for i, c in enumerate(class_names)}).astype(int)
+    df = df.reset_index(drop=True)
+    df.attrs["class_names"] = class_names
+    df.attrs["root"] = str(root)
+    df.attrs["task"] = task
+    return df
+
+
+def summary(df: pd.DataFrame | None = None, **kwargs) -> pd.DataFrame:
+    """Image counts per dataset / class / split. Pass a ``load_index`` result or its arguments."""
+    df = load_index(**kwargs) if df is None else df
+    df = df.reset_index(drop=True)
+    t = pd.crosstab([df.dataset, df.class_name], df.split).reindex(columns=list(SPLITS), fill_value=0)
+    t["total"] = t.sum(axis=1)
+    t.loc[("ALL", ""), :] = t.sum()
+    return t.astype(int)
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+def _f(v):
+    try:
+        v = float(v)
+        return v
+    except (TypeError, ValueError):
+        return float("nan")
+
+
+class BruiseDataset(Dataset):
+    """One split of the index as a PyTorch Dataset.
+
+    Each item is a dict:
+        image        float tensor (3, H, W)
+        mask         float tensor (1, H, W), 1 = bruise   (segmentation only; all-zero for mimics)
+        label        int, index into ``dataset.class_names``  (0 = bruise; with binary=True, 1 = bruise)
+        class_name, dataset, source_dataset, subject_id, image_id   (strings)
+        is_bruise    bool
+        skin_L, skin_a, skin_b, skin_ita, lesion_L, lesion_a, lesion_b, fitzpatrick, monk_tone
+                     floats (NaN when unknown)
+    """
+
+    META_FLOATS = ("skin_L", "skin_a", "skin_b", "skin_ita", "lesion_L", "lesion_a", "lesion_b", "fitzpatrick", "monk_tone")
+
+    def __init__(self, df: pd.DataFrame, root: str | Path | None = None, task: str | None = None,
+                 transform: Callable | None = None, fast_decode: int | None = None):
+        self.df = df.reset_index(drop=True)
+        self.root = Path(root or df.attrs.get("root") or DEFAULT_ROOT)
+        self.task = task or df.attrs.get("task", "segmentation")
+        self.class_names = list(df.attrs.get("class_names", sorted(df.class_name.unique())))
+        self.transform = transform
+        self.fast_decode = fast_decode
+
+    def __len__(self):
+        return len(self.df)
+
+    def __repr__(self):
+        return f"BruiseDataset({len(self)} images, task={self.task}, classes={self.class_names})"
+
+    def load_image(self, i: int) -> Image.Image:
+        im = Image.open(self.root / self.df.at[i, "image_path"])
+        if self.fast_decode and im.format == "JPEG":
+            im.draft("RGB", (self.fast_decode, self.fast_decode))  # decode big JPEGs at reduced scale
+        return im.convert("RGB")
+
+    def load_mask(self, i: int, size: tuple[int, int]) -> np.ndarray:
+        p = self.df.at[i, "mask_path"]
+        if not isinstance(p, str) or not p:
+            return np.zeros((size[1], size[0]), np.uint8)
+        m = Image.open(self.root / p).convert("L")
+        if m.size != size:
+            m = m.resize(size, Image.NEAREST)
+        return (np.asarray(m) > 0).astype(np.uint8) * 255
+
+    def __getitem__(self, i: int) -> dict:
+        r = self.df.iloc[i]
+        img = self.load_image(i)
+        mask = self.load_mask(i, img.size) if self.task == "segmentation" else None
+        image = np.asarray(img)
+        if self.transform is not None:
+            out = self.transform(image=image, mask=mask) if mask is not None else self.transform(image=image)
+            image, mask = out["image"], out.get("mask", mask)
+        if isinstance(image, np.ndarray):  # custom transform returned numpy
+            image = torch.from_numpy(np.ascontiguousarray(image)).permute(2, 0, 1).float() / 255.0
+        item = {"image": image}
+        if self.task == "segmentation":
+            if isinstance(mask, np.ndarray):
+                mask = torch.from_numpy((mask > 0).astype(np.float32))
+            mask = mask.float()
+            if mask.ndim == 2:
+                mask = mask[None]
+            item["mask"] = (mask > 0).float()
+        item.update(label=int(r.label), class_name=r.class_name, dataset=r.dataset,
+                    source_dataset=str(r.source_dataset), subject_id=str(r.subject_id) if pd.notna(r.subject_id) else "",
+                    image_id=r.image_id, is_bruise=bool(r.is_bruise))
+        for k in self.META_FLOATS:
+            item[k] = _f(r[k])
+        return item
+
+
+# ---------------------------------------------------------------------------
+# Loaders
+# ---------------------------------------------------------------------------
+def _sampler(df: pd.DataFrame, balance: str | None, seed: int):
+    if balance in (None, "none"):
+        return None
+    key = {"dataset": df.dataset, "class": df.class_name, "source": df.dataset + "/" + df.source_dataset.astype(str)}.get(balance)
+    if key is None:
+        raise ValueError("balance must be None, 'dataset', 'class' or 'source'")
+    w = 1.0 / key.map(key.value_counts()).to_numpy(dtype=np.float64)
+    g = torch.Generator().manual_seed(seed)
+    return WeightedRandomSampler(torch.as_tensor(w, dtype=torch.double), num_samples=len(df), replacement=True, generator=g)
+
+
+def get_datasets(datasets=("nij", "fenwick"), image_size: int = 512, resize: str = "pad", augment: bool = True,
+                 train_transform: Callable | None = None, eval_transform: Callable | None = None,
+                 fast_decode: bool = True, **index_kwargs) -> tuple[BruiseDataset, BruiseDataset, BruiseDataset]:
+    """Same arguments as ``get_dataloaders`` minus the DataLoader ones; returns (train, val, test) Datasets."""
+    df = load_index(datasets, **index_kwargs)
+    ttr = train_transform or BruiseTransform(image_size, train=augment, resize=resize)
+    tev = eval_transform or BruiseTransform(image_size, train=False, resize=resize)
+    fd = 2 * image_size if (fast_decode and train_transform is None) else None
+    fde = 2 * image_size if (fast_decode and eval_transform is None) else None
+    out = []
+    for s, t, f in (("train", ttr, fd), ("val", tev, fde), ("test", tev, fde)):
+        part = df[df.split == s].copy()
+        part.attrs = dict(df.attrs)
+        out.append(BruiseDataset(part, transform=t, fast_decode=f))
+    return tuple(out)
+
+
+def get_dataloaders(
+    datasets: Iterable[str] = ("nij", "fenwick"),
+    *,
+    task: str = "segmentation",
+    mimic_classes: str | Iterable[str] = "all",
+    mimic_sources: Iterable[str] | None = None,
+    exclude_mimic_sources: Iterable[str] | None = None,
+    merge: dict | None = None,
+    min_class_images: int = 50,
+    annotator: str | dict = "primary",
+    test_annotator: str | dict = "majority",
+    val_fold: int = 0,
+    commercial_only: bool = False,
+    binary: bool = False,
+    image_size: int = 512,
+    resize: str = "pad",
+    augment: bool = True,
+    train_transform: Callable | None = None,
+    eval_transform: Callable | None = None,
+    batch_size: int = 8,
+    eval_batch_size: int | None = None,
+    num_workers: int = 4,
+    balance: str | None = "dataset",
+    seed: int = 42,
+    pin_memory: bool | None = None,
+    root: str | Path | None = None,
+    verbose: bool = True,
+) -> tuple[DataLoader, DataLoader, DataLoader]:
+    """Return ``(train_loader, val_loader, test_loader)`` for any mix of datasets.
+
+    Example:
+        train_dl, val_dl, test_dl = get_dataloaders(["nij", "fenwick", "mimics"])
+        class_names = train_dl.dataset.class_names
+
+    See README.md for every option.
+    """
+    torch.manual_seed(seed)
+    tr, va, te = get_datasets(datasets, image_size=image_size, resize=resize, augment=augment,
+                              train_transform=train_transform, eval_transform=eval_transform, root=root,
+                              task=task, mimic_classes=mimic_classes, mimic_sources=mimic_sources,
+                              exclude_mimic_sources=exclude_mimic_sources, merge=merge,
+                              min_class_images=min_class_images, annotator=annotator,
+                              test_annotator=test_annotator, val_fold=val_fold, commercial_only=commercial_only,
+                              binary=binary)
+    if verbose:
+        print(summary(pd.concat([tr.df, va.df, te.df], ignore_index=True)).to_string())
+        print("classes:", tr.class_names)
+    if len(tr) == 0:
+        raise ValueError("the training split is empty - check your dataset / class selection")
+    pin = torch.cuda.is_available() if pin_memory is None else pin_memory
+    sampler = _sampler(tr.df, balance, seed)
+    g = torch.Generator().manual_seed(seed)
+    common = dict(num_workers=num_workers, pin_memory=pin, persistent_workers=num_workers > 0)
+    train_dl = DataLoader(tr, batch_size=batch_size, sampler=sampler, shuffle=sampler is None, generator=g,
+                          drop_last=len(tr) > batch_size, worker_init_fn=_seed_worker, **common)
+    eb = eval_batch_size or batch_size
+    val_dl = DataLoader(va, batch_size=eb, shuffle=False, **common)
+    test_dl = DataLoader(te, batch_size=eb, shuffle=False, **common)
+    return train_dl, val_dl, test_dl
+
+
+def get_test_loaders(datasets=("nij", "fenwick"), batch_size: int = 8, num_workers: int = 4,
+                     **kwargs) -> dict[str, DataLoader]:
+    """One test DataLoader per dataset (e.g. {"nij": ..., "fenwick": ..., "mimics": ...})."""
+    kwargs.pop("augment", None)
+    _, _, te = get_datasets(datasets, augment=False, **kwargs)
+    out = {}
+    for d in te.df.dataset.unique():
+        part = te.df[te.df.dataset == d].copy()
+        part.attrs = dict(te.df.attrs)
+        ds = BruiseDataset(part, transform=te.transform, fast_decode=te.fast_decode)
+        out[d] = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+    return out
+
+
+def _seed_worker(worker_id: int):
+    import random
+    s = torch.initial_seed() % 2**32
+    np.random.seed(s)
+    random.seed(s)