evaluate.py

"""Evaluate trained weights on a processed labeled dataset: timing errors, confusion matrices, F1.

Example:
    python evaluate.py --data processed_data/test --weights weights/encoder.h5
"""
import argparse

import numpy as np
import tensorflow as tf
from sklearn.metrics import multilabel_confusion_matrix, f1_score

from callbacks import asynchrony_flags
from config import SEG_LEN, FS
from data import set_gpu, load_labeled
from networks import resnet9_domain

NAMES = ["DI", "LC", "EC", "IE"]


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('--data', required=True, help='processed labeled data directory')
    p.add_argument('--weights', default='weights/encoder.h5', help='encoder weights (.h5)')
    p.add_argument('--gpu', type=int, default=None)
    args = p.parse_args()

    tf.get_logger().setLevel('ERROR')
    set_gpu(args.gpu)
    with tf.device('/CPU:0'):
        data, labels = load_labeled(args.data, to_dataset=False)
    labels = labels.numpy()
    print(f"{len(labels)} breaths loaded from {args.data}")

    encoder = resnet9_domain()
    encoder.load_weights(args.weights)
    tip_out, tep_out, ie_out, _, _ = encoder.predict(data, batch_size=512)

    # Mean absolute timing errors in ms
    to_ms = SEG_LEN * 1000 / FS
    print("\nMean absolute error [ms]")
    print(f"  mechanical start (aux): {np.mean(np.abs(tip_out[:, 0] - labels[:, 0])) * to_ms:.1f}")
    print(f"  mechanical end   (aux): {np.mean(np.abs(tep_out[:, 0] - labels[:, 1])) * to_ms:.1f}")
    print(f"  patient start inspiration: {np.mean(np.abs(tip_out[:, 1] - labels[:, 2])) * to_ms:.1f}")
    print(f"  patient end inspiration:   {np.mean(np.abs(tep_out[:, 1] - labels[:, 3])) * to_ms:.1f}")

    # Asynchronies: predicted patient timings vs the known mechanical cycle
    di_p, lc_p, ec_p = asynchrony_flags(labels, tip_out[:, 1], tep_out[:, 1])
    di_t, lc_t, ec_t = asynchrony_flags(labels, labels[:, 2], labels[:, 3])
    ie_p = np.round(ie_out[:, 2])
    ie_t = labels[:, -1]

    target = np.stack([di_t, lc_t, ec_t, ie_t], axis=1).astype(np.float32)
    preds = np.stack([di_p, lc_p, ec_p, ie_p], axis=1).astype(np.float32)

    print("\nConfusion matrices [[TN FP] [FN TP]]")
    for name, cm in zip(NAMES, multilabel_confusion_matrix(target, preds)):
        print(f"{name}:\n{cm}")

    print("\nF1 per asynchrony")
    for name, f1 in zip(NAMES, f1_score(target, preds, average=None, zero_division=0)):
        print(f"  {name}: {f1:.4f}")
    print(f"  Average (macro): {f1_score(target, preds, average='macro', zero_division=0):.4f}")


if __name__ == '__main__':
    main()