prpy 0.2.2__py3-none-any.whl

prpy/numpy/face.py

@@ -0,0 +1,213 @@
+# Copyright (c) 2024 Philipp Rouast
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import numpy as np
+from prpy.numpy.image import crop_slice_resize
+from typing import Tuple, Union
+
+def _get_roi_from_det(
+    det: tuple,
+    rel_change: tuple,
+    clip_dims: Union[tuple, None] = None
+  ) -> Tuple[int, int, int, int]:
+  """Convert face detection to roi by relative add/reduce.
+
+  Args:
+    det: The face detection [0, H/W] in format (x0, y0, x1, y1).
+    rel_change: The relative change to make in format (left, top, right, bottom).
+    clip_dims: tuple (frame_w, frame_h) to clip the result to (optional).
+  Returns:
+    out: The roi [0, H/W] in format (x0, y0, x1, y1)
+  """
+  assert isinstance(det, tuple) and len(det) == 4 and all(isinstance(i, int) for i in det)
+  assert det[2] > det[0]
+  assert det[3] > det[1]
+  assert isinstance(rel_change, tuple) and len(rel_change) == 4 and all(isinstance(i, float) for i in rel_change)
+  assert clip_dims is None or (isinstance(clip_dims, tuple) and len(clip_dims) == 2 and all(isinstance(i, int) for i in clip_dims))
+  def _clip_dims(val, min_dim, max_dim):
+    return min(max(val, min_dim), max_dim)
+  det_w = det[2]-det[0]
+  det_h = det[3]-det[1]
+  rel_ch_l, rel_ch_t, rel_ch_r, rel_ch_b = rel_change
+  abs_ch_l = int(rel_ch_l * det_w)
+  abs_ch_t = int(rel_ch_t * det_h)
+  abs_ch_r = int(rel_ch_r * det_w)
+  abs_ch_b = int(rel_ch_b * det_h)
+  if clip_dims is not None:
+    return (_clip_dims(det[0] - abs_ch_l, 0, clip_dims[0]),
+            _clip_dims(det[1] - abs_ch_t, 0, clip_dims[1]),
+            _clip_dims(det[2] + abs_ch_r, 0, clip_dims[0]),
+            _clip_dims(det[3] + abs_ch_b, 0, clip_dims[1]))
+  else:
+    return (det[0]-abs_ch_l, det[1]-abs_ch_t, det[2]+abs_ch_r, det[3]+abs_ch_b)
+
+def get_face_roi_from_det(det: tuple) -> tuple:
+  """Convert face detection into face roi.
+  Reduces width to 60% and height to 80%. 
+  
+  Args:
+    det: The face detection [0, H/W] in form (x0, y0, x1, y1)
+  Returns:
+    out: The roi [0, H/W] in form (x0, y0, x1, y1)
+  """
+  return _get_roi_from_det(det=det, rel_change=(-0.2, -0.1, -0.2, -0.1))
+
+def get_forehead_roi_from_det(det: tuple) -> tuple:
+  """Convert face detection into forehead roi.
+  Reduces det to forehead as 35% to 65% of width, and 15% to 25% of height. 
+  
+  Args:
+    det: The face detection [0, H/W] in form (x0, y0, x1, y1)
+  Returns:
+    out: The roi [0, H/W] in form (x0, y0, x1, y1)
+  """
+  return _get_roi_from_det(det=det, rel_change=(-0.35, -0.15, -0.35, -0.75))
+
+def get_upper_body_roi_from_det(
+    det: tuple,
+    clip_dims: tuple,
+    cropped: bool = False,
+    v: int = 1
+  ) -> tuple:
+  """Convert face detection into upper body roi and clip to frame constraints.
+
+  Args:
+    det: The face detection [0, H/W] in form (x0, y0, x1, y1)
+    clip_dims: constraints (frame_w, frame_h) to clip the result to
+    cropped: Create cropped variant?
+    v: Version of ROI definition (0, 1, 2, or 3)
+  Returns:
+    out: The roi [0, H/W] in form (x0, y0, x1, y1)
+  """
+  assert isinstance(cropped, bool)
+  assert isinstance(v, int)
+  if v == 0:
+    # V0: (.25, .3, .25, .5) -> (.175, .27, .175, .45)
+    if not cropped:
+      return _get_roi_from_det(
+        det=det, rel_change=(.25, .3, .25, .5), clip_dims=clip_dims)
+    else:
+      return _get_roi_from_det(
+        det=det, rel_change=(.175, .27, .175, .45), clip_dims=clip_dims)
+  elif v == 1:
+    # V1: (.25, .2, .25, .4) -> (.175, .15, .175, .3)
+    if not cropped:
+      return _get_roi_from_det(
+        det=det, rel_change=(.25, .2, .25, .4), clip_dims=clip_dims)
+    else:
+      return _get_roi_from_det(
+        det=det, rel_change=(.175, .15, .175, .3), clip_dims=clip_dims)
+  elif v == 2:
+    # V2: (.25, .1, .25, .5) -> (.175, .075, .175, .375)
+    if not cropped:
+      return _get_roi_from_det(
+        det=det, rel_change=(.25, .1, .25, .5), clip_dims=clip_dims)
+    else:
+      return _get_roi_from_det(
+        det=det, rel_change=(.175, .075, .175, .375), clip_dims=clip_dims)
+  elif v == 3:
+    # V3: (.2, .3, .2, .45) -> (.15, .25, .15, .35)
+    if not cropped:
+      return _get_roi_from_det(
+        det=det, rel_change=(.2, .3, .2, .45), clip_dims=clip_dims)
+    else:
+      return _get_roi_from_det(
+        det=det, rel_change=(.15, .25, .15, .35), clip_dims=clip_dims)
+  else:
+    raise ValueError("v {} is not defined".format(v))
+
+def get_meta_roi_from_det(
+    det: tuple,
+    clip_dims: tuple
+  ) -> tuple:
+  """Convert face detection into meta roi and clip to frame constraints.
+
+  Args:
+    det: The face detection [0, H/W] in form (x0, y0, x1, y1)
+    clip_dims: constraints (frame_w, frame_h) to clip the result to
+  Returns:
+    out: The roi [0, H/W] in form (x0, y0, x1, y1)
+  """
+  return _get_roi_from_det(
+    det=det, rel_change=(.2, .2, .2, .2), clip_dims=clip_dims)
+
+def get_roi_from_det(
+    det: tuple,
+    roi_method: Union[str, None],
+    clip_dims: Union[tuple, None] = None
+  ) -> tuple:
+  """Convert face detection into specified roi.
+
+  Args:
+    det: The face detection [0, H/W] in form (x0, y0, x1, y1)
+    roi_method: Which roi method to use. Either 'forehead', 'face',
+      'upper_body', 'upper_body_cropped', 'meta', None (directly use det)
+    clip_dims: Constraints (frame_w, frame_h) to clip the result to (optional).
+  Returns:
+    out: The roi [0, H/W] in form (x0, y0, x1, y1)
+  """
+  assert roi_method is None or isinstance(roi_method, str)
+  if roi_method == 'face':
+    return get_face_roi_from_det(det)
+  elif roi_method == 'forehead':
+    return get_forehead_roi_from_det(det)
+  elif roi_method == 'upper_body':
+    assert clip_dims is not None
+    return get_upper_body_roi_from_det(det, clip_dims=clip_dims, cropped=False)
+  elif roi_method == 'upper_body_cropped':
+    assert clip_dims is not None
+    return get_upper_body_roi_from_det(det, clip_dims=clip_dims, cropped=True)
+  elif roi_method == 'meta':
+    assert clip_dims is not None
+    return get_meta_roi_from_det(det, clip_dims=clip_dims)
+  elif roi_method is None or roi_method == 'det':
+    return det
+  else:
+    raise ValueError("roi method {} is not supported".format(roi_method))
+
+def crop_resize_from_det(
+    video: np.ndarray,
+    det: tuple,
+    size: tuple,
+    roi_method: str,
+    library: str,
+    scale_algorithm: str
+  ) -> np.ndarray:
+  """Crop and resize a video according to a single face detection.
+  Resize to specified size with specified method.
+
+  Args:
+    video: The video. Shape (n_frames, h, w, c)
+    det: The face detection in form (x_0, y_0, x_1, y_1)
+    size: The target size for resize - (h, w)
+    roi_method: Which roi method to use. Either 'forehead', 'face',
+      'upper_body', 'upper_body_cropped', 'meta', None (directly use det)
+    library: The library used for resize (PIL, cv2, or tf - returns tf.Tensor)
+    scale_algorithm: The algorithm used for scaling. Supports: bicubic,
+      bilinear, area (not for PIL!), lanczos
+  Returns:
+    result: Cropped and resized video. Shape [n_frames, size[0], size[1], c]
+  """
+  assert isinstance(video, np.ndarray) and len(video.shape) == 4
+  _, height, width, _ = video.shape
+  roi = get_roi_from_det(det, roi_method=roi_method, clip_dims=(width, height))
+  return crop_slice_resize(
+    inputs=video, target_size=size, roi=roi, library=library,
+    preserve_aspect_ratio=False, scale_algorithm=scale_algorithm)

