"""Metrics for ultrasound images."""
from functools import partial
from typing import List
import keras
import numpy as np
from keras import ops
from zea import log
from zea.backend import func_on_device, jit
from zea.func import tensor
from zea.func.tensor import translate
from zea.internal.registry import metrics_registry
from zea.internal.utils import reduce_to_signature
from zea.models.lpips import LPIPS
from zea.models.ultrapips import UltraPIPS
[docs]
def get_metric(name, **kwargs):
"""Get metric function given name."""
metric_fn = metrics_registry[name]
if not metric_fn.__name__.startswith("get_"):
return partial(metric_fn, **kwargs)
log.info(f"Initializing metric: {log.green(name)}")
return metric_fn(**kwargs)
def _reduce_mean(array, keep_batch_dim=True):
"""Reduce array by taking the mean.
Args:
array (tensor): Input tensor of shape (..., height, width, channels)
keep_batch_dim (bool): Whether to keep the batch dimensions when reducing.
Default is True.
"""
if keep_batch_dim:
return ops.mean(array, axis=(-3, -2, -1))
else:
return ops.mean(array)
[docs]
@metrics_registry(name="cnr", paired=True, jittable=True)
def cnr(x, y):
"""Calculate contrast to noise ratio"""
mu_x = ops.mean(x)
mu_y = ops.mean(y)
var_x = ops.var(x)
var_y = ops.var(y)
return 20 * ops.log10(ops.abs(mu_x - mu_y) / ops.sqrt((var_x + var_y) / 2))
[docs]
@metrics_registry(name="contrast", paired=True, jittable=True)
def contrast(x, y):
"""Contrast ratio"""
return 20 * ops.log10(ops.mean(x) / ops.mean(y))
[docs]
@metrics_registry(name="gcnr", paired=True, jittable=False)
def gcnr(x, y, bins=256):
"""Generalized contrast-to-noise-ratio"""
x = ops.convert_to_numpy(x)
y = ops.convert_to_numpy(y)
x = np.ravel(x)
y = np.ravel(y)
_, bins = np.histogram(np.concatenate((x, y)), bins=bins)
f, _ = np.histogram(x, bins=bins, density=True)
g, _ = np.histogram(y, bins=bins, density=True)
f /= np.sum(f)
g /= np.sum(g)
return 1 - np.sum(np.minimum(f, g))
[docs]
@metrics_registry(name="fwhm", paired=False, jittable=False)
def fwhm(img):
"""Resolution full width half maxima"""
mask = ops.nonzero(img >= 0.5 * ops.amax(img))[0]
return mask[-1] - mask[0]
[docs]
@metrics_registry(name="snr", paired=False, jittable=True)
def snr(img):
"""Signal to noise ratio"""
return ops.mean(img) / ops.std(img)
[docs]
@metrics_registry(name="wopt_mae", paired=True, jittable=True)
def wopt_mae(ref, img):
"""Find the optimal weight that minimizes the mean absolute error"""
wopt = ops.median(ref / img)
return wopt
[docs]
@metrics_registry(name="wopt_mse", paired=True, jittable=True)
def wopt_mse(ref, img):
"""Find the optimal weight that minimizes the mean squared error"""
wopt = ops.sum(ref * img) / ops.sum(img * img)
return wopt
[docs]
@metrics_registry(name="psnr", paired=True, jittable=True)
def psnr(y_true, y_pred, *, max_val=255):
"""Peak Signal to Noise Ratio (PSNR) for two input tensors.
PSNR = 20 * log10(max_val) - 10 * log10(mean(square(y_true - y_pred)))
Args:
y_true (tensor): input tensor of shape (height, width, channels)
with optional batch dimension.
y_pred (tensor): input tensor of shape (height, width, channels)
with optional batch dimension.
max_val: The dynamic range of the images
Returns:
Tensor (float): PSNR score for each image in the batch.
"""
mse = _reduce_mean(ops.square(y_true - y_pred))
psnr = 20 * ops.log10(max_val) - 10 * ops.log10(mse)
return psnr
[docs]
@metrics_registry(name="mse", paired=True, jittable=True)
def mse(y_true, y_pred):
"""Gives the MSE for two input tensors.
Args:
y_true (tensor): input tensor of shape (height, width, channels)
with optional batch dimension.
y_pred (tensor): input tensor of shape (height, width, channels)
with optional batch dimension.
Returns:
(float): mean squared error between y_true and y_pred. L2 loss.
"""
return _reduce_mean(ops.square(y_true - y_pred))
[docs]
@metrics_registry(name="mae", paired=True, jittable=True)
def mae(y_true, y_pred):
"""Gives the MAE for two input tensors.
Args:
y_true (tensor): input tensor of shape (height, width, channels)
with optional batch dimension.
y_pred (tensor): input tensor of shape (height, width, channels)
with optional batch dimension.
Returns:
(float): mean absolute error between y_true and y_pred. L1 loss.
"""
return _reduce_mean(ops.abs(y_true - y_pred))
[docs]
@metrics_registry(name="ssim", paired=True, jittable=True)
def ssim(
a,
b,
*,
max_val: float = 255.0,
filter_size: int = 11,
filter_sigma: float = 1.5,
k1: float = 0.01,
k2: float = 0.03,
):
"""Computes the structural similarity index (SSIM) between image pairs.
Thin wrapper around :func:`keras.ops.image.ssim`, which implements the standard
SSIM from: Z. Wang, A. C. Bovik, H. R. Sheikh and E. P. Simoncelli,
"Image quality assessment: from error visibility to structural similarity",
in IEEE Transactions on Image Processing, vol. 13, no. 4, pp. 600-612, 2004.
Note: the true SSIM is only defined on grayscale. This function does not
perform any colorspace transform. If the input is in a color space, then it
will compute the average SSIM.
Args:
a: First image (or batch of images), of shape ``((batch,) height, width, channels)``.
b: Second image (or batch of images), with the same shape as `a`.
max_val: The maximum magnitude that `a` or `b` can have.
filter_size: Window size (>= 1). Image dims must be at least this small.
filter_sigma: The bandwidth of the Gaussian used for filtering (> 0.).
k1: One of the SSIM dampening parameters (> 0.).
k2: One of the SSIM dampening parameters (> 0.).
Returns:
Each image's mean SSIM: a scalar for a single image, or one value per image
in the batch.
"""
return ops.image.ssim(
a,
b,
max_val=max_val,
filter_size=filter_size,
filter_sigma=filter_sigma,
k1=k1,
k2=k2,
data_format="channels_last",
)
[docs]
@metrics_registry(name="smsle", paired=True, jittable=True)
def smsle(y_true, y_pred, *, dynamic_range: float = 60.0):
r"""Signed Mean Squared Logarithmic Error (SMSLE).
Loss for regressing radio-frequency data, where a plain MSE is dominated by the
few high-amplitude samples. SMSLE instead compares the signals after log
compression, which weighs the low-amplitude samples that carry the speckle
texture just as heavily. The positive and negative half-waves are compressed
separately, so the sign of the signal (and with it the phase information needed
for further processing) is preserved.
Writing :math:`L_D(x) = \mathrm{clip}\left(20 \log_{10}
\frac{\mathrm{clip}(x, \epsilon, \max|x|)}{\max|x|}, -D, 0\right)` for log
compression to a dynamic range of :math:`D` dB:
.. math::
\mathrm{SMSLE} = \tfrac{1}{2} \mathrm{MSE}\big(L_D(y), L_D(\hat{y})\big)
+ \tfrac{1}{2} \mathrm{MSE}\big(L_D(-y), L_D(-\hat{y})\big)
Both inputs are normalized by their own maximum absolute value before log
compression, so the metric is invariant to the overall scale of either input.
.. citation:: luijten2020adaptive
Args:
y_true (tensor): Ground truth values.
y_pred (tensor): Predicted values.
dynamic_range (float): Dynamic range :math:`D` in dB used to clip the
log-compressed signals. Defaults to 60.
Returns:
(float): SMSLE loss value.
"""
eps = keras.config.epsilon()
def _log_compress(x, x_max):
return ops.clip(
20 * ops.log10(ops.clip(x / x_max, eps, 1)),
-dynamic_range,
0,
)
y_pred_max = ops.maximum(ops.max(ops.abs(y_pred)), eps)
y_true_max = ops.maximum(ops.max(ops.abs(y_true)), eps)
positive = ops.mean(
ops.square(_log_compress(y_pred, y_pred_max) - _log_compress(y_true, y_true_max))
)
negative = ops.mean(
ops.square(_log_compress(-y_pred, y_pred_max) - _log_compress(-y_true, y_true_max))
)
return 0.5 * positive + 0.5 * negative
[docs]
@metrics_registry(name="ncc", paired=True, jittable=True)
def ncc(x, y):
"""Normalized cross correlation"""
num = ops.sum(x * y)
denom = ops.sqrt(ops.sum(x**2) * ops.sum(y**2))
return num / ops.maximum(denom, keras.config.epsilon())
[docs]
@metrics_registry(
name="lpips", paired=True, jittable=True, torch_vmappable=False, torch_jittable=False
)
def get_lpips(image_range, clip=False):
"""
Get the Learned Perceptual Image Patch Similarity (LPIPS) metric.
Args:
image_range (list): The range of the images. Will be translated to [-1, 1] for LPIPS.
clip (bool): Whether to clip the images to `image_range`.
Returns:
The LPIPS metric function.
"""
# Get the LPIPS model
_lpips = LPIPS.from_preset("lpips")
_lpips.trainable = False
_lpips.disable_checks = True
def lpips(img1, img2):
"""
The LPIPS metric function.
Args:
img1 (tensor) with shape (height, width, channels) with optional batch dimension
img2 (tensor) with shape (height, width, channels) with optional batch dimension
Returns (float): The LPIPS metric between img1 and img2 with shape (batch_size,)
or scalar if no batch dimension.
"""
# clip and translate images to [-1, 1]
if clip:
img1 = ops.clip(img1, *image_range)
img2 = ops.clip(img2, *image_range)
img1 = translate(img1, image_range, [-1, 1])
img2 = translate(img2, image_range, [-1, 1])
return _lpips([img1, img2])
return lpips
[docs]
@metrics_registry(
name="ultrapips", paired=True, jittable=True, torch_vmappable=False, torch_jittable=False
)
def get_ultrapips(image_range, clip=False, preset="ultrapips-tusa"):
"""
Get the UltraPIPS perceptual similarity metric for B-mode ultrasound.
.. note::
Unlike LPIPS, UltraPIPS resizes all images to a fixed square size, so
non-square images are stretched (a warning is logged). Crop or pad them
to a square first if the aspect ratio matters.
Args:
image_range (list): The range of the images. Will be translated to [-1, 1] for UltraPIPS.
clip (bool): Whether to clip the images to `image_range`.
preset (str): The UltraPIPS preset (backbone) to load.
Returns:
The UltraPIPS metric function.
"""
_ultrapips = UltraPIPS.from_preset(preset)
_ultrapips.trainable = False
_ultrapips.disable_checks = True
def ultrapips(img1, img2):
"""
The UltraPIPS metric function.
Args:
img1 (tensor) with shape (height, width, channels) with optional batch dimension
img2 (tensor) with shape (height, width, channels) with optional batch dimension
Returns (float): The UltraPIPS distance between img1 and img2 with shape (batch_size,)
or scalar if no batch dimension.
"""
# cast to float first so translate does not subtract in an integer dtype
img1 = ops.cast(img1, "float32")
img2 = ops.cast(img2, "float32")
if clip:
img1 = ops.clip(img1, *image_range)
img2 = ops.clip(img2, *image_range)
img1 = translate(img1, image_range, [-1, 1])
img2 = translate(img2, image_range, [-1, 1])
return _ultrapips([img1, img2])
return ultrapips
[docs]
class Metrics:
"""Class for calculating multiple paired metrics. Also useful for batch processing.
Will preprocess images by translating to [0, 255], clipping, and quantizing to uint8
if specified.
Example:
.. doctest::
>>> from zea import metrics
>>> import numpy as np
>>> metrics = metrics.Metrics(["psnr", "lpips"], image_range=[0, 255])
>>> y_true = np.random.rand(4, 128, 128, 1)
>>> y_pred = np.random.rand(4, 128, 128, 1)
>>> result = metrics(y_true, y_pred)
>>> result = {k: float(v) for k, v in result.items()}
>>> print(result) # doctest: +ELLIPSIS
{'psnr': ..., 'lpips': ...}
"""
def __init__(
self,
metrics: List[str],
image_range: tuple,
quantize: bool = False,
clip: bool = False,
jit_compile: bool = True,
**kwargs,
):
"""Initialize the Metrics class.
Args:
metrics (list): List of metric names to calculate.
image_range (tuple): The range of the images. Used for metrics like PSNR and LPIPS.
quantize (bool): Whether to quantize the images to uint8 before calculating metrics.
clip (bool): Whether to clip the images to `image_range` before calculating metrics.
kwargs: Additional keyword arguments to pass to the metric functions.
"""
# Assert all metrics are paired
for m in metrics:
assert metrics_registry.get_parameter(m, "paired"), (
f"Metric {m} is not a paired metric."
)
# Add image_range to kwargs for metrics that require it
kwargs["image_range"] = image_range
self.image_range = image_range
# Initialize all metrics
self.metrics = {}
self.torch_vmappable = {}
for m in metrics:
jittable = metrics_registry.get_parameter(m, "jittable")
metric_fn = get_metric(m, **reduce_to_signature(metrics_registry[m], kwargs))
if jit_compile and jittable:
metric_fn = jit(
metric_fn,
torch=metrics_registry.get_parameter(m, "torch_jittable", default=True),
)
self.metrics[m] = metric_fn
self.torch_vmappable[m] = metrics_registry.get_parameter(
m, "torch_vmappable", default=True
)
# Other settings
self.quantize = quantize
self.clip = clip
@staticmethod
def _call_metric_fn(
fun,
y_true,
y_pred,
average_batches,
return_numpy,
device,
mapped_batch_size=None,
torch_vmappable=True,
):
num_batch_axes = max(0, ops.ndim(y_true) - 3)
# Because most metric functions do not support batching, we vmap over the batch axes.
# This does assume that the metric function can handle single images of shape (h, w, c).
# Some metrics (e.g. lpips) are not compatible with torch.vmap (e.g. due to Conv2D
# channels_last memory format checks), so we fall back to a sequential loop.
metric_fn = fun
for _ in range(num_batch_axes):
# recursively vmap the leading axis
# disable_jit only when on the torch backend and the metric is not
# torch-vmappable (e.g. LPIPS), so that other backends (JAX, etc.)
# still benefit from their native vmap / vectorized_map paths.
disable_jit = not torch_vmappable and keras.backend.backend() == "torch"
metric_fn = tensor.vmap(
metric_fn,
in_axes=0,
disable_jit=disable_jit,
batch_size=mapped_batch_size,
)
out = func_on_device(metric_fn, device, y_true, y_pred)
if average_batches:
out = ops.mean(out)
if return_numpy:
out = ops.convert_to_numpy(out)
return out
def _preprocess(self, tensor):
tensor = translate(tensor, self.image_range, [0, 255])
if self.clip:
tensor = ops.clip(tensor, 0, 255)
if self.quantize:
tensor = ops.cast(tensor, "uint8")
tensor = ops.cast(tensor, "float32") # Some metrics require float32
return tensor
[docs]
def __call__(
self,
y_true,
y_pred,
average_batches=True,
mapped_batch_size=None,
return_numpy=True,
device=None,
):
"""Calculate all metrics and return as a dictionary.
Assumes input shape (..., h, w, c), i.e. images of shape (h, w, c) with
any number of leading batch dimensions. The metrics will be calculated
on these 2d images and mapped across all leading batch dimensions.
Args:
y_true (tensor): Ground truth images with shape (..., h, w, c)
y_pred (tensor): Predicted images with shape (..., h, w, c)
average_batches (bool): Whether to average the metrics over the batch dimensions.
mapped_batch_size (optional int): The batch size to use for computing
metric values in parallel.
You may want to decrease this if you run into memory issues, e.g. with LPIPS.
return_numpy (bool): Whether to return the metrics as numpy arrays. If False, will
return as tensors.
device (str): The device to run the metric calculations on. If None, will use the
default device.
"""
results = {}
for name, metric in self.metrics.items():
results[name] = self._call_metric_fn(
metric,
self._preprocess(y_true),
self._preprocess(y_pred),
average_batches,
return_numpy,
device,
mapped_batch_size=mapped_batch_size,
torch_vmappable=self.torch_vmappable[name],
)
return results
def _sector_reweight_image(image, sector_angle, axis):
"""
Reweights image according to the amount of area each
row of pixels will occupy if that image is scan converted
with angle sector_angle.
This 'image' could be e.g. a pixelwise loss or metric.
We can compute this by viewing the scan converted image as the sector
of a circle with a known central angle, and radius given by depth.
See: https://en.wikipedia.org/wiki/Circular_sector
Params:
image (ndarray or Tensor): image to be re-weighted, any shape
sector_angle (float | int): angle in degrees
axis (int): axis corresponding to the height/depth dimension.
Returns:
reweighted_image (ndarray): image with pixels reweighted to area occupied by each
pixel post-scan-conversion.
"""
height = image.shape[axis]
depths = ops.arange(height, dtype="float32") + 0.5 # center of the pixel as its depth
reweighting_factors = (sector_angle / 360) * 2 * np.pi * depths
# Reshape reweighting_factors to broadcast along the specified axis
shape = [1] * ops.ndim(image)
shape[axis] = height
reweighting_factors = ops.reshape(reweighting_factors, shape)
return reweighting_factors * image