zea 0.0.4__py3-none-any.whl → 0.0.6__py3-none-any.whl

zea/metrics.py

@@ -1,131 +1,423 @@
-"""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.models.lpips import LPIPS
+from zea.utils import reduce_to_signature, 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:
+        .. code-block:: python
+
+            metrics = zea.metrics.Metrics(["psnr", "lpips"], image_range=[0, 255])
+            result = metrics(y_true, y_pred)
+            print(result)  # {"psnr": 30.5, "lpips": 0.15}
     """
-    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)
+
+        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 +441,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

@@ -4,8 +4,9 @@
 
 Currently, the following models are available (all inherited from :class:`zea.models.BaseModel`):
 
-- :class:`zea.models.echonet.EchoNetDynamic`: A model for echocardiography segmentation.
+- :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.
@@ -16,7 +17,7 @@ To use these models, you can import them directly from the :mod:`zea.models` mod
 
 .. code-block:: python
 
-    from zea.models import UNet
+    from zea.models.unet import UNet
 
     model = UNet.from_preset("unet-echonet-inpainter")
 
@@ -48,7 +49,7 @@ An example of how to use the :class:`zea.models.diffusion.DiffusionModel` is sho
 
 .. code-block:: python
 
-    from zea.models import DiffusionModel
+    from zea.models.diffusion import DiffusionModel
 
     model = DiffusionModel.from_preset("diffusion-echonet-dynamic")
     samples = model.sample(n_samples=4)
@@ -74,9 +75,11 @@ The following steps are recommended when adding a new model:
 
 from . import (
     carotid_segmenter,
+    deeplabv3,
     dense,
     diffusion,
     echonet,
+    echonetlvh,
     generative,
     gmm,
     layers,

zea/models/deeplabv3.py

@@ -0,0 +1,131 @@
+"""DeepLabV3+ architecture for multi-class segmentation. For more details see https://arxiv.org/abs/1802.02611."""
+
+import keras
+from keras import layers, ops
+
+
+def convolution_block(
+    block_input,
+    num_filters=256,
+    kernel_size=3,
+    dilation_rate=1,
+    use_bias=False,
+):
+    """
+    Create a convolution block with batch normalization and ReLU activation.
+
+    This is a standard building block used throughout the DeepLabV3+ architecture,
+    consisting of Conv2D -> BatchNormalization -> ReLU.
+
+    Args:
+        block_input (Tensor): Input tensor to the convolution block
+        num_filters (int): Number of output filters/channels. Defaults to 256.
+        kernel_size (int): Size of the convolution kernel. Defaults to 3.
+        dilation_rate (int): Dilation rate for dilated convolution. Defaults to 1.
+        use_bias (bool): Whether to use bias in the convolution layer. Defaults to False.
+
+    Returns:
+        Tensor: Output tensor after convolution, batch normalization, and ReLU
+    """
+    x = layers.Conv2D(
+        num_filters,
+        kernel_size=kernel_size,
+        dilation_rate=dilation_rate,
+        padding="same",
+        use_bias=use_bias,
+        kernel_initializer=keras.initializers.HeNormal(),
+    )(block_input)
+    x = layers.BatchNormalization()(x)
+    return ops.nn.relu(x)
+
+
+def DilatedSpatialPyramidPooling(dspp_input):
+    """
+    Implement Atrous Spatial Pyramid Pooling (ASPP) module.
+
+    ASPP captures multi-scale context by applying parallel atrous convolutions
+    with different dilation rates. This helps the model understand objects
+    at multiple scales.
+
+    The module consists of:
+    - Global average pooling branch
+    - 1x1 convolution branch
+    - 3x3 convolutions with dilation rates 6, 12, and 18
+
+    Reference: https://arxiv.org/abs/1706.05587
+
+    Args:
+        dspp_input (Tensor): Input feature tensor from encoder
+
+    Returns:
+        Tensor: Multi-scale feature representation
+    """
+    dims = dspp_input.shape
+    x = layers.AveragePooling2D(pool_size=(dims[-3], dims[-2]))(dspp_input)
+    x = convolution_block(x, kernel_size=1, use_bias=True)
+    out_pool = layers.UpSampling2D(
+        size=(dims[-3] // x.shape[1], dims[-2] // x.shape[2]),
+        interpolation="bilinear",
+    )(x)
+
+    out_1 = convolution_block(dspp_input, kernel_size=1, dilation_rate=1)
+    out_6 = convolution_block(dspp_input, kernel_size=3, dilation_rate=6)
+    out_12 = convolution_block(dspp_input, kernel_size=3, dilation_rate=12)
+    out_18 = convolution_block(dspp_input, kernel_size=3, dilation_rate=18)
+
+    x = layers.Concatenate(axis=-1)([out_pool, out_1, out_6, out_12, out_18])
+    output = convolution_block(x, kernel_size=1)
+    return output
+
+
+def DeeplabV3Plus(image_shape, num_classes, pretrained_weights=None):
+    """
+    Build DeepLabV3+ model for semantic segmentation.
+
+    DeepLabV3+ combines the benefits of spatial pyramid pooling and encoder-decoder
+    architecture. It uses a ResNet50 backbone as encoder, ASPP for multi-scale
+    feature extraction, and a simple decoder for recovering spatial details.
+
+    Architecture:
+    1. Encoder: ResNet50 backbone with atrous convolutions
+    2. ASPP: Multi-scale feature extraction
+    3. Decoder: Simple decoder with skip connections
+    4. Output: Final segmentation prediction
+
+    Reference: https://arxiv.org/abs/1802.02611
+
+    Args:
+        image_shape (tuple): Input image shape as (height, width, channels)
+        num_classes (int): Number of output classes for segmentation
+        pretrained_weights (str, optional): Pretrained weights for ResNet50 backbone.
+                                          Defaults to None.
+
+    Returns:
+        keras.Model: Complete DeepLabV3+ model
+    """
+    model_input = keras.Input(shape=image_shape)
+    # 3-channel grayscale as repeated single channel for ResNet50
+    model_input_3_channel = ops.concatenate([model_input, model_input, model_input], axis=-1)
+    preprocessed = keras.applications.resnet50.preprocess_input(model_input_3_channel)
+    resnet50 = keras.applications.ResNet50(
+        weights=pretrained_weights, include_top=False, input_tensor=preprocessed
+    )
+    x = resnet50.get_layer("conv4_block6_2_relu").output
+    x = DilatedSpatialPyramidPooling(x)
+
+    input_a = layers.UpSampling2D(
+        size=(image_shape[0] // 4 // x.shape[1], image_shape[1] // 4 // x.shape[2]),
+        interpolation="bilinear",
+    )(x)
+    input_b = resnet50.get_layer("conv2_block3_2_relu").output
+    input_b = convolution_block(input_b, num_filters=48, kernel_size=1)
+
+    x = layers.Concatenate(axis=-1)([input_a, input_b])
+    x = convolution_block(x)
+    x = convolution_block(x)
+    x = layers.UpSampling2D(
+        size=(image_shape[0] // x.shape[1], image_shape[1] // x.shape[2]),
+        interpolation="bilinear",
+    )(x)
+    model_output = layers.Conv2D(num_classes, kernel_size=(1, 1), padding="same")(x)
+    return keras.Model(inputs=model_input, outputs=model_output)