Added the benchmarking script (Not working properly)

Added metrics for age bracket (`n_off_accuracy`), as well as
binary/multiclass false-{positive,negative} rates (`binary_fpr_fnr`
and `multiclass_fpr_fnr`). Also, added mutual agreement between models
`agreement_fraction` and `agreement_elements`.

.gitignore

@@ -3,3 +3,6 @@
 *.egg-info
 landmark_models
 *__pycache__/
+data/*
+models/*
+gpt*.py

README.md

@@ -0,0 +1,6 @@
+# ISR Face Dataset/Model Bias Check API
+
+We use [FairFace](https://github.com/dchen236/FairFace) and [MiVOLO](https://github.com/wildchlamydia/mivolo) version 1, face-only checkpoint.
+
+## Dataset
+Download the tar file `VISTEAM-NAS/Public_Data/facing2-skin-tone-train-images.tar.bz2` to the `data` directory, and extract it. This dataset has balanced sex and skin-tone, and unbalanced age.

benchmark.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import logging
+from pathlib import Path
+
+import mediapipe as mp
+import numpy as np
+import pandas as pd
+import torch.nn as nn
+
+from facebias import load_dataset
+from facebias.detectors import get_face_boxes
+from facebias.detectors.mediapipe import MediapipeDetector
+from facebias.estimators.fairface import FairFace
+from facebias.estimators.mivolov1 import MiVOLOv1
+from facebias.metrics import calc_metrics_per_subgroup, calc_metrics
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(f"facebias:{__file__}")
+
+
+if __name__ == "__main__":
+    import os
+    logger.info(os.getcwd())
+
+    DATASET_PATH = Path("data/facing2-train/")
+    METADATA_PATH = DATASET_PATH / "meta-w-age.csv"
+    DETECTOR_PATH = Path("models/blaze_face_full_range.tflite")
+    # TEST_IMS: list[str] = ["10", "12", "14", "9"]
+
+    detector = MediapipeDetector(str(DETECTOR_PATH))
+    imdict, meta = load_dataset(
+        DATASET_PATH, meta_path=METADATA_PATH, imname_proc_fn=lambda x: x.split("_")[0]
+    )
+    for im, feats in meta.items():
+        age = int(feats["age"])
+        d = age // 10 * 10
+        feats["age_group"] = "{}-{}".format(d, d + 9)
+        feats["age"] = float(feats["age"])
+
+    face_bboxes = get_face_boxes(imdict, detector)
+
+    # for t in TEST_IMS:
+    #     logger.info("-- {} - {}".format(t, meta[str(t)]))
+
+    print(FairFace.capabilities())
+    model_ff = FairFace(Path("models/fairface_alldata_4race_20191111.pt"), device="cpu")
+    preds_ff = model_ff.predict(imdict)
+    metrics_ff = calc_model_performance(meta, preds_ff)
+
+    # logger.info("FairFace -- Test Images")
+    # for t in TEST_IMS:
+    #     logger.info("--{} - {}".format(t, preds_ff[str(t)]))
+
+    metrics_ff_groups = calc_metrics_per_subgroup(meta, preds_ff)
+    for k, v in metrics_ff_groups.items():
+        for kv, vv in v.items():
+            print(k, kv, vv)
+
+    print(MiVOLOv1.capabilities())
+    model_mv = MiVOLOv1(
+        Path("models/volo-v1_model_imdb_age_gender_4.22.pth.tar"), device="cpu"
+    )
+    preds_mv = model_mv.predict(imdict)
+
+    # logger.info("MiVOLOv1(Face Only) -- Test Images")
+    # for t in TEST_IMS:
+    #     logger.info("{} - {}".format(t, preds_mv[str(t)]))
+
+    metrics_mv = calc_model_performance(meta, preds_mv)
+    metrics_mv_groups = calc_metrics_per_subgroup(meta, preds_mv)
+    for k, v in metrics_mv_groups.items():
+        for kv, vv in v.items():
+            print(k, kv, vv)

src/facebias/__init__.py

@@ -10,8 +10,7 @@ from typing import Any, Callable
 import cv2
 import numpy as np
 
-logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(f"facebias:{__name__}")
-logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -22,7 +21,12 @@ class FaceBox:
     y2: int
 
 
-def load_metadata(p: Path, key_id="image", key_proc_fn=None) -> dict[str, dict[str, str]]:
+# TODO(gschardong): Convert all CSV reading functions to pandas
+
+
+def load_metadata(
+    p: Path, key_id="image", key_proc_fn=None
+) -> dict[str, dict[str, str]]:
     lines = []
     with open(p, newline="") as csvfile:
         dialect = csv.Sniffer().sniff(csvfile.read(1024))
@@ -40,9 +44,7 @@ def load_metadata(p: Path, key_id="image", key_proc_fn=None) -> dict[str, dict[s
 
 
 def load_dataset(
-        root: Path,
+    root: Path, meta_path: Path | None, imname_proc_fn: Callable | None
-        meta_path: Path | None,
-        imname_proc_fn: Callable |None
 ) -> tuple[dict[str, np.ndarray], dict[str, dict[str, Any]] | None]:
     """
     if `meta_path` is `None`, we won't attempt to read it.
@@ -82,7 +84,9 @@ def load_dataset(
         except cv2.error:
             logger.info(f'File "{p}" is not an image. Skipping.')
         else:
-            proc_imname = imname_proc_fn(p.name) if imname_proc_fn is not None else str(p.name)
+            proc_imname = (
+                imname_proc_fn(p.name) if imname_proc_fn is not None else str(p.name)
+            )
             ims[proc_imname] = im
 
     if not metadata:
@@ -91,3 +95,157 @@ def load_dataset(
             logger.error(f'Metadata file not found at "{meta_path}".')
 
     return ims, metadata
+
+
+# def calc_model_performance(
+#     gt: pd.DataFrame,
+#     preds: pd.DataFrame,
+#     keys: list[str] | None = None,
+#     possible_caps: dict[Capability, Any] | None = None,
+# ) -> pd.DataFrame:
+#     """
+#     We assume that both `gt` and `preds` have the same structure. They should
+#     be indexed by individual ID, such as the image name, and each value is a
+#     dictionary with model prediction capabilities as keys (e.g., "age_group",
+#     "sex", "skin-color", etc.), and the values are the predictions, or ground-truth
+#     values for each ID/capability.
+
+#     if `keys` is empty, then we infer from common keys present in `preds` and `gt`.
+
+#     Parameters
+#     ----------
+#     gt: pd.DataFrame
+#     preds: pd.DataFrame
+#     keys: list[str] | None
+
+#     Returns
+#     -------
+#     metrics: pd.DataFrame
+#     """
+#     common_caps = keys
+#     if keys is None:
+#         common_caps = set(gt.columns) & set(preds.columns)
+#         if not common_caps:
+#             logger.error(
+#                 f'No common capabilities found. Predictions has "{preds.columns}",'
+#                 f' ground-truth has "{gt.columns}".'
+#             )
+#             return None
+
+#     # Finding common images between predictions and ground-truth.
+#     common_inds = set(preds.index) & set(gt.index)
+#     if not common_inds:
+#         logger.error("No common images found between predictions and ground-truth.")
+#         return None
+
+#     metric_vals = dict()
+#     for cap in common_caps:
+#         if isinstance(preds[cap].iloc[0], (float, int)):
+#             metric_vals[cap] = {
+#                 "mean_absolute_error": mean_absolute_error(gt[cap], preds[cap]),
+#                 "max_error": max_error(gt[cap], preds[cap]),
+#             }
+#         else:
+#             labels = possible_caps[cap]
+#             if possible_caps is None:
+#                 labels = sorted(list(set(preds[cap].unique()) | set(gt[cap].unique())))
+#             metric_vals[cap] = {
+#                 "accuracy": accuracy_score(gt[cap], preds[cap]),
+#                 "balanced_accuracy": balanced_accuracy_score(gt[cap], preds[cap]),
+#                 "cohen-kappa": cohen_kappa_score(gt[cap], preds[cap], labels=labels),
+#             }
+
+#     return pd.DataFrame.from_dict(metric_vals)
+
+
+# def calc_metrics_per_subgroup(
+#     gt: pd.DataFrame,
+#     preds: pd.DataFrame,
+#     model_cls: BaseEstimator,
+#     metrics: list[str] = [
+#         "accuracy",
+#     ],
+# ) -> pd.DataFrame:
+#     """Calculate performance metrics per sub-group for each capability.
+
+#     Parameters
+#     ----------
+#     gt: pd.DataFrame
+
+#     preds: pd.DataFrame
+
+#     model_cls: BaseEstimator-derived class
+
+#     Returns
+#     -------
+#     metrics: dict[Capability, dict[Any, dict]]
+#     """
+#     common_caps = set(gt.columns) & set(preds.columns)
+
+#     # metrics = {}
+#     index = sorted(
+#         [
+#             model_cls.possible_capability_values(c)
+#             for c in common_caps
+#             if c != Capability.AGE
+#         ],
+#         key=len,
+#     )
+#     index = product(*index)
+
+#     metrics = ["accuracy", ""]
+#     df = pd.DataFrame(index=index, columns=metrics)
+
+#     for cap in common_caps:
+#         # TODO(gschardong): Better to store the "type" of each capability
+#         # somewhere and test all numeric types here.
+#         if cap == Capability.AGE:
+#             continue
+
+#         other_caps = common_caps - set([cap])
+
+#         # TODO(gschardong): Do we only need the values that occur in the data,
+#         # or all possible values? If the first is true, then we need to fetch
+#         # from the model class itself, else, we keep it as is.
+#         unique_values_cap = set(gt[cap].unique()) | set(preds[cap].unique())
+
+#         metrics[cap] = {}
+#         for val in unique_values_cap:
+#             ids = gt.index[gt[cap] == val]
+#             metrics[cap][val] = {"number_of_elements": len(ids)}
+#             for ocap in other_caps:
+#                 metrics[cap][val][ocap] = {}
+#                 fpred_data = preds[ocap][ids]
+#                 fgt_data = gt[ocap][ids]
+
+#                 if isinstance(fpred_data[0], (float, int)):
+#                     # If data is numeric, we calculate regression-based metrics
+#                     metrics[cap][val][ocap] = {
+#                         "mean_absolute_error": mean_absolute_error(
+#                             fgt_data, fpred_data
+#                         ),
+#                         "max_error": max_error(fgt_data, fpred_data),
+#                     }
+#                 else:
+#                     unique_values_ocap = sorted(
+#                         list(set(gt[ocap].unique()) | set(preds[ocap].unique()))
+#                     )
+#                     unique_values_ocap = np.array(unique_values_ocap)
+
+#                     metrics_small = {}
+#                     # if len(fgt_data.unique()) == 2:
+#                     #     cm = confusion_matrix(fgt_data, fpred_data, labels=unique_values_ocap)
+#                     # else:
+#                     cm = multilabel_confusion_matrix(
+#                         fgt_data, fpred_data, labels=unique_values_ocap
+#                     )
+#                     for m, oval in zip(cm, unique_values_ocap):
+#                         # metrics[cap][val][ocap][oval] = m
+#                         metrics_small[oval] = m
+#                     return m
+
+#                     metrics[cap][val][ocap] = {
+#                         "accuracy": accuracy_score(fgt_data, fpred_data),
+#                     }
+
+#     return metrics

src/facebias/estimators/__init__.py

@@ -12,7 +12,7 @@ class Capability(StrEnum):
     AGE = "age"
     AGEGROUP = "age_group"
     SEX = "sex"
-    SKINCOLOR = "skin_color"
+    SKINTONE = "skin_tone"
     ETHNICITY = "ethinicity"
 
 
@@ -23,7 +23,9 @@ class InvalidCapabilityError(Exception):
 
 class BaseEstimator(ABC):
     @abstractmethod
-    def predict(self, images: dict[Path, np.ndarray]) -> dict[Path, dict[Capability, Any]]:
+    def predict(
+        self, images: dict[Path, np.ndarray]
+    ) -> dict[Path, dict[Capability, Any]]:
         """Runs the estimator on a batch of images.
 
         The input `images` is a dictionary indexed by the image name, or ID,

src/facebias/estimators/fairface.py

@@ -70,9 +70,9 @@ class FairFace(BaseEstimator):
             y = y.cpu().detach().squeeze().numpy()
 
             # Ethnicity prediction
-            y_ethno = y[:4]
+            # y_ethno = y[:4]
-            ethno_score = np.exp(y_ethno) / np.sum(np.exp(y_ethno))
+            # ethno_score = np.exp(y_ethno) / np.sum(np.exp(y_ethno))
-            ethno_pred = np.argmax(ethno_score)
+            # ethno_pred = np.argmax(ethno_score)
 
             # Age prediction
             y_age = y[9:18]
@@ -96,9 +96,9 @@ class FairFace(BaseEstimator):
     def capabilities() -> list[Capability]:
         return [Capability.AGEGROUP, Capability.SEX] #, Capability.ETHNICITY]
 
-    def possible_capability_values(cap: Capability) -> list:
+    def possible_capability_values(cap: Capability) -> list[str]:
         if cap == Capability.AGEGROUP:
-            return ["0-2", "3-9", "10-19", "20-29", "30-39", "40-49", "50-59", "60-69", "70+"]
+            return ["00-02", "03-09", "10-19", "20-29", "30-39", "40-49", "50-59", "60-69", "70+"]
         elif cap == Capability.SEX:
             return ['m', 'f']
         # elif cap == Capability.ETHNICITY:
@@ -108,8 +108,7 @@ class FairFace(BaseEstimator):
 
 
 def _to_age_label(age):
-    return FairFace.possible_classes(Capability.AGEGROUP)[age]
+    return FairFace.possible_capability_values(Capability.AGEGROUP)[age]
-    # return "{}-{}".format(age * 10, age * 10 + 9)
 
 
 def _to_ethno_label(val):

src/facebias/estimators/mivolo/README.md

@@ -0,0 +1,3 @@
+# Subset of MiVOLO code
+
+This is a subset of the [MiVOLO](https://github.com/wildchlamydia/mivolo) code necessary to instantiate the face-only attribute model.

src/facebias/evaluation.py

@@ -0,0 +1,214 @@
+# -*- coding: utf-8 -*-
+
+"""Dataset/model evaluation functions."""
+
+import logging
+from itertools import permutations, combinations
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from scipy.stats import entropy
+
+from facebias.estimators import Capability
+from facebias.metrics import gini
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("facebias:evaluation.py")
+
+
+# TODO(gschardong): Move to the same file as `load_dataset`
+def _to_age_bracket(row):
+    iage = int(row["age"])
+    if iage < 3:
+        return "00-02"
+    elif iage < 10:
+        return "03-09"
+    elif iage > 69:
+        return "70+"
+
+    d = iage // 10 * 10
+    return "{}-{}".format(d, d + 9)
+
+
+def generate_description(dataset: pd.DataFrame, name: str | None = None) -> tuple[str]:
+    """Generates a textual description of `dataset`, its variables and values.
+
+    For all supported variables in `dataset`, this function lists their names,
+    types, value ranges and descriptive statistics. It also groups each
+    variable, comparing it to the others, thus creating a subgroup description
+    as well.
+
+    Note that we only support the variables defined by the `Capability` enum.
+
+    Parameters
+    ----------
+    dataset: pd.DataFrame
+
+    name: str, optional
+        The dataset name. Only ever used in the first line of the returned
+        description, for pretty printing.
+
+    Returns
+    -------
+    desc: str
+        The textual description as a human-readable string.
+
+    See Also
+    --------
+    `Capability`
+    """
+    outstr = ""
+    if name is not None:
+        outstr = f"The dataset {name}"
+    else:
+        outstr = f"The dataset"
+
+    outstr += f" has a total of {len(dataset)} {dataset.index.name}s, with the following supported features/capabilities:"
+
+    caps = []
+    for c in Capability:
+        if c.value in dataset:
+            caps.append(c)
+            outstr += f"\n- {c.value}"
+
+    outstr += "\n\nEach feature/capability has the following types and values:"
+    for c in caps:
+        if c == Capability.AGE:
+            outstr += f"\n{c.value}: numeric"
+        else:
+            outstr += f"\n{c.value}: categorical"
+            outstr += f"\n - {sorted(dataset[c].unique())}"
+
+    outstr += "\n\nData distribution statistics."
+    for c in caps:
+        outstr += (
+            f'\nThe feature/capability "{c}" has the following distribution of values:'
+        )
+        if c == Capability.AGE:
+            m1 = dataset[c].min()
+            m2 = dataset[c].max()
+            mean = dataset[c].mean()
+            std = dataset[c].std()
+            p25 = dataset[c].quantile(0.25)
+            median = dataset[c].median()
+            p75 = dataset[c].quantile(0.75)
+
+            outstr += f"\n  - min = {m1}"
+            outstr += f"\n  - max = {m2}"
+            outstr += f"\n  - mean = {mean:.2f}"
+            outstr += f"\n  - std = {std:.2f}"
+            outstr += f"\n  - p25 = {p25}"
+            outstr += f"\n  - p50 = {median}"
+            outstr += f"\n  - p75 = {p75}"
+            outstr += "\n Interqualtile ranges:"
+            outstr += f"\n  - p25-min = {p25 - m1}"
+            outstr += f"\n  - p50-p25 = {median - p25}"
+            outstr += f"\n  - p75-p50 = {p75 - median}"
+            outstr += f"\n  - max-p75 = {m2 - p75}"
+        else:
+            series = dataset[c].value_counts().sort_index()
+            for s in series.index:
+                outstr += f"\n - {s}: {series[s]}"
+
+    outstr += "\n\nPer capability/class data distribution statistics."
+    if Capability.AGE in caps and caps[-1] != Capability.AGE:
+        # Rotating AGE to be the last in the list, as it cannot be the grouping
+        #  variable, since it is numeric.
+        idx = caps.index(Capability.AGE)
+        del caps[idx]
+        caps.append(Capability.AGE)
+        # caps = caps[1:] + caps[:1]
+
+    for c1, c2 in combinations(caps, 2):
+        # Here we test for age, but this should really be a test for any
+        # numerical variable.
+        if c1 == Capability.AGE:
+            continue
+
+        gb = dataset.groupby(c1)[[c2]]
+        tmpdf = None
+        tmpstr = '\nGrouping by "{}", the dataset has the following {} for each value of "{}"'
+
+        if c2 != Capability.AGE:
+            outstr += tmpstr.format(c1, "number of elements", c2)
+            tmpdf = gb.value_counts().sort_index().unstack(level=c1)
+            tmpdf = tmpdf.fillna(0).astype(int)
+            # tmpdf = pd.concat([tmpdf, gb.apply(gini)], axis=1)
+        elif c1 != Capability.AGEGROUP:
+            outstr += tmpstr.format(c1, "statistics", c2)
+            tmpdf = pd.concat(
+                [
+                    gb.min(),
+                    gb.quantile(0.25),
+                    gb.median(),
+                    gb.quantile(0.75),
+                    gb.max(),
+                    gb.apply(gini),
+                ],
+                axis=1,
+            )
+            tmpdf.columns = ["min", "p25", "p50", "p75", "max", "gini_impurity"]
+        else:
+            continue
+
+        outstr += f"\n{tmpdf}\n"
+
+    return outstr
+
+
+if __name__ == "__main__":
+    import os
+
+    logger.info(os.getcwd())
+
+    DATASET_PATH = Path("data/facing2-train/")
+    METADATA_PATH = DATASET_PATH / "meta-w-age.csv"
+
+    meta = pd.read_csv(METADATA_PATH, sep=",", index_col="image")
+    meta[Capability.AGEGROUP.value] = meta.apply(_to_age_bracket, axis=1)
+    meta = meta.sort_index()
+    meta[Capability.AGEGROUP + "_cat"], _ = pd.factorize(
+        meta[Capability.AGEGROUP], sort=True
+    )
+    meta[Capability.SEX + "_cat"], _ = pd.factorize(meta[Capability.SEX], sort=True)
+
+    # GINI IMPURITY
+    # Lower values means a concentration of values around a single class, i.e. bias.
+    age_gini = gini(meta[Capability.AGE])
+
+    agegroup_gini = gini(meta[Capability.AGEGROUP + "_cat"])
+
+    # Should be close to 0.5, indicating a 50/50 split of males and females,
+    # representing maximum uncertainty.
+    sex_gini = gini(meta[Capability.SEX + "_cat"])
+
+    # SHANNON'S ENTROPY
+    count_per_agegroup = meta["age_group"].value_counts()
+    prob_per_agegroup = count_per_agegroup / count_per_agegroup.sum()
+    H_agegroup = entropy(prob_per_agegroup)
+
+    count_per_sex = meta["sex"].value_counts()
+    prob_per_sex = count_per_sex / count_per_sex.sum()
+    H_sex = entropy(prob_per_sex)
+
+    # Now, onto the subgroup metrics.
+    # The goal is to be able to answer the following types of questions:
+    # 1) How many women are in each age-bracket?
+    # 2) Given the population in age-bracket 20-49 years, how is their gender distribution?
+    # 3) Do we need to collect more images of new individuals? If so, what population should we focus on?
+    sex_gb = meta.groupby(Capability.SEX)[["age_group_cat"]]
+    agegroup_gb = meta.groupby(Capability.AGEGROUP)[["sex_cat"]]
+
+    gini_per_sex = sex_gb.apply(gini)
+    gini_per_agegroup = agegroup_gb.apply(gini)
+
+    # Prototype textual description of the dataset.
+    s = generate_description(meta, name=DATASET_PATH.name)
+    print(s)
+
+
+# Diagnostics of biases in the dataset. To be incorporated into a
+# "generate_diagnostics" function later on.
+def generate_bias_diagnostics():
+    pass

src/facebias/metrics.py

@@ -1,231 +1,241 @@
 # -*- coding: utf-8 -*-
 
+"""Model evaluation metrics."""
+
 import logging
+from itertools import combinations
 from pathlib import Path
 from typing import Any
 
 import numpy as np
 import pandas as pd
-from sklearn.metrics import (accuracy_score, cohen_kappa_score,
+from sklearn.metrics import confusion_matrix, multilabel_confusion_matrix
-                             confusion_matrix, max_error, mean_absolute_error,
-                             multilabel_confusion_matrix, balanced_accuracy_score)
 
 from facebias.estimators import Capability
 
 logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(f"facebias:{__name__}")
+logger = logging.getLogger("facebias:metrics.py")
 
 
-def calc_model_performance(
+# Needed for the n_off_accuracy family of functions
-        gt: pd.DataFrame,
+_agegroup_int_map = {
-        preds: pd.DataFrame,
+    "00-02": 0,
-        keys: list[str] = [],
+    "03-09": 1,
-) -> pd.DataFrame:
+    "10-19": 2,
+    "20-29": 3,
+    "30-39": 4,
+    "40-49": 5,
+    "50-59": 6,
+    "60-69": 7,
+    "70+": 8,
+}
+
+
+def n_off_accuracy(gt: list[int], pred: list[int], n: int = 1) -> float:
+    """Returns the n-off accuracy for ordinal class labels encoded as consecutive integers.
+
+    A prediction is counted as correct if it is exact or off by at most `n`.
     """
-    We assume that both `gt` and `preds` have the same structure. They should
+    if len(gt) != len(pred):
-    be indexed by individual ID, such as the image name, and each value is a
+        raise ValueError("Both arrays must have the same length.")
-    dictionary with model prediction capabilities as keys (e.g., "age_group",
-    "sex", "skin-color", etc.), and the values are the predictions, or ground-truth
-    values for each ID/capability.
 
-    if `keys` is empty, then we infer from common keys present in `preds` and `gt`.
+    if len(gt) == 0 or len(pred) == 0:
+        raise ValueError("The arrays must be populated.")
+
+    arr_gt = np.asarray(gt)
+    arr_pred = np.asarray(pred)
+
+    return float(np.mean(np.abs(arr_gt - arr_pred) <= n))
+
+
+def one_off_accuracy(gt: list[int], pred: list[int]) -> float:
+    """1-off accuracy for ordinal class labels encoded as consecutive integers."""
+    return n_off_accuracy(gt, pred, n=1)
+
+
+def two_off_accuracy(gt: list[int], pred: list[int]) -> float:
+    """2-off accuracy for ordinal class labels encoded as consecutive integers."""
+    return n_off_accuracy(gt, pred, n=2)
+
+
+def binary_fpr_fnr(cm: np.ndarray) -> dict[str, np.number]:
+    """Given a confusion matrix, calculates the false-positive and negative rates.
 
     Parameters
     ----------
-    gt: dict[str, dict[str, Any]]
+    cm: np.ndarray
-    preds: dict[str, dict[str, Any]]
+        The 2x2 confusion matrix returned by `scikit-learn`.
-    keys: list[str] | None
 
     Returns
     -------
-    metrics: dict[str, dict[str, float]]
+    metrics: dict[str, np.number]
-    """
+        The calculated metrics keyed under "FPR", "FNR", "TP", "TN", "FP", "FN".
-    common_caps = keys
-    if not keys:
-        common_caps = set(gt.columns) & set(preds.columns)
-        if not common_caps:
-            logger.error(
-                f'No common capabilities found. Predictions has "{preds.columns}",'
-                f' ground-truth has "{gt.columns}".'
-            )
-            return None
 
-    # Finding common images between predictions and ground-truth.
+    See Also
-    common_inds = set(preds.index) & set(gt.index)
+    --------
-    if not common_inds:
+    sklearn.metrics.multilabel_confusion_matrix
-        logger.error("No common images found between predictions and ground-truth.")
+    sklearn.metrics.confusion_matrix
+    """
+    tn, fp, fn, tp = cm.ravel()
+    return {
+        "FPR": fp / (fp + tn) if (fp + tn) != 0 else np.nan,
+        "FNR": fn / (fn + tp) if (fn + tp) != 0 else np.nan,
+        "TN": int(tn),
+        "FP": int(fp),
+        "FN": int(fn),
+        "TP": int(tp),
+    }
+
+
+def multiclass_fpr_fnr(
+    gt: pd.Series, preds: pd.Series, labels: list[Any] | None = None
+) -> tuple[dict, list[str]]:
+    """Calculates one-vs-rest false-positive and negative rates for each class.
+
+    Also returns the counts of false-positives, false-negatives, true-positives
+    and true-negatives, i.e., the confusion matrix for each class.
+
+    Parameters
+    ----------
+    gt: pd.Series
+        Indexed ground-truth data.
+
+    preds: pd.Series
+        Indexed predictions.
+
+    labels: list[Any], optional
+        List of labels in the data. If left empty, then we infer from the
+        union of unique elements in `gt` and `preds`.
+
+    Returns
+    -------
+    results: dict[str, dict[str, number]]
+        The false-{positive,negative}-rates, and true/false positive/negative
+        values for each label in `labels`.
+
+    labels: list[str]
+        The labels used when calculating the confusion matrix. If labels was
+        passed as an argument, then it is returned unchanged, else, we return
+        the labels inferred from the data.
+    """
+    if labels is None or not labels:
+        logging.info("Labels not provided. Inferring from the data.")
+        labels = sorted(list(set(gt.unique()) | set(preds.unique())))
+
+    mcm = multilabel_confusion_matrix(gt, preds, labels=labels)
+
+    results = {}
+    for label, m in zip(labels, mcm):
+        results[label] = binary_fpr_fnr(m)
+
+    return results, labels
+
+
+def _agreement_sanity_checks(x_pred: pd.Series, y_pred: pd.Series) -> bool:
+    if len(x_pred) != len(y_pred):
+        raise ValueError(
+            f"Predictions have different lengths. len(x_pred) = {len(x_pred)}"
+            f" len(y_pred) = {len(y_pred)}"
+        )
+    if not all(x_pred.index == y_pred.index):
+        raise ValueError("Index mismatch between series")
+
+    return True
+
+
+def agreement_fraction(x_pred: pd.Series, y_pred: pd.Series) -> float:
+    """Calculates the fraction of agreement between predictions by two models.
+
+    Note that the predictions must both have the same indices and lengths.
+
+    Parameters
+    ----------
+    x_pred: pd.Series
+        Predictions of the first model.
+
+    y_pred: pd.Series
+        Second model predictions.
+
+    Returns
+    -------
+    fraction: float
+        The fraction of agreement between the results.
+    """
+    try:
+        _agreement_sanity_checks(x_pred, y_pred)
+    except ValueError as e:
+        logger.error(f"Cannot calculate agreement fraction -- {str(e)}.")
         return None
 
-    metric_vals = dict()
+    return (x_pred == y_pred).sum() / len(x_pred)
-    for cap in common_caps:
-        if isinstance(preds[cap].iloc[0], (float, int)):
-            metric_vals[cap] = {
-                "mean_absolute_error": mean_absolute_error(gt[cap], preds[cap]),
-                "max_error": max_error(gt[cap], preds[cap]),
-            }
-        else:
-            labels = sorted(list(set(preds[cap].unique()) | set(gt[cap].unique())))
-            metric_vals[cap] = {
-                "accuracy": accuracy_score(gt[cap], preds[cap]),
-                "balanced_accuracy": balanced_accuracy_score(gt[cap], preds[cap]),
-                "cohen-kappa": cohen_kappa_score(gt[cap], preds[cap], labels=labels),
-            }
-
-    return pd.DataFrame.from_dict(metric_vals)
 
 
-def _find_unique_values_per_capability(
+def agreement_elements(
-    class_output: dict[str, dict[Capability, Any]], caps: list[Capability] | None = None
+    x_pred: pd.Series, y_pred: pd.Series, return_disagreement: bool = True
-) -> dict[Capability, str]:
+) -> tuple[pd.Series, pd.Series | None]:
-    """Returns the set of values per capability in `class_output`.
+    """Returns the elements of agreement, and optionally, disagreement between models.
+
+    Note that, as in `agreement_fraction`, the predictions must have the same
+    lengths and matching indices.
 
     Parameters
     ----------
-    class_output: dict[str, dict[Capability, Any]]
+    x_pred: pd.Series
-        The classification results, or ground-truth data indexed by element.
+        Predictions of the first model.
 
-    caps: list[Capability] | None
+    y_pred: pd.Series
-        The list of capabilities to find unique values for. If left as `None`,
+        Second model predictions.
-        we will find unique values for all of them.
 
-    Results
+    return_disagreement: bool, optional
-    -------
+        Returns the disagreement as well if set (default behavior), otherwise,
-    unique_vals: dict[Capability, str]
+        returns `None` for the `disagreement_idx`.
-        The unique values indexed by capability.
-    """
-    if caps is None:
-        caps = list(next(iter(class_output.values())).keys())
-    elif not isinstance(caps, (list, tuple)):
-        caps = [caps]
-
-    unique_vals = dict()
-    for cap in caps:
-        unique_vals[cap] = set()
-        for res in class_output.values():
-            unique_vals[cap].add(res[cap])
-
-    return unique_vals
-
-
-def _get_capability_data(
-    class_outputs: dict[str, dict[Capability, Any]], cap: Capability
-) -> dict[str, Any]:
-    """Returns data for all individuals regarding a capability.
-
-    Parameters
-    ----------
-    class_outputs: dict[str, dict[Capability, Any]]
-        The estimator outputs indexed by individual.
-
-    cap: Capability
-        The desired capability.
 
     Returns
     -------
-    data: dict[str, Any]
+    agreement_idx: pd.Series
-        The capability data indexed by individual.
+        The index of elements where models have the same results.
+
+    disagreement_idx: pd.Series | None
+        The index of elements where models diverge on their results. Only
+        returned when `return_disagreement` is True, else it returns `None.
     """
-    data_per_id = dict()
+    try:
+        _agreement_sanity_checks(x_pred, y_pred)
+    except ValueError as e:
+        logger.error(f"Cannot get the (dis)agreement elements -- {str(e)}.")
+        return None
 
-    for ind, data in class_outputs.items():
+    idx = x_pred.index
-        if cap not in data:
+    if return_disagreement:
-            logger.warning(
+        return idx[x_pred == y_pred], idx[x_pred != y_pred]
-                f'Entry for capability "{cap.value}" not found for individual "{ind}". Skipping.'
+    else:
-            )
+        return idx[x_pred == y_pred], None
-            continue
-        data_per_id[ind] = data[cap]
-
-    return data_per_id
 
 
-def _filter_by_index(data: dict[str, Any], indx: Any):
+def gini(x: list[float]) -> float:
-    return dict((k, v) for k, v in data.items() if k in indx)
+    x = np.array(x, dtype=np.float32)
-
+    n = len(x)
-
+    diffs = sum(abs(i - j) for i, j in combinations(x, r=2))
-def calc_metrics_per_subgroup(
+    return (diffs / (n**2 * x.mean())).item()
-    gt: dict[str, dict[str, Any]], preds: dict[str, dict[str, Any]]
-) -> dict[Capability, dict[Any, dict]]:
-    """Calculate performance metrics per sub-group for each capability.
-
-    Parameters
-    ----------
-    gt: dict[str, dict[str, Any]]
-
-    preds: dict[str, dict[str, Any]]
-
-    Returns
-    -------
-    metrics: dict[Capability, dict[Any, dict]]
-    """
-    common_caps = set(_find_common_capabilities(gt, preds))
-
-    metrics = {}
-    for cap in common_caps:
-        if cap == Capability.AGE:
-            continue
-
-        other_caps = common_caps - set([cap])
-        unique_values_cap = _find_unique_values_per_capability(gt, cap)[cap]
-
-        metrics[cap] = {}
-        for val in unique_values_cap:
-            ids = [k for k, v in gt.items() if v[cap] == val]
-
-            metrics[cap][val] = {"number_of_elements": len(ids)}
-            for ocap in other_caps:
-                metrics[cap][val][ocap] = {}
-                filtered_pred = _filter_by_index(_get_capability_data(preds, ocap), ids)
-                filtered_gt = _filter_by_index(_get_capability_data(gt, ocap), ids)
-
-                filtered_pred_data = np.array([filtered_pred[i] for i in ids])
-                filtered_gt_data = np.array([filtered_gt[i] for i in ids])
-
-                if isinstance(filtered_pred_data[0], float):
-                    # If data is numeric, we calculate regression-based metrics
-                    metrics[cap][val][ocap] = {
-                        "mean_absolute_error": mean_absolute_error(
-                            filtered_gt_data, filtered_pred_data
-                        ),
-                        "max_error": max_error(filtered_gt_data, filtered_pred_data),
-                    }
-                else:
-                    if len(np.unique(filtered_gt_data)) == 2:
-                        print(cap, val, ocap)
-                        cm = confusion_matrix(filtered_gt_data, filtered_pred_data, labels=get_unique_labels(filtered_gt_data, ocap))
-                    else:
-                        print(cap, val, ocap)
-                        cm = multilabel_confusion_matrix(filtered_gt_data, filtered_pred_data, labels=get_unique_labels(filtered_gt, ocap))
-                    return cm
-
-                    metrics[cap][val][ocap] = {
-                        "accuracy": accuracy_score(
-                            filtered_gt_data, filtered_pred_data
-                        ),
-                    }
-
-    return metrics
-
-
-def get_unique_labels(data, cap: Capability):
-    caps = set()
-    for id, vals in data.items():
-        if cap not in vals:
-            logger.warning(f'Capability "{cap.value}" not registered for individual "{id}". Skipping')
-            continue
-
-        caps.add(vals[cap])
-
-    return caps
 
 
+# TODO(gschardong): Move to the same file as `load_dataset`
 def _to_age_bracket(row):
     iage = int(row["age"])
+    if iage < 3:
+        return "00-02"
+    elif iage < 10:
+        return "03-09"
+    elif iage > 69:
+        return "70+"
+
     d = iage // 10 * 10
     return "{}-{}".format(d, d + 9)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     import os
+
     logger.info(os.getcwd())
 
     from facebias import load_dataset
@@ -237,53 +247,89 @@ if __name__ == '__main__':
     DATASET_PATH = Path("../../data/facing2-train/")
     METADATA_PATH = DATASET_PATH / "meta-w-age.csv"
     DETECTOR_PATH = Path("../../models/blaze_face_full_range.tflite")
-    # TEST_IMS: list[str] = ["10", "12", "14", "9"]
 
     detector = MediapipeDetector(str(DETECTOR_PATH))
     imdict, _ = load_dataset(
         DATASET_PATH, meta_path=None, imname_proc_fn=lambda x: x.split("_")[0]
     )
-    meta = pd.read_csv(METADATA_PATH, sep=',', index_col="image")
+    meta = pd.read_csv(METADATA_PATH, sep=",", index_col="image")
     meta[Capability.AGEGROUP.value] = meta.apply(_to_age_bracket, axis=1)
     meta = meta.sort_index()
 
     face_bboxes = get_face_boxes(imdict, detector)
 
-    # for t in TEST_IMS:
-    #     logger.info("-- {} - {}".format(t, meta[str(t)]))
-
     print(FairFace.capabilities())
-    model_ff = FairFace(Path("../../models/fairface_alldata_4race_20191111.pt"), device="cpu")
-    preds_ff = model_ff.predict(imdict)
-    preds_ff = pd.DataFrame.from_dict(preds_ff).T
-    preds_ff.index.rename("image", inplace=True)
-    preds_ff.index = preds_ff.index.astype(meta.index.dtype)
-    preds_ff = preds_ff.sort_index()
-
-    metrics_ff = calc_model_performance(meta, preds_ff)
-
-    # logger.info("FairFace -- Test Images")
-    # for t in TEST_IMS:
-    #     logger.info("--{} - {}".format(t, preds_ff[str(t)]))
-
-    metrics_ff_groups = calc_metrics_per_subgroup(meta, preds_ff)
-    for k, v in metrics_ff_groups.items():
-        for kv, vv in v.items():
-            print(k, kv, vv)
-
-
     print(MiVOLOv1.capabilities())
-    model_mv = MiVOLOv1(
-        Path("../../models/volo-v1_model_imdb_age_gender_4.22.pth.tar"), device="cpu"
-    )
-    preds_mv = model_mv.predict(imdict)
-    preds_mv = pd.DataFrame.from_dict(preds_mv).T
-    preds_mv.index.rename("image", inplace=True)
-    preds_mv.index = preds_mv.index.astype(meta.index.dtype)
-    preds_mv = preds_mv.sort_index()
 
-    # logger.info("MiVOLOv1(Face Only) -- Test Images")
+    models = {
-    # for t in TEST_IMS:
+        "fairface": FairFace(
-    #     logger.info("{} - {}".format(t, preds_mv[str(t)]))
+            Path("../../models/fairface_alldata_4race_20191111.pt"), device="cpu"
+        ),
+        "mivolo": MiVOLOv1(
+            Path("../../models/volo-v1_model_imdb_age_gender_4.22.pth.tar"),
+            device="cpu",
+        ),
+    }
 
-    metrics_mv = calc_model_performance(meta, preds_mv)
+    preds_per_model = dict((k, None) for k in models.keys())
+    for model_name, model in models.items():
+        preds = model.predict(imdict)
+        preds = pd.DataFrame.from_dict(preds).T
+        preds.index.rename("image", inplace=True)
+        preds.index = preds.index.astype(meta.index.dtype)
+        preds = preds.sort_index()
+        preds_per_model[model_name] = preds
+
+    for model_name, preds in preds_per_model.items():
+        gt_age_group_ord = meta["age_group"].apply(lambda x: _agegroup_int_map[x])
+        preds_age_group_ord = preds["age_group"].apply(lambda x: _agegroup_int_map[x])
+
+        acc_one_off = one_off_accuracy(gt_age_group_ord, preds_age_group_ord)
+        acc_two_off = two_off_accuracy(gt_age_group_ord, preds_age_group_ord)
+
+        agegroup_subclass, labels = multiclass_fpr_fnr(
+            meta["age_group"], preds["age_group"]
+        )
+
+        print("==== Age group metrics by class ====")
+        # Print as table.
+        for k, v in agegroup_subclass.items():
+            print(f"Class {k}")
+            for m, vv in v.items():
+                print(f"\t{m} -- {vv}")
+
+        agegroup_subclass = pd.DataFrame.from_dict(agegroup_subclass)
+        print(agegroup_subclass)
+
+        model_cls = type(model)
+        ordered_labels = model_cls.possible_capability_values(Capability.SEX)
+        metrics_sex = binary_fpr_fnr(
+            confusion_matrix(meta["sex"], preds["sex"], labels=ordered_labels)
+        )
+
+        print(
+            "==== Sex metrics ===="
+            f"\nPositive class -- {ordered_labels[0]},"
+            f" Negative class -- {ordered_labels[1]}"
+        )
+        for k, v in metrics_sex.items():
+            print(f"\t{k} -- {v}")
+
+    # Agreement tests
+    model_list = list(models.keys())
+    for i in range(len(model_list)):
+        for j in range(i + 1, len(model_list)):
+            first, second = model_list[i], model_list[j]
+            print(f"{first} -- {second}")
+
+            for cap in model_cls.capabilities():
+                if cap == Capability.AGE:
+                    continue
+                frac = agreement_fraction(
+                    preds_per_model[first][cap], preds_per_model[second][cap]
+                )
+                print(f'Agreement fraction for capability: "{cap}" - {frac}')
+                agreement, disagreement = agreement_elements(
+                    meta[cap], preds[cap], return_disagreement=True
+                )
+                print(disagreement)