prpy/numpy/image.py

@@ -0,0 +1,141 @@
+# Copyright (c) 2024 Philipp Rouast
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import numpy as np
+from typing import Union
+
+def crop_slice_resize(
+    inputs: np.ndarray,
+    target_size: Union[int, tuple, list], 
+    roi: Union[tuple, list, None] = None,
+    target_idxs: Union[tuple, list, np.ndarray, None] = None,
+    preserve_aspect_ratio: bool = False,
+    keepdims: bool = True,
+    library: str = 'PIL',
+    scale_algorithm: str = 'bicubic'
+  ) -> np.ndarray:
+  """Crop, slice, and resize image(s) with all same settings.
+
+  Args:
+    inputs: The inputs as uint8 shape (h, w, 3) or (n_frames, h, w, 3)
+    target_size: The target size; scalar or (H, W) if preserve_aspect_ratio=False
+    roi: The region of interest in format (x0, y0, x1, y1). Use None to keep all.
+    target_idxs: The frame indices to be used. Use None to keep all.
+    preserve_aspect_ratio: Preserve the aspect ratio?
+    keepdims: If True, always keep n_frames dim. Otherwise, may drop n_frames dim.
+    library: The library to use. `cv2` or `PIL` (return np ndarray), or `tf` (returns tf Tensor)
+    scale_algorithm: The algorithm used for scaling.
+      Supports: bicubic, bilinear, area (not for PIL!), lanczos. Default: bicubic
+  Returns:
+    result: The processed frame(s) as float32 shape (h, w, 3) or (n_frames, h, w, 3)
+  """
+  assert isinstance(inputs, np.ndarray) and (len(inputs.shape) == 3 or len(inputs.shape) == 4)
+  assert isinstance(target_size, int) or (isinstance(target_size, (tuple, list)) and len(target_size) == 2 and all(isinstance(i, int) for i in target_size))
+  assert roi is None or (isinstance(roi, (tuple, list)) and len(roi) == 4 and all(isinstance(i, int) for i in roi) and roi[2] > roi[0] and roi[3] > roi[1])
+  assert target_idxs is None or isinstance(target_idxs, np.ndarray) or (isinstance(target_idxs, (tuple, list)) and all(isinstance(i, int) for i in target_idxs))
+  assert isinstance(preserve_aspect_ratio, bool)
+  assert isinstance(keepdims, bool)
+  assert isinstance(library, str)
+  assert isinstance(scale_algorithm, str)
+  unpack_target_size = lambda x: (x[0], x[1]) if isinstance(x, (list, tuple)) else (x, x)
+  target_height, target_width = unpack_target_size(target_size)
+  inputs_shape = inputs.shape
+  # Add temporal dim if necessary
+  if len(inputs_shape) == 3: inputs = inputs[np.newaxis,:,:,:]
+  # Apply target_idxs and roi
+  inputs = inputs[(target_idxs if target_idxs is not None else slice(None)), 
+                  (slice(roi[1], roi[3]) if isinstance(roi, (tuple, list)) else slice(None)),
+                  (slice(roi[0], roi[2]) if isinstance(roi, (tuple, list)) else slice(None))]
+  in_shape = inputs.shape
+  # Compute out size
+  def _out_size(in_shape, target_height, target_width, preserve_aspect_ratio):
+    _, height, width, _ = in_shape
+    if preserve_aspect_ratio:
+      # Determine critical side
+      inputs_r = float(width)/height
+      target_r = float(target_width)/target_height
+      if inputs_r < target_r:
+        # Height is critical side
+        out_size = (target_height, int(target_height*inputs_r))
+      else:
+        # Height is critical side
+        out_size = (int(target_width/inputs_r), target_width)
+    else:
+      out_size = (target_height, target_width)
+    return out_size
+  out_size = _out_size(in_shape, target_height, target_width, preserve_aspect_ratio)
+  # Distinguish between different cases
+  if out_size == (in_shape[1], in_shape[2]):
+    # No resizing necessary
+    if library == 'tf':
+      import tensorflow as tf
+      out = tf.convert_to_tensor(inputs, dtype=tf.float32)
+    else:
+      out = inputs.astype(np.float32) 
+  else:
+    # Resize to out_size
+    if library == 'tf':
+      import tensorflow as tf
+      # https://www.tensorflow.org/api_docs/python/tf/image/ResizeMethod
+      mapping = {"bicubic": "bicubic", "bilinear": "bilinear", "lanczos": "lanczos3", "area": "area"}
+      try:
+        library_algorithm = mapping[scale_algorithm]
+      except KeyError:
+        raise ValueError("Scaling algorithm {} is not supported by {}".format(scale_algorithm, library))
+      out = tf.image.resize(
+        images=inputs, size=(target_height, target_width),
+        preserve_aspect_ratio=preserve_aspect_ratio,
+        method=library_algorithm, antialias=False)
+    elif library == 'PIL':
+      from PIL import Image
+      # https://pillow.readthedocs.io/en/stable/releasenotes/2.7.0.html#image-resizing-filters
+      mapping = {"bicubic": Image.BICUBIC, "bilinear": Image.BILINEAR, "lanczos": Image.LANCZOS}
+      try:
+        library_algorithm = mapping[scale_algorithm]
+      except KeyError:
+        raise ValueError("Scaling algorithm {} is not supported by {}".format(scale_algorithm, library))
+      # PIL requires (width, height)
+      out_size = (out_size[1], out_size[0])
+      out = np.asarray([
+        np.asarray(Image.fromarray(f).resize(out_size, resample=library_algorithm)) for f in inputs])
+    elif library == 'cv2':
+      import cv2
+      # https://docs.opencv.org/3.4/da/d54/group__imgproc__transform.html
+      mapping = {"bicubic": cv2.INTER_CUBIC, "bilinear": cv2.INTER_LINEAR, "lanczos": cv2.INTER_LANCZOS4, "area": cv2.INTER_AREA}
+      try:
+        library_algorithm = mapping[scale_algorithm]
+      except KeyError:
+        raise ValueError("Scaling algorithm {} is not supported by {}".format(scale_algorithm, library))
+      # cv2 requires (width, height)
+      out_size = (out_size[1], out_size[0])
+      out = np.asarray([
+        cv2.resize(src=f, dsize=out_size, interpolation=library_algorithm) for f in inputs])
+    else:
+      raise ValueError("Library {} not supported".format(library))
+  if keepdims and len(out.shape) == 3 and len(inputs_shape) == 4:
+    # Add temporal dim back - might have been lost when slicing
+    newaxis = tf.newaxis if library == 'tf' else np.newaxis
+    out = out[newaxis,:,:,:]
+  elif not keepdims and len(out.shape) == 4:
+    # Remove temporal dim if necessary
+    squeeze = tf.squeeze if library == 'tf' else np.squeeze
+    if out.shape[0] == 1:
+      out = squeeze(out, axis=0)
+  return out

