zea 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

zea/metrics.py

@@ -1,131 +1,431 @@
-"""Quality metrics for ultrasound images."""
+"""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, tensor_ops
+from zea.backend import func_on_device
 from zea.internal.registry import metrics_registry
+from zea.internal.utils import reduce_to_signature
+from zea.models.lpips import LPIPS
+from zea.tensor_ops import translate
 
 
-def get_metric(name):
+def get_metric(name, **kwargs):
     """Get metric function given name."""
-    return metrics_registry[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)
 
 
-@metrics_registry(name="cnr", framework="numpy", supervised=True)
+def _reduce_mean(array, keep_batch_dim=True):
+    """Reduce array by taking the mean.
+    Preserves batch dimension if keep_batch_dim=True.
+    """
+    if keep_batch_dim:
+        ndim = ops.ndim(array)
+        axis = tuple(range(max(0, ndim - 3), ndim))
+    else:
+        axis = None
+    return ops.mean(array, axis=axis)
+
+
+@metrics_registry(name="cnr", paired=True)
 def cnr(x, y):
     """Calculate contrast to noise ratio"""
-    mu_x = np.mean(x)
-    mu_y = np.mean(y)
+    mu_x = ops.mean(x)
+    mu_y = ops.mean(y)
 
-    var_x = np.var(x)
-    var_y = np.var(y)
+    var_x = ops.var(x)
+    var_y = ops.var(y)
 
-    return 20 * np.log10(np.abs(mu_x - mu_y) / np.sqrt((var_x + var_y) / 2))
+    return 20 * ops.log10(ops.abs(mu_x - mu_y) / ops.sqrt((var_x + var_y) / 2))
 
 
-@metrics_registry(name="contrast", framework="numpy", supervised=True)
+@metrics_registry(name="contrast", paired=True)
 def contrast(x, y):
     """Contrast ratio"""
-    return 20 * np.log10(x.mean() / y.mean())
+    return 20 * ops.log10(ops.mean(x) / ops.mean(y))
 
 
-@metrics_registry(name="gcnr", framework="numpy", supervised=True)
+@metrics_registry(name="gcnr", paired=True)
 def gcnr(x, y, bins=256):
     """Generalized contrast-to-noise-ratio"""
-    x = x.flatten()
-    y = y.flatten()
+    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 /= f.sum()
-    g /= g.sum()
+    f /= np.sum(f)
+    g /= np.sum(g)
     return 1 - np.sum(np.minimum(f, g))
 
 
-@metrics_registry(name="fwhm", framework="numpy", supervised=False)
+@metrics_registry(name="fwhm", paired=False)
 def fwhm(img):
     """Resolution full width half maxima"""
-    mask = np.nonzero(img >= 0.5 * np.amax(img))[0]
+    mask = ops.nonzero(img >= 0.5 * ops.amax(img))[0]
     return mask[-1] - mask[0]
 
 
-@metrics_registry(name="speckle_res", framework="numpy", supervised=False)
-def speckle_res(img):
-    """TODO: Write speckle edge-spread function resolution code"""
-    raise NotImplementedError
-
-
-@metrics_registry(name="snr", framework="numpy", supervised=False)
+@metrics_registry(name="snr", paired=False)
 def snr(img):
     """Signal to noise ratio"""
-    return img.mean() / img.std()
+    return ops.mean(img) / ops.std(img)
 
 
-@metrics_registry(name="wopt_mae", framework="numpy", supervised=True)
+@metrics_registry(name="wopt_mae", paired=True)
 def wopt_mae(ref, img):
     """Find the optimal weight that minimizes the mean absolute error"""
-    wopt = np.median(ref / img)
+    wopt = ops.median(ref / img)
     return wopt
 
 
-@metrics_registry(name="wopt_mse", framework="numpy", supervised=True)
+@metrics_registry(name="wopt_mse", paired=True)
 def wopt_mse(ref, img):
     """Find the optimal weight that minimizes the mean squared error"""
-    wopt = np.sum(ref * img) / np.sum(img * img)
+    wopt = ops.sum(ref * img) / ops.sum(img * img)
     return wopt
 
 
-@metrics_registry(name="l1loss", framework="numpy", supervised=True)
-def l1loss(x, y):
-    """L1 loss"""
-    return np.abs(x - y).mean()
+@metrics_registry(name="psnr", paired=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)))
 
-@metrics_registry(name="l2loss", framework="numpy", supervised=True)
-def l2loss(x, y):
-    """L2 loss"""
-    return np.sqrt(((x - y) ** 2).mean())
+    Args:
+        y_true (tensor): [None, height, width, channels]
+        y_pred (tensor): [None, height, width, channels]
+        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
 
-@metrics_registry(name="psnr", framework="numpy", supervised=True)
-def psnr(x, y):
-    """Peak signal to noise ratio"""
-    dynamic_range = max(x.max(), y.max()) - min(x.min(), y.min())
-    return 20 * np.log10(dynamic_range / l2loss(x, y))
 
+@metrics_registry(name="mse", paired=True)
+def mse(y_true, y_pred):
+    """Gives the MSE for two input tensors.
+    Args:
+        y_true (tensor)
+        y_pred (tensor)
+    Returns:
+        (float): mean squared error between y_true and y_pred. L2 loss.
+
+    """
+    return _reduce_mean(ops.square(y_true - y_pred))
 
-@metrics_registry(name="ncc", framework="numpy", supervised=True)
-def ncc(x, y):
-    """Normalized cross correlation"""
-    return (x * y).sum() / np.sqrt((x**2).sum() * (y**2).sum())
 
+@metrics_registry(name="mae", paired=True)
+def mae(y_true, y_pred):
+    """Gives the MAE for two input tensors.
+    Args:
+        y_true (tensor)
+        y_pred (tensor)
+    Returns:
+        (float): mean absolute error between y_true and y_pred. L1 loss.
 
-@metrics_registry(name="image_entropy", framework="numpy", supervised=False)
-def image_entropy(image):
-    """Calculate the entropy of the image
+    """
+    return _reduce_mean(ops.abs(y_true - y_pred))
+
+
+@metrics_registry(name="ssim", paired=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,
+    return_map: bool = False,
+    filter_fn=None,
+):
+    """Computes the structural similarity index (SSIM) between image pairs.
+
+    This function is based on the standard SSIM implementation 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.
+
+    This function copied from [`dm_pix.ssim`](https://dm-pix.readthedocs.io/en/latest/api.html#dm_pix.ssim),
+    which is part of the DeepMind's `dm_pix` library. They modeled their implementation
+    after the `tf.image.ssim` function.
+
+    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:
-        image (ndarray): The image for which the entropy is calculated
+        a: First image (or set of images).
+        b: Second image (or set of images).
+        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.).
+        return_map: If True, will cause the per-pixel SSIM "map" to be returned.
+        filter_fn: An optional argument for overriding the filter function used by
+            SSIM, which would otherwise be a 2D Gaussian blur specified by filter_size
+            and filter_sigma.
 
     Returns:
-        float: The entropy of the image
+        Each image's mean SSIM, or a tensor of individual values if `return_map`.
     """
-    marg = np.histogramdd(np.ravel(image), bins=256)[0] / image.size
-    marg = list(filter(lambda p: p > 0, np.ravel(marg)))
-    entropy = -np.sum(np.multiply(marg, np.log2(marg)))
-    return entropy
+
+    if filter_fn is None:
+        # Construct a 1D Gaussian blur filter.
+        hw = filter_size // 2
+        shift = (2 * hw - filter_size + 1) / 2
+        f_i = ((ops.cast(ops.arange(filter_size), "float32") - hw + shift) / filter_sigma) ** 2
+        filt = ops.exp(-0.5 * f_i)
+        filt /= ops.sum(filt)
+
+        # Construct a 1D convolution.
+        def filter_fn_1(z):
+            return tensor_ops.correlate(z, ops.flip(filt), mode="valid")
+
+        # Apply the vectorized filter along the y axis.
+        def filter_fn_y(z):
+            z_flat = ops.reshape(ops.moveaxis(z, -3, -1), (-1, z.shape[-3]))
+            z_filtered_shape = ((z.shape[-4],) if z.ndim == 4 else ()) + (
+                z.shape[-2],
+                z.shape[-1],
+                -1,
+            )
+            _z_filtered = ops.vectorized_map(filter_fn_1, z_flat)
+            z_filtered = ops.moveaxis(ops.reshape(_z_filtered, z_filtered_shape), -1, -3)
+            return z_filtered
+
+        # Apply the vectorized filter along the x axis.
+        def filter_fn_x(z):
+            z_flat = ops.reshape(ops.moveaxis(z, -2, -1), (-1, z.shape[-2]))
+            z_filtered_shape = ((z.shape[-4],) if z.ndim == 4 else ()) + (
+                z.shape[-3],
+                z.shape[-1],
+                -1,
+            )
+            _z_filtered = ops.vectorized_map(filter_fn_1, z_flat)
+            z_filtered = ops.moveaxis(ops.reshape(_z_filtered, z_filtered_shape), -1, -2)
+            return z_filtered
+
+        # Apply the blur in both x and y.
+        filter_fn = lambda z: filter_fn_y(filter_fn_x(z))
+
+    mu0 = filter_fn(a)
+    mu1 = filter_fn(b)
+    mu00 = mu0 * mu0
+    mu11 = mu1 * mu1
+    mu01 = mu0 * mu1
+    sigma00 = filter_fn(a**2) - mu00
+    sigma11 = filter_fn(b**2) - mu11
+    sigma01 = filter_fn(a * b) - mu01
+
+    # Clip the variances and covariances to valid values.
+    # Variance must be non-negative:
+    epsilon = keras.config.epsilon()
+    sigma00 = ops.maximum(epsilon, sigma00)
+    sigma11 = ops.maximum(epsilon, sigma11)
+    sigma01 = ops.sign(sigma01) * ops.minimum(ops.sqrt(sigma00 * sigma11), ops.abs(sigma01))
+
+    c1 = (k1 * max_val) ** 2
+    c2 = (k2 * max_val) ** 2
+    numer = (2 * mu01 + c1) * (2 * sigma01 + c2)
+    denom = (mu00 + mu11 + c1) * (sigma00 + sigma11 + c2)
+    ssim_map = numer / denom
+    ssim_value = ops.mean(ssim_map, axis=tuple(range(-3, 0)))
+    return ssim_map if return_map else ssim_value
+
+
+@metrics_registry(name="ncc", paired=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())
 
 
-@metrics_registry(name="image_sharpness", framework="numpy", supervised=False)
-def image_sharpness(image):
-    """Calculate the sharpness of the image
+@metrics_registry(name="lpips", paired=True)
+def get_lpips(image_range, batch_size=None, clip=False):
+    """
+    Get the Learned Perceptual Image Patch Similarity (LPIPS) metric.
 
     Args:
-        image (ndarray): The image for which the sharpness is calculated
+        image_range (list): The range of the images. Will be translated to [-1, 1] for LPIPS.
+        batch_size (int): The batch size for the LPIPS model.
+        clip (bool): Whether to clip the images to `image_range`.
 
     Returns:
-        float: The sharpness of the image
+        The LPIPS metric function which can be used with [..., h, w, c] tensors in
+        the range `image_range`.
+    """
+    # Get the LPIPS model
+    _lpips = LPIPS.from_preset("lpips")
+    _lpips.trainable = False
+    _lpips.disable_checks = True
+
+    def unstack_lpips(imgs):
+        """Unstack the images and calculate the LPIPS metric."""
+        img1, img2 = ops.unstack(imgs, num=2, axis=-1)
+        return _lpips([img1, img2])
+
+    def lpips(img1, img2, **kwargs):
+        """
+        The LPIPS metric function.
+        Args:
+            img1 (tensor) with shape (..., h, w, c)
+            img2 (tensor) with shape (..., h, w, c)
+        Returns (float): The LPIPS metric between img1 and img2 with shape [...]
+        """
+        # 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])
+
+        imgs = ops.stack([img1, img2], axis=-1)
+        n_batch_dims = ops.ndim(img1) - 3
+        return tensor_ops.func_with_one_batch_dim(
+            unstack_lpips, imgs, n_batch_dims, batch_size=batch_size
+        )
+
+    return lpips
+
+
+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': ...}
     """
-    return np.mean(np.abs(np.gradient(image)))
+
+    def __init__(
+        self,
+        metrics: List[str],
+        image_range: tuple,
+        quantize: bool = False,
+        clip: bool = False,
+        **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.
+            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 = {
+            m: get_metric(m, **reduce_to_signature(metrics_registry[m], kwargs)) for m in metrics
+        }
+
+        # Other settings
+        self.quantize = quantize
+        self.clip = clip
+
+    @staticmethod
+    def _call_metric_fn(fun, y_true, y_pred, average_batch, batch_axes, return_numpy, device):
+        if batch_axes is None:
+            batch_axes = tuple(range(ops.ndim(y_true) - 3))
+        elif not isinstance(batch_axes, (list, tuple)):
+            batch_axes = (batch_axes,)
+
+        # Because most metric functions do not support batching, we vmap over the batch axes.
+        metric_fn = fun
+        for ax in reversed(batch_axes):
+            metric_fn = tensor_ops.vmap(metric_fn, in_axes=ax, _use_torch_vmap=True)
+
+        out = func_on_device(metric_fn, device, y_true, y_pred)
+
+        if average_batch:
+            out = ops.mean(out)
+
+        if return_numpy:
+            out = ops.convert_to_numpy(out)
+        return out
+
+    def _prepocess(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
+
+    def __call__(
+        self,
+        y_true,
+        y_pred,
+        average_batch=True,
+        batch_axes=None,
+        return_numpy=True,
+        device=None,
+    ):
+        """Calculate all metrics and return as a dictionary.
+
+        Args:
+            y_true (tensor): Ground truth images with shape [..., h, w, c]
+            y_pred (tensor): Predicted images with shape [..., h, w, c]
+            average_batch (bool): Whether to average the metrics over the batch dimensions.
+            batch_axes (tuple): The axes corresponding to the batch dimensions. If None, will
+                assume all leading dimensions except the last 3 are batch dimensions.
+            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._prepocess(y_true),
+                self._prepocess(y_pred),
+                average_batch,
+                batch_axes,
+                return_numpy,
+                device,
+            )
+        return results
 
 
 def _sector_reweight_image(image, sector_angle, axis):
@@ -149,10 +449,10 @@ def _sector_reweight_image(image, sector_angle, axis):
             pixel post-scan-conversion.
     """
     height = image.shape[axis]
-    depths = np.arange(height) + 0.5  # center of the pixel as its depth
+    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] * image.ndim
+    shape = [1] * ops.ndim(image)
     shape[axis] = height
-    reweighting_factors = np.reshape(reweighting_factors, shape)
+    reweighting_factors = ops.reshape(reweighting_factors, shape)
     return reweighting_factors * image

zea/models/__init__.py

@@ -2,31 +2,37 @@
 
 ``zea`` contains a collection of models for various tasks, all located in the :mod:`zea.models` package.
 
-Currently, the following models are available (all inherited from :class:`zea.models.BaseModel`):
+See the following dropdown for a list of available models:
 
-- :class:`zea.models.echonet.EchoNetDynamic`: A model for left ventricle segmentation.
-- :class:`zea.models.carotid_segmenter.CarotidSegmenter`: A model for carotid artery segmentation.
-- :class:`zea.models.echonetlvh.EchoNetLVH`: A model for left ventricle hypertrophy segmentation.
-- :class:`zea.models.unet.UNet`: A simple U-Net implementation.
-- :class:`zea.models.lpips.LPIPS`: A model implementing the perceptual similarity metric.
-- :class:`zea.models.taesd.TinyAutoencoder`: A tiny autoencoder model for image compression.
+.. dropdown:: **Available models**
+
+    - :class:`zea.models.echonet.EchoNetDynamic`: A model for left ventricle segmentation.
+    - :class:`zea.models.carotid_segmenter.CarotidSegmenter`: A model for carotid artery segmentation.
+    - :class:`zea.models.echonetlvh.EchoNetLVH`: A model for left ventricle hypertrophy segmentation.
+    - :class:`zea.models.unet.UNet`: A simple U-Net implementation.
+    - :class:`zea.models.lpips.LPIPS`: A model implementing the perceptual similarity metric.
+    - :class:`zea.models.taesd.TinyAutoencoder`: A tiny autoencoder model for image compression.
+    - :class:`zea.models.regional_quality.MobileNetv2RegionalQuality`: A scoring model for myocardial regions in apical views.
+    - :class:`zea.models.lv_segmentation.AugmentedCamusSeg`: A nnU-Net based left ventricle and myocardium segmentation model.
 
 Presets for these models can be found in :mod:`zea.models.presets`.
 
 To use these models, you can import them directly from the :mod:`zea.models` module and load the pretrained weights using the :meth:`from_preset` method. For example:
 
-.. code-block:: python
+.. doctest::
 
-    from zea.models.unet import UNet
+    >>> from zea.models.unet import UNet
 
-    model = UNet.from_preset("unet-echonet-inpainter")
+    >>> model = UNet.from_preset("unet-echonet-inpainter")
 
 You can list all available presets using the :attr:`presets` attribute:
 
-.. code-block:: python
+.. doctest::
 
-    presets = list(UNet.presets.keys())
-    print(f"Available built-in zea presets for UNet: {presets}")
+    >>> from zea.models.unet import UNet
+    >>> presets = list(UNet.presets.keys())
+    >>> print(f"Available built-in zea presets for UNet: {presets}")
+    Available built-in zea presets for UNet: ['unet-echonet-inpainter']
 
 
 zea generative models
@@ -40,19 +46,21 @@ Typically, these models have some additional methods, such as:
 - :meth:`posterior_sample` for drawing samples from the posterior given measurements
 - :meth:`log_density` for computing the log-probability of data under the model
 
-The following generative models are currently available:
+See the following dropdown for a list of available *generative* models:
+
+.. dropdown:: **Available models**
 
-- :class:`zea.models.diffusion.DiffusionModel`: A deep generative diffusion model for ultrasound image generation.
-- :class:`zea.models.gmm.GaussianMixtureModel`: A Gaussian Mixture Model.
+    - :class:`zea.models.diffusion.DiffusionModel`: A deep generative diffusion model for ultrasound image generation.
+    - :class:`zea.models.gmm.GaussianMixtureModel`: A Gaussian Mixture Model.
 
 An example of how to use the :class:`zea.models.diffusion.DiffusionModel` is shown below:
 
-.. code-block:: python
+.. doctest::
 
-    from zea.models.diffusion import DiffusionModel
+    >>> from zea.models.diffusion import DiffusionModel
 
-    model = DiffusionModel.from_preset("diffusion-echonet-dynamic")
-    samples = model.sample(n_samples=4)
+    >>> model = DiffusionModel.from_preset("diffusion-echonet-dynamic")  # doctest: +SKIP
+    >>> samples = model.sample(n_samples=4)  # doctest: +SKIP
 
 
 Contributing and adding new models
@@ -84,7 +92,9 @@ from . import (
     gmm,
     layers,
     lpips,
+    lv_segmentation,
     presets,
+    regional_quality,
     taesd,
     unet,
     utils,

zea/models/base.py

@@ -8,6 +8,7 @@ import importlib
 import keras
 from keras.src.saving.serialization_lib import record_object_after_deserialization
 
+from zea import log
 from zea.internal.core import classproperty
 from zea.models.preset_utils import builtin_presets, get_preset_loader, get_preset_saver
 
@@ -77,7 +78,7 @@ class BaseModel(keras.models.Model):
                 initialized.
             **kwargs: Additional keyword arguments.
 
-        Examples:
+        Example:
             .. code-block:: python
 
                 # Load a Gemma backbone with pre-trained weights.
@@ -96,14 +97,29 @@ class BaseModel(keras.models.Model):
 
         """
         loader = get_preset_loader(preset)
-        model_cls = loader.check_model_class()
-        if not issubclass(model_cls, cls):
-            raise ValueError(
-                f"Saved preset has type `{model_cls.__name__}` which is not "
-                f"a subclass of calling class `{cls.__name__}`. Call "
-                f"`from_preset` directly on `{model_cls.__name__}` instead."
-            )
-        return loader.load_model(model_cls, load_weights, **kwargs)
+        loader_cls = loader.check_model_class()
+        if cls != loader_cls:
+            full_cls_name = f"{cls.__module__}.{cls.__name__}"
+            full_loader_cls_name = f"{loader_cls.__module__}.{loader_cls.__name__}"
+            if issubclass(cls, loader_cls):
+                log.warning(
+                    f"The preset '{preset}' is for model class '{full_loader_cls_name}', but you "
+                    f"are calling from a subclass '{full_cls_name}', so the returned object will "
+                    f"be of type '{full_cls_name}'."
+                )
+            elif issubclass(loader_cls, cls):
+                log.warning(
+                    f"The preset '{preset}' is for model class '{full_loader_cls_name}', "
+                    f"which is a subclass of the calling class '{full_cls_name}', "
+                    f"so the returned object will be of type '{full_cls_name}'."
+                )
+            else:
+                raise ValueError(
+                    f"The preset '{preset}' is for model class '{full_loader_cls_name}', "
+                    f"which is not compatible with the calling class '{full_cls_name}'. "
+                    f"Please call '{full_loader_cls_name}.from_preset()' instead."
+                )
+        return loader.load_model(cls, load_weights, **kwargs)
 
     def save_to_preset(self, preset_dir):
         """Save backbone to a preset directory.
@@ -115,7 +131,7 @@ class BaseModel(keras.models.Model):
         saver.save_model(self)
 
 
-def deserialize_zea_object(config):
+def deserialize_zea_object(config, cls=None):
     """Retrieve the object by deserializing the config dict.
 
     Need to borrow this function from keras and customize a bit to allow
@@ -132,10 +148,10 @@ def deserialize_zea_object(config):
     class_name = config["class_name"]
     inner_config = config["config"] or {}
 
-    module = config.get("module", None)
-    registered_name = config.get("registered_name", class_name)
-
-    cls = _retrieve_class(module, registered_name, config)
+    if cls is None:
+        module = config.get("module", None)
+        registered_name = config.get("registered_name", class_name)
+        cls = _retrieve_class(module, registered_name, config)
 
     if not hasattr(cls, "from_config"):
         raise TypeError(