Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,009 Bytes
1b34a12 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
"""
Copyright (c) Facebook, Inc. and its affiliates.
This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""
import argparse
import pathlib
from argparse import ArgumentParser
from typing import Optional
import h5py
import numpy as np
from runstats import Statistics
from skimage.metrics import peak_signal_noise_ratio, structural_similarity
from fastmri import transforms
def mse(gt: np.ndarray, pred: np.ndarray) -> np.ndarray:
"""Compute Mean Squared Error (MSE)"""
return np.mean((gt - pred) ** 2)
def nmse(gt: np.ndarray, pred: np.ndarray) -> np.ndarray:
"""Compute Normalized Mean Squared Error (NMSE)"""
return np.array(np.linalg.norm(gt - pred) ** 2 / np.linalg.norm(gt) ** 2)
def psnr(
gt: np.ndarray, pred: np.ndarray, maxval: Optional[float] = None
) -> np.ndarray:
"""Compute Peak Signal to Noise Ratio metric (PSNR)"""
if maxval is None:
maxval = gt.max()
return peak_signal_noise_ratio(gt, pred, data_range=maxval)
def ssim(
gt: np.ndarray, pred: np.ndarray, maxval: Optional[float] = None
) -> np.ndarray:
"""Compute Structural Similarity Index Metric (SSIM)"""
if not gt.ndim == 3:
raise ValueError("Unexpected number of dimensions in ground truth.")
if not gt.ndim == pred.ndim:
raise ValueError("Ground truth dimensions does not match pred.")
maxval = gt.max() if maxval is None else maxval
ssim = np.array([0])
for slice_num in range(gt.shape[0]):
ssim = ssim + structural_similarity(
gt[slice_num], pred[slice_num], data_range=maxval
)
return ssim / gt.shape[0]
METRIC_FUNCS = dict(
MSE=mse,
NMSE=nmse,
PSNR=psnr,
SSIM=ssim,
)
class Metrics:
"""
Maintains running statistics for a given collection of metrics.
"""
def __init__(self, metric_funcs):
"""
Parameters
----------
metric_funcs : dict
A dictionary where the keys are metric names (as strings) and the values
are Python functions for evaluating the corresponding metrics.
"""
self.metrics = {metric: Statistics() for metric in metric_funcs}
def push(self, target, recons):
for metric, func in METRIC_FUNCS.items():
self.metrics[metric].push(func(target, recons))
def means(self):
return {metric: stat.mean() for metric, stat in self.metrics.items()}
def stddevs(self):
return {metric: stat.stddev() for metric, stat in self.metrics.items()}
def __repr__(self):
means = self.means()
stddevs = self.stddevs()
metric_names = sorted(list(means))
return " ".join(
f"{name} = {means[name]:.4g} +/- {2 * stddevs[name]:.4g}"
for name in metric_names
)
def evaluate(args, recons_key):
metrics = Metrics(METRIC_FUNCS)
for tgt_file in args.target_path.iterdir():
with h5py.File(tgt_file, "r") as target, h5py.File(
args.predictions_path / tgt_file.name, "r"
) as recons:
if args.acquisition and args.acquisition != target.attrs["acquisition"]:
continue
if args.acceleration and target.attrs["acceleration"] != args.acceleration:
continue
target = target[recons_key][()]
recons = recons["reconstruction"][()]
target = transforms.center_crop(
target, (target.shape[-1], target.shape[-1])
)
recons = transforms.center_crop(
recons, (target.shape[-1], target.shape[-1])
)
metrics.push(target, recons)
return metrics
if __name__ == "__main__":
parser = ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
"--target-path",
type=pathlib.Path,
required=True,
help="Path to the ground truth data",
)
parser.add_argument(
"--predictions-path",
type=pathlib.Path,
required=True,
help="Path to reconstructions",
)
parser.add_argument(
"--challenge",
choices=["singlecoil", "multicoil"],
required=True,
help="Which challenge",
)
parser.add_argument("--acceleration", type=int, default=None)
parser.add_argument(
"--acquisition",
choices=[
"CORPD_FBK",
"CORPDFS_FBK",
"AXT1",
"AXT1PRE",
"AXT1POST",
"AXT2",
"AXFLAIR",
],
default=None,
help=(
"If set, only volumes of the specified acquisition type are used "
"for evaluation. By default, all volumes are included."
),
)
args = parser.parse_args()
recons_key = (
"reconstruction_rss" if args.challenge == "multicoil" else "reconstruction_esc"
)
metrics = evaluate(args, recons_key)
print(metrics)
|