prpy/numpy/metric.py

@@ -0,0 +1,179 @@
+# Copyright (c) 2024 Philipp Rouast
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import numpy as np
+from scipy import signal
+from typing import Union
+
+def mag2db(mag: Union[np.ndarray, np.float64]) -> np.ndarray:
+  """Magnitude to decibels element-wise.
+
+  Args:
+    mag: Magnitude. Arbitrary shape.
+  Returns:
+    out: Decibels. Same shape as input.
+  """
+  assert isinstance(mag, (np.ndarray, np.float64))
+  return 20. * np.log10(mag)
+
+def mae(
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    axis: Union[int, None] = -1
+  ) -> np.ndarray:
+  """Mean absolute error
+
+  Args:
+    y_true: True values. Shape (..., dim_n, dim_axis)
+    y_pred: Predicted values. Shape (..., dim_n, dim_axis)
+    axis: Axis along which the means are computed
+  Returns:
+    mae: The mean absolute error. Shape (..., dim_n)
+  """
+  assert axis is None or isinstance(axis, int)
+  y_true = np.asarray(y_true)
+  y_pred = np.asarray(y_pred)
+  assert y_true.shape == y_pred.shape
+  if np.all(np.isfinite(y_true)) and np.all(np.isfinite(y_pred)):
+    return np.mean(np.abs(y_true - y_pred), axis=axis)
+  else:
+    return np.full(y_true.shape[:axis], fill_value=np.nan, dtype=y_true.dtype)
+
+def mse(
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    axis: Union[int, None] = -1
+  ) -> np.ndarray:
+  """Mean squared error
+
+  Args:
+    y_true: True values. Shape (..., dim_n, dim_axis)
+    y_pred: Predicted values. Shape (..., dim_n, dim_axis)
+    axis: Axis along which the means are computed
+  Returns:
+    mse: The mean squared error. Shape (..., dim_n)
+  """
+  assert axis is None or isinstance(axis, int)
+  y_true = np.asarray(y_true)
+  y_pred = np.asarray(y_pred)
+  assert y_true.shape == y_pred.shape
+  if np.all(np.isfinite(y_true)) and np.all(np.isfinite(y_pred)):
+    return np.mean(np.square(y_true - y_pred), axis=axis)
+  else:
+    return np.full(y_true.shape[:axis], fill_value=np.nan, dtype=y_true.dtype)
+
+def rmse(
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    axis: Union[int, None] = -1
+  ) -> np.ndarray:
+  """Root mean squared error
+
+  Args:
+    y_true: True values. Shape (..., dim_n, dim_axis)
+    y_pred: Predicted values. Shape (..., dim_n, dim_axis)
+    axis: Axis along which the means are computed
+  Returns:
+    rmse: The root mean squared error. Shape (..., dim_n)
+  """
+  assert axis is None or isinstance(axis, int)
+  y_true = np.asarray(y_true)
+  y_pred = np.asarray(y_pred)
+  assert y_true.shape == y_pred.shape
+  if np.all(np.isfinite(y_true)) and np.all(np.isfinite(y_pred)):
+    return np.sqrt(np.mean(np.square(y_true - y_pred), axis=axis))
+  else:
+    return np.full(y_true.shape[:axis], fill_value=np.nan, dtype=y_true.dtype)
+
+def cor(
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    axis: Union[int, None] = -1
+  ) -> np.ndarray:
+  """Pearson's correlation coefficient
+
+  Args:
+    y_true: True values. Shape (..., dim_n, dim_axis)
+    y_pred: Predicted values. Shape (..., dim_n, dim_axis)
+    axis: Axis along which correlations are computed
+  Returns:
+    cor: The correlation coefficients. Shape (..., dim_n)
+  """
+  assert axis is None or isinstance(axis, int)
+  y_true = np.asarray(y_true)
+  y_pred = np.asarray(y_pred)
+  assert y_true.shape == y_pred.shape
+  if np.all(np.isfinite(y_true)) and np.all(np.isfinite(y_pred)):
+    res_true = y_true - np.mean(y_true, axis=axis, keepdims=True)
+    res_pred = y_pred - np.mean(y_pred, axis=axis, keepdims=True)
+    cov = np.mean(res_true * res_pred, axis=axis)
+    var_true = np.mean(res_true**2, axis=axis)
+    var_pred = np.mean(res_pred**2, axis=axis)
+    sigma_true = np.sqrt(var_true)
+    sigma_pred = np.sqrt(var_pred)
+    return cov / (sigma_true * sigma_pred)
+  else:
+    return np.full(y_true.shape[:axis], fill_value=np.nan, dtype=y_true.dtype)
+
+def snr(
+    f_true: np.ndarray,
+    y_pred: np.ndarray,
+    f_s: float,
+    f_res: float,
+    tol: float = .1,
+    f_min: float = .5,
+    f_max: float = 4.):
+  """Signal-to-noise ratio
+  Args:
+    f_true: The true frequencies. Shape (b,) or ()
+    y_pred: Predicted vals. Shape (b, t) or (t,)
+    f_s: Sampling frequency
+    f_res: Frequency resolution
+    tol: Frequency domain tolerance
+    f_min: Minimum frequency included in metric calculation
+    f_max: Maximum frequency included in metric calculation
+  Returns:
+    snr: The signal to noise ratio. Shape (b,) or ()
+  """
+  assert isinstance(f_s, float)
+  assert isinstance(f_res, float)
+  assert isinstance(tol, float)
+  assert isinstance(f_min, float)
+  assert isinstance(f_max, float)
+  f_true = np.asarray(f_true)
+  y_pred = np.asarray(y_pred)
+  assert len(y_pred.shape) == 1 or len(y_pred.shape) == 2
+  assert f_true.shape == y_pred.shape[:-1]
+  if np.all(np.isfinite(y_pred)) and np.all(np.isfinite(f_true)):
+    n = f_s // f_res
+    f, pxx = signal.periodogram(y_pred, fs=f_s, nfft=n, detrend=False, axis=-1)
+    if len(y_pred.shape) == 2:
+      f = np.broadcast_to(f[np.newaxis], pxx.shape)
+      f_true = f_true[...,np.newaxis]
+    gt_mask_1 = (f >= f_true - tol) & (f <= f_true + tol)
+    gt_mask_2 = (f >= f_true * 2 - tol) & (f <= f_true * 2 + tol)
+    gt_mask = gt_mask_1 | gt_mask_2
+    s_power = np.sum(pxx * gt_mask, axis=-1)
+    f_mask = (f >= f_min) & (f <= f_max)
+    all_power = np.sum(pxx * f_mask, axis=-1)
+    snr = mag2db(s_power / (all_power - s_power))
+    return np.squeeze(snr)
+  else:
+    return np.full(y_pred.shape[:-1], fill_value=np.nan, dtype=y_pred.dtype)