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

zea/ops/ultrasound.py

@@ -0,0 +1,890 @@
+import uuid
+
+import keras
+import numpy as np
+from keras import ops
+
+from zea import log
+from zea.beamform.beamformer import tof_correction
+from zea.display import scan_convert
+from zea.func.tensor import (
+    apply_along_axis,
+    correlate,
+    extend_n_dims,
+    reshape_axis,
+)
+from zea.func.ultrasound import (
+    channels_to_complex,
+    complex_to_channels,
+    demodulate,
+    envelope_detect,
+    get_low_pass_iq_filter,
+    log_compress,
+    upmix,
+)
+from zea.internal.core import (
+    DEFAULT_DYNAMIC_RANGE,
+    DataTypes,
+)
+from zea.internal.registry import ops_registry
+from zea.ops.base import (
+    ImageOperation,
+    Operation,
+)
+from zea.ops.tensor import (
+    GaussianBlur,
+)
+from zea.simulator import simulate_rf
+
+
+@ops_registry("simulate_rf")
+class Simulate(Operation):
+    """Simulate RF data."""
+
+    # Define operation-specific static parameters
+    STATIC_PARAMS = ["n_ax", "apply_lens_correction"]
+
+    def __init__(self, **kwargs):
+        super().__init__(
+            output_data_type=DataTypes.RAW_DATA,
+            additional_output_keys=["n_ch"],
+            **kwargs,
+        )
+
+    def call(
+        self,
+        scatterer_positions,
+        scatterer_magnitudes,
+        probe_geometry,
+        apply_lens_correction,
+        lens_thickness,
+        lens_sound_speed,
+        sound_speed,
+        n_ax,
+        center_frequency,
+        sampling_frequency,
+        t0_delays,
+        initial_times,
+        element_width,
+        attenuation_coef,
+        tx_apodizations,
+        **kwargs,
+    ):
+        return {
+            self.output_key: simulate_rf(
+                ops.convert_to_tensor(scatterer_positions),
+                ops.convert_to_tensor(scatterer_magnitudes),
+                probe_geometry=probe_geometry,
+                apply_lens_correction=apply_lens_correction,
+                lens_thickness=lens_thickness,
+                lens_sound_speed=lens_sound_speed,
+                sound_speed=sound_speed,
+                n_ax=n_ax,
+                center_frequency=center_frequency,
+                sampling_frequency=sampling_frequency,
+                t0_delays=t0_delays,
+                initial_times=initial_times,
+                element_width=element_width,
+                attenuation_coef=attenuation_coef,
+                tx_apodizations=tx_apodizations,
+            ),
+            "n_ch": 1,  # Simulate always returns RF data (so single channel)
+        }
+
+
+@ops_registry("tof_correction")
+class TOFCorrection(Operation):
+    """Time-of-flight correction operation for ultrasound data."""
+
+    # Define operation-specific static parameters
+    STATIC_PARAMS = ["f_number", "apply_lens_correction"]
+
+    def __init__(self, **kwargs):
+        super().__init__(
+            input_data_type=DataTypes.RAW_DATA,
+            output_data_type=DataTypes.ALIGNED_DATA,
+            **kwargs,
+        )
+
+    def call(
+        self,
+        flatgrid,
+        sound_speed,
+        polar_angles,
+        focus_distances,
+        sampling_frequency,
+        f_number,
+        demodulation_frequency,
+        t0_delays,
+        tx_apodizations,
+        initial_times,
+        probe_geometry,
+        t_peak,
+        tx_waveform_indices,
+        apply_lens_correction=None,
+        lens_thickness=None,
+        lens_sound_speed=None,
+        **kwargs,
+    ):
+        """Perform time-of-flight correction on raw RF data.
+
+        Args:
+            raw_data (ops.Tensor): Raw RF data to correct
+            flatgrid (ops.Tensor): Grid points at which to evaluate the time-of-flight
+            sound_speed (float): Sound speed in the medium
+            polar_angles (ops.Tensor): Polar angles for scan lines
+            focus_distances (ops.Tensor): Focus distances for scan lines
+            sampling_frequency (float): Sampling frequency
+            f_number (float): F-number for apodization
+            demodulation_frequency (float): Demodulation frequency
+            t0_delays (ops.Tensor): T0 delays
+            tx_apodizations (ops.Tensor): Transmit apodizations
+            initial_times (ops.Tensor): Initial times
+            probe_geometry (ops.Tensor): Probe element positions
+            t_peak (float): Time to peak of the transmit pulse
+            tx_waveform_indices (ops.Tensor): Index of the transmit waveform for each
+                transmit. (All zero if there is only one waveform)
+            apply_lens_correction (bool): Whether to apply lens correction
+            lens_thickness (float): Lens thickness
+            lens_sound_speed (float): Sound speed in the lens
+
+        Returns:
+            dict: Dictionary containing tof_corrected_data
+        """
+
+        raw_data = kwargs[self.key]
+
+        tof_kwargs = {
+            "flatgrid": flatgrid,
+            "t0_delays": t0_delays,
+            "tx_apodizations": tx_apodizations,
+            "sound_speed": sound_speed,
+            "probe_geometry": probe_geometry,
+            "initial_times": initial_times,
+            "sampling_frequency": sampling_frequency,
+            "demodulation_frequency": demodulation_frequency,
+            "f_number": f_number,
+            "polar_angles": polar_angles,
+            "focus_distances": focus_distances,
+            "t_peak": t_peak,
+            "tx_waveform_indices": tx_waveform_indices,
+            "apply_lens_correction": apply_lens_correction,
+            "lens_thickness": lens_thickness,
+            "lens_sound_speed": lens_sound_speed,
+        }
+
+        if not self.with_batch_dim:
+            tof_corrected = tof_correction(raw_data, **tof_kwargs)
+        else:
+            tof_corrected = ops.map(
+                lambda data: tof_correction(data, **tof_kwargs),
+                raw_data,
+            )
+
+        return {self.output_key: tof_corrected}
+
+
+@ops_registry("pfield_weighting")
+class PfieldWeighting(Operation):
+    """Weighting aligned data with the pressure field."""
+
+    def __init__(self, **kwargs):
+        super().__init__(
+            input_data_type=DataTypes.ALIGNED_DATA,
+            output_data_type=DataTypes.ALIGNED_DATA,
+            **kwargs,
+        )
+
+    def call(self, flat_pfield=None, **kwargs):
+        """Weight data with pressure field.
+
+        Args:
+            flat_pfield (ops.Tensor): Pressure field weight mask of shape (n_pix, n_tx)
+
+        Returns:
+            dict: Dictionary containing weighted data
+        """
+        data = kwargs[self.key]  # must start with ((batch_size,) n_tx, n_pix, ...)
+
+        if flat_pfield is None:
+            return {self.output_key: data}
+
+        # Swap (n_pix, n_tx) to (n_tx, n_pix)
+        flat_pfield = ops.swapaxes(flat_pfield, 0, 1)
+
+        # Add batch dimension if needed
+        if self.with_batch_dim:
+            pfield_expanded = ops.expand_dims(flat_pfield, axis=0)
+        else:
+            pfield_expanded = flat_pfield
+
+        append_n_dims = ops.ndim(data) - ops.ndim(pfield_expanded)
+        pfield_expanded = extend_n_dims(pfield_expanded, axis=-1, n_dims=append_n_dims)
+
+        # Perform element-wise multiplication with the pressure weight mask
+        weighted_data = data * pfield_expanded
+
+        return {self.output_key: weighted_data}
+
+
+@ops_registry("scan_convert")
+class ScanConvert(Operation):
+    """Scan convert images to cartesian coordinates."""
+
+    STATIC_PARAMS = ["fill_value"]
+
+    def __init__(self, order=1, **kwargs):
+        """Initialize the ScanConvert operation.
+
+        Args:
+            order (int, optional): Interpolation order. Defaults to 1. Currently only
+                GPU support for order=1.
+        """
+        if order > 1:
+            jittable = False
+            log.warning(
+                "GPU support for order > 1 is not available. " + "Disabling jit for ScanConvert."
+            )
+        else:
+            jittable = True
+
+        super().__init__(
+            input_data_type=DataTypes.IMAGE,
+            output_data_type=DataTypes.IMAGE_SC,
+            jittable=jittable,
+            additional_output_keys=[
+                "resolution",
+                "x_lim",
+                "y_lim",
+                "z_lim",
+                "rho_range",
+                "theta_range",
+                "phi_range",
+                "d_rho",
+                "d_theta",
+                "d_phi",
+            ],
+            **kwargs,
+        )
+        self.order = order
+
+    def call(
+        self,
+        rho_range=None,
+        theta_range=None,
+        phi_range=None,
+        resolution=None,
+        coordinates=None,
+        fill_value=None,
+        **kwargs,
+    ):
+        """Scan convert images to cartesian coordinates.
+
+        Args:
+            rho_range (Tuple): Range of the rho axis in the polar coordinate system.
+                Defined in meters.
+            theta_range (Tuple): Range of the theta axis in the polar coordinate system.
+                Defined in radians.
+            phi_range (Tuple): Range of the phi axis in the polar coordinate system.
+                Defined in radians.
+            resolution (float): Resolution of the output image in meters per pixel.
+                if None, the resolution is computed based on the input data.
+            coordinates (Tensor): Coordinates for scan convertion. If None, will be computed
+                based on rho_range, theta_range, phi_range and resolution. If provided, this
+                operation can be jitted.
+            fill_value (float): Value to fill the image with outside the defined region.
+
+        """
+        if fill_value is None:
+            fill_value = np.nan
+
+        data = kwargs[self.key]
+
+        if self._jit_compile and self.jittable:
+            assert coordinates is not None, (
+                "coordinates must be provided to jit scan conversion."
+                "You can set ScanConvert(jit_compile=False) to disable jitting."
+            )
+
+        data_out, parameters = scan_convert(
+            data,
+            rho_range,
+            theta_range,
+            phi_range,
+            resolution,
+            coordinates,
+            fill_value,
+            self.order,
+            with_batch_dim=self.with_batch_dim,
+        )
+
+        return {self.output_key: data_out, **parameters}
+
+
+@ops_registry("demodulate")
+class Demodulate(Operation):
+    """Demodulates the input data to baseband. After this operation, the carrier frequency
+    is removed (0 Hz) and the data is in IQ format stored in two real valued channels."""
+
+    def __init__(self, axis=-3, **kwargs):
+        super().__init__(
+            input_data_type=DataTypes.RAW_DATA,
+            output_data_type=DataTypes.RAW_DATA,
+            jittable=True,
+            additional_output_keys=[
+                "demodulation_frequency",
+                "center_frequency",
+                "n_ch",
+            ],
+            **kwargs,
+        )
+        self.axis = axis
+
+    def call(self, center_frequency=None, sampling_frequency=None, **kwargs):
+        data = kwargs[self.key]
+
+        demodulation_frequency = center_frequency
+
+        # Split the complex signal into two channels
+        iq_data_two_channel = demodulate(
+            data=data,
+            center_frequency=center_frequency,
+            sampling_frequency=sampling_frequency,
+            axis=self.axis,
+        )
+
+        return {
+            self.output_key: iq_data_two_channel,
+            "demodulation_frequency": demodulation_frequency,
+            "center_frequency": 0.0,
+            "n_ch": 2,
+        }
+
+
+@ops_registry("fir_filter")
+class FirFilter(Operation):
+    """Apply a FIR filter to the input signal using convolution.
+
+    Looks for the filter taps in the input dictionary using the specified ``filter_key``.
+    """
+
+    def __init__(
+        self,
+        axis: int,
+        complex_channels: bool = False,
+        filter_key: str = "fir_filter_taps",
+        **kwargs,
+    ):
+        """
+        Args:
+            axis (int): Axis along which to apply the filter. Cannot be the batch dimension.
+                When using ``complex_channels=True``, the complex channels are removed to convert
+                to complex numbers before filtering, so adjust the ``axis`` accordingly!
+            complex_channels (bool): Whether the last dimension of the input signal represents
+                complex channels (real and imaginary parts). When True, it will convert the signal
+                to ``complex`` dtype before filtering and convert it back to two channels
+                after filtering.
+            filter_key (str): Key in the input dictionary where the FIR filter taps are stored.
+                Default is "fir_filter_taps".
+        """
+        super().__init__(**kwargs)
+        self._check_axis(axis)
+
+        self.axis = axis
+        self.complex_channels = complex_channels
+        self.filter_key = filter_key
+
+    def _check_axis(self, axis, ndim=None):
+        """Check if the axis is valid."""
+        if ndim is not None:
+            if axis < -ndim or axis >= ndim:
+                raise ValueError(f"Axis {axis} is out of bounds for array of dimension {ndim}.")
+
+        if self.with_batch_dim and (axis == 0 or (ndim is not None and axis == -ndim)):
+            raise ValueError("Cannot apply FIR filter along batch dimension.")
+
+    @property
+    def valid_keys(self):
+        """Get the valid keys for the `call` method."""
+        return self._valid_keys.union({self.filter_key})
+
+    def call(self, **kwargs):
+        signal = kwargs[self.key]
+        fir_filter_taps = kwargs[self.filter_key]
+
+        if self.complex_channels:
+            signal = channels_to_complex(signal)
+
+        self._check_axis(self.axis, ndim=ops.ndim(signal))
+
+        def _convolve(signal):
+            """Apply the filter to the signal using correlation."""
+            return correlate(signal, fir_filter_taps[::-1], mode="same")
+
+        filtered_signal = apply_along_axis(_convolve, self.axis, signal)
+
+        if self.complex_channels:
+            filtered_signal = complex_to_channels(filtered_signal)
+
+        return {self.output_key: filtered_signal}
+
+
+@ops_registry("low_pass_filter")
+class LowPassFilter(FirFilter):
+    """Apply a low-pass FIR filter to the input signal using convolution.
+
+    It is recommended to use :class:`FirFilter` with pre-computed filter taps for jittable
+    operations. The :class:`LowPassFilter` operation itself is not jittable and is provided
+    for convenience only.
+
+    Uses :func:`get_low_pass_iq_filter` to compute the filter taps.
+    """
+
+    def __init__(self, axis: int, complex_channels: bool = False, num_taps: int = 128, **kwargs):
+        """Initialize the LowPassFilter operation.
+
+        Args:
+            axis (int): Axis along which to apply the filter. Cannot be the batch dimension.
+                When using ``complex_channels=True``, the complex channels are removed to convert
+                to complex numbers before filtering, so adjust the ``axis`` accordingly.
+            complex_channels (bool): Whether the last dimension of the input signal represents
+                complex channels (real and imaginary parts). When True, it will convert the signal
+                to ``complex`` dtype before filtering and convert it back to two channels
+                after filtering.
+            num_taps (int): Number of taps in the FIR filter. Default is 128.
+        """
+        self._random_suffix = str(uuid.uuid4())
+        kwargs.pop("filter_key", None)
+        kwargs.pop("jittable", None)
+        super().__init__(
+            axis=axis,
+            complex_channels=complex_channels,
+            filter_key=f"low_pass_{self._random_suffix}",
+            jittable=False,
+            **kwargs,
+        )
+        self.num_taps = num_taps
+
+    def call(self, bandwidth, sampling_frequency, center_frequency, **kwargs):
+        lpf = get_low_pass_iq_filter(
+            self.num_taps,
+            ops.convert_to_numpy(sampling_frequency).item(),
+            ops.convert_to_numpy(center_frequency).item(),
+            ops.convert_to_numpy(bandwidth).item(),
+        )
+        kwargs[self.filter_key] = lpf
+        return super().call(**kwargs)
+
+
+@ops_registry("channels_to_complex")
+class ChannelsToComplex(Operation):
+    def call(self, **kwargs):
+        data = kwargs[self.key]
+        output = channels_to_complex(data)
+        return {self.output_key: output}
+
+
+@ops_registry("complex_to_channels")
+class ComplexToChannels(Operation):
+    def __init__(self, axis=-1, **kwargs):
+        super().__init__(**kwargs)
+        self.axis = axis
+
+    def call(self, **kwargs):
+        data = kwargs[self.key]
+        output = complex_to_channels(data, axis=self.axis)
+        return {self.output_key: output}
+
+
+@ops_registry("lee_filter")
+class LeeFilter(ImageOperation):
+    """
+    The Lee filter is a speckle reduction filter commonly used in synthetic aperture radar (SAR)
+    and ultrasound image processing. It smooths the image while preserving edges and details.
+    This implementation uses Gaussian filter for local statistics and treats channels independently.
+
+    Lee, J.S. (1980). Digital image enhancement and noise filtering by use of local statistics.
+    IEEE Transactions on Pattern Analysis and Machine Intelligence, (2), 165-168.
+    """
+
+    def __init__(self, sigma=3, kernel_size=None, pad_mode="symmetric", **kwargs):
+        """
+        Args:
+            sigma (float): Standard deviation for Gaussian kernel. Default is 3.
+            kernel_size (int, optional): Size of the Gaussian kernel. If None,
+                it will be calculated based on sigma.
+            pad_mode (str): Padding mode to be used for Gaussian blur. Default is "symmetric".
+        """
+        super().__init__(**kwargs)
+        self.sigma = sigma
+        self.kernel_size = kernel_size
+        self.pad_mode = pad_mode
+
+        # Create a GaussianBlur instance for computing local statistics
+        self.gaussian_blur = GaussianBlur(
+            sigma=self.sigma,
+            kernel_size=self.kernel_size,
+            pad_mode=self.pad_mode,
+            with_batch_dim=self.with_batch_dim,
+            jittable=self._jittable,
+            key="data",
+        )
+
+    @property
+    def with_batch_dim(self):
+        """Get the with_batch_dim property of the LeeFilter operation."""
+        return self._with_batch_dim
+
+    @with_batch_dim.setter
+    def with_batch_dim(self, value):
+        """Set the with_batch_dim property of the LeeFilter operation."""
+        self._with_batch_dim = value
+        if hasattr(self, "gaussian_blur"):
+            self.gaussian_blur.with_batch_dim = value
+
+    def call(self, **kwargs):
+        """Apply the Lee filter to the input data.
+
+        Args:
+            data (ops.Tensor): Input image data of shape (height, width, channels) with
+                optional batch dimension if ``self.with_batch_dim``.
+        """
+        super().call(**kwargs)
+        data = kwargs.pop(self.key)
+
+        # Apply Gaussian blur to get local mean
+        img_mean = self.gaussian_blur.call(data=data, **kwargs)[self.gaussian_blur.output_key]
+
+        # Apply Gaussian blur to squared data to get local squared mean
+        img_sqr_mean = self.gaussian_blur.call(
+            data=data**2,
+            **kwargs,
+        )[self.gaussian_blur.output_key]
+
+        # Calculate local variance
+        img_variance = img_sqr_mean - img_mean**2
+
+        # Calculate global variance (per channel)
+        overall_variance = ops.var(data, axis=(-3, -2), keepdims=True)
+
+        # Calculate adaptive weights
+        eps = keras.config.epsilon()
+        img_weights = img_variance / (img_variance + overall_variance + eps)
+
+        # Apply Lee filter formula
+        img_output = img_mean + img_weights * (data - img_mean)
+
+        return {self.output_key: img_output}
+
+
+@ops_registry("companding")
+class Companding(Operation):
+    """Companding according to the A- or μ-law algorithm.
+
+    Invertible compressing operation. Used to compress
+    dynamic range of input data (and subsequently expand).
+
+    μ-law companding:
+    https://en.wikipedia.org/wiki/%CE%9C-law_algorithm
+    A-law companding:
+    https://en.wikipedia.org/wiki/A-law_algorithm
+
+    Args:
+        expand (bool, optional): If set to False (default),
+            data is compressed, else expanded.
+        comp_type (str): either `a` or `mu`.
+        mu (float, optional): compression parameter. Defaults to 255.
+        A (float, optional): compression parameter. Defaults to 87.6.
+    """
+
+    def __init__(self, expand=False, comp_type="mu", **kwargs):
+        super().__init__(**kwargs)
+        self.expand = expand
+        self.comp_type = comp_type.lower()
+        if self.comp_type not in ["mu", "a"]:
+            raise ValueError("comp_type must be 'mu' or 'a'.")
+
+        if self.comp_type == "mu":
+            self._compand_func = self._mu_law_expand if self.expand else self._mu_law_compress
+        else:
+            self._compand_func = self._a_law_expand if self.expand else self._a_law_compress
+
+    @staticmethod
+    def _mu_law_compress(x, mu=255, **kwargs):
+        x = ops.clip(x, -1, 1)
+        return ops.sign(x) * ops.log(1.0 + mu * ops.abs(x)) / ops.log(1.0 + mu)
+
+    @staticmethod
+    def _mu_law_expand(y, mu=255, **kwargs):
+        y = ops.clip(y, -1, 1)
+        return ops.sign(y) * ((1.0 + mu) ** ops.abs(y) - 1.0) / mu
+
+    @staticmethod
+    def _a_law_compress(x, A=87.6, **kwargs):
+        x = ops.clip(x, -1, 1)
+        x_sign = ops.sign(x)
+        x_abs = ops.abs(x)
+        A_log = ops.log(A)
+        val1 = x_sign * A * x_abs / (1.0 + A_log)
+        val2 = x_sign * (1.0 + ops.log(A * x_abs)) / (1.0 + A_log)
+        y = ops.where((x_abs >= 0) & (x_abs < (1.0 / A)), val1, val2)
+        return y
+
+    @staticmethod
+    def _a_law_expand(y, A=87.6, **kwargs):
+        y = ops.clip(y, -1, 1)
+        y_sign = ops.sign(y)
+        y_abs = ops.abs(y)
+        A_log = ops.log(A)
+        val1 = y_sign * y_abs * (1.0 + A_log) / A
+        val2 = y_sign * ops.exp(y_abs * (1.0 + A_log) - 1.0) / A
+        x = ops.where((y_abs >= 0) & (y_abs < (1.0 / (1.0 + A_log))), val1, val2)
+        return x
+
+    def call(self, mu=255, A=87.6, **kwargs):
+        data = kwargs[self.key]
+
+        mu = ops.cast(mu, data.dtype)
+        A = ops.cast(A, data.dtype)
+
+        data_out = self._compand_func(data, mu=mu, A=A)
+        return {self.output_key: data_out}
+
+
+@ops_registry("downsample")
+class Downsample(Operation):
+    """Downsample data along a specific axis."""
+
+    def __init__(self, factor: int = 1, phase: int = 0, axis: int = -3, **kwargs):
+        super().__init__(
+            additional_output_keys=["sampling_frequency", "n_ax"],
+            **kwargs,
+        )
+        if factor < 1:
+            raise ValueError("Downsample factor must be >= 1.")
+        if phase < 0 or phase >= factor:
+            raise ValueError("phase must satisfy 0 <= phase < factor.")
+        self.factor = factor
+        self.phase = phase
+        self.axis = axis
+
+    def call(self, sampling_frequency=None, n_ax=None, **kwargs):
+        data = kwargs[self.key]
+        length = ops.shape(data)[self.axis]
+        sample_idx = ops.arange(self.phase, length, self.factor)
+        data_downsampled = ops.take(data, sample_idx, axis=self.axis)
+
+        output = {self.output_key: data_downsampled}
+        # downsampling also affects the sampling frequency
+        if sampling_frequency is not None:
+            sampling_frequency = sampling_frequency / self.factor
+            output["sampling_frequency"] = sampling_frequency
+        if n_ax is not None:
+            n_ax = n_ax // self.factor
+            output["n_ax"] = n_ax
+        return output
+
+
+@ops_registry("anisotropic_diffusion")
+class AnisotropicDiffusion(Operation):
+    """Speckle Reducing Anisotropic Diffusion (SRAD) filter.
+
+    Reference:
+    - https://www.researchgate.net/publication/5602035_Speckle_reducing_anisotropic_diffusion
+    - https://nl.mathworks.com/matlabcentral/fileexchange/54044-image-despeckle-filtering-toolbox
+    """
+
+    def call(self, niter=100, lmbda=0.1, rect=None, eps=1e-6, **kwargs):
+        """Anisotropic diffusion filter.
+
+        Assumes input data is non-negative.
+
+        Args:
+            niter: Number of iterations.
+            lmbda: Lambda parameter.
+            rect: Rectangle [x1, y1, x2, y2] for homogeneous noise (optional).
+            eps: Small epsilon for stability.
+        Returns:
+            Filtered image (2D tensor or batch of images).
+        """
+        data = kwargs[self.key]
+
+        if not self.with_batch_dim:
+            data = ops.expand_dims(data, axis=0)
+
+        batch_size = ops.shape(data)[0]
+
+        results = []
+        for i in range(batch_size):
+            image = data[i]
+            image_out = self._anisotropic_diffusion_single(image, niter, lmbda, rect, eps)
+            results.append(image_out)
+
+        result = ops.stack(results, axis=0)
+
+        if not self.with_batch_dim:
+            result = ops.squeeze(result, axis=0)
+
+        return {self.output_key: result}
+
+    def _anisotropic_diffusion_single(self, image, niter, lmbda, rect, eps):
+        """Apply anisotropic diffusion to a single image (2D)."""
+        image = ops.exp(image)
+        M, N = image.shape
+
+        for _ in range(niter):
+            iN = ops.concatenate([image[1:], ops.zeros((1, N), dtype=image.dtype)], axis=0)
+            iS = ops.concatenate([ops.zeros((1, N), dtype=image.dtype), image[:-1]], axis=0)
+            jW = ops.concatenate([image[:, 1:], ops.zeros((M, 1), dtype=image.dtype)], axis=1)
+            jE = ops.concatenate([ops.zeros((M, 1), dtype=image.dtype), image[:, :-1]], axis=1)
+
+            if rect is not None:
+                x1, y1, x2, y2 = rect
+                imageuniform = image[x1:x2, y1:y2]
+                q0_squared = (ops.std(imageuniform) / (ops.mean(imageuniform) + eps)) ** 2
+
+            dN = iN - image
+            dS = iS - image
+            dW = jW - image
+            dE = jE - image
+
+            G2 = (dN**2 + dS**2 + dW**2 + dE**2) / (image**2 + eps)
+            L = (dN + dS + dW + dE) / (image + eps)
+            num = (0.5 * G2) - ((1 / 16) * (L**2))
+            den = (1 + ((1 / 4) * L)) ** 2
+            q_squared = num / (den + eps)
+
+            if rect is not None:
+                den = (q_squared - q0_squared) / (q0_squared * (1 + q0_squared) + eps)
+            c = 1.0 / (1 + den)
+            cS = ops.concatenate([ops.zeros((1, N), dtype=image.dtype), c[:-1]], axis=0)
+            cE = ops.concatenate([ops.zeros((M, 1), dtype=image.dtype), c[:, :-1]], axis=1)
+
+            D = (cS * dS) + (c * dN) + (cE * dE) + (c * dW)
+            image = image + (lmbda / 4) * D
+
+        result = ops.log(image)
+        return result
+
+
+@ops_registry("envelope_detect")
+class EnvelopeDetect(Operation):
+    """Envelope detection of RF signals."""
+
+    def __init__(
+        self,
+        axis=-3,
+        **kwargs,
+    ):
+        super().__init__(
+            input_data_type=DataTypes.BEAMFORMED_DATA,
+            output_data_type=DataTypes.ENVELOPE_DATA,
+            **kwargs,
+        )
+        self.axis = axis
+
+    def call(self, **kwargs):
+        """
+        Args:
+            - data (Tensor): The beamformed data of shape (..., grid_size_z, grid_size_x, n_ch).
+        Returns:
+            - envelope_data (Tensor): The envelope detected data
+                of shape (..., grid_size_z, grid_size_x).
+        """
+        data = kwargs[self.key]
+
+        data = envelope_detect(data, axis=self.axis)
+
+        return {self.output_key: data}
+
+
+@ops_registry("upmix")
+class UpMix(Operation):
+    """Upmix IQ data to RF data."""
+
+    def __init__(
+        self,
+        upsampling_rate=1,
+        **kwargs,
+    ):
+        super().__init__(
+            **kwargs,
+        )
+        self.upsampling_rate = upsampling_rate
+
+    def call(
+        self,
+        sampling_frequency=None,
+        center_frequency=None,
+        **kwargs,
+    ):
+        data = kwargs[self.key]
+
+        if data.shape[-1] == 1:
+            log.warning("Upmixing is not applicable to RF data.")
+            return {self.output_key: data}
+        elif data.shape[-1] == 2:
+            data = channels_to_complex(data)
+
+        data = upmix(data, sampling_frequency, center_frequency, self.upsampling_rate)
+        data = ops.expand_dims(data, axis=-1)
+        return {self.output_key: data}
+
+
+@ops_registry("log_compress")
+class LogCompress(Operation):
+    """Logarithmic compression of data."""
+
+    def __init__(self, clip: bool = True, **kwargs):
+        """Initialize the LogCompress operation.
+
+        Args:
+            clip (bool): Whether to clip the output to a dynamic range. Defaults to True.
+        """
+        super().__init__(
+            input_data_type=DataTypes.ENVELOPE_DATA,
+            output_data_type=DataTypes.IMAGE,
+            **kwargs,
+        )
+        self.clip = clip
+
+    def call(self, dynamic_range=None, **kwargs):
+        """Apply logarithmic compression to data.
+
+        Args:
+            dynamic_range (tuple, optional): Dynamic range in dB. Defaults to (-60, 0).
+
+        Returns:
+            dict: Dictionary containing log-compressed data
+        """
+        data = kwargs[self.key]
+
+        if dynamic_range is None:
+            dynamic_range = ops.array(DEFAULT_DYNAMIC_RANGE)
+        dynamic_range = ops.cast(dynamic_range, data.dtype)
+
+        compressed_data = log_compress(data)
+        if self.clip:
+            compressed_data = ops.clip(compressed_data, dynamic_range[0], dynamic_range[1])
+
+        return {self.output_key: compressed_data}
+
+
+@ops_registry("reshape_grid")
+class ReshapeGrid(Operation):
+    """Reshape flat grid data to grid shape."""
+
+    def __init__(self, axis=0, **kwargs):
+        super().__init__(**kwargs)
+        self.axis = axis
+
+    def call(self, grid, **kwargs):
+        """
+        Args:
+            - data (Tensor): The flat grid data of shape (..., n_pix, ...).
+        Returns:
+            - reshaped_data (Tensor): The reshaped data of shape (..., grid.shape, ...).
+        """
+        data = kwargs[self.key]
+        reshaped_data = reshape_axis(data, grid.shape[:-1], self.axis + int(self.with_batch_dim))
+        return {self.output_key: reshaped_data}