waveletvit 0.0.1__py3-none-any.whl

waveletvit/BandDecoder.py

@@ -0,0 +1,146 @@
+"""BandDecoder: reconstruct full-resolution wavelet bands from compressed grids.
+
+Reverse of the separable strided conv tokenization in WaveletViTBlockND.
+Per-level decoder shared across HP bands. Uses axis-by-axis ConvTranspose1D.
+
+Copyright 2026 Kishore Kumar Tarafdar.
+Licensed under the Apache License, Version 2.0. See LICENSE for details.
+"""
+from __future__ import annotations
+
+from typing import List
+
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras import layers
+
+
+def _apply_convtranspose1d_along_axis(x, conv, axis, dims=3):
+    """Apply Conv1D transpose along a specific spatial axis of an N-D tensor."""
+    spatial = list(range(1, 1 + dims))
+    target = 1 + axis
+    perm = [0, target] + [s for s in spatial if s != target] + [dims + 1]
+    x_t = tf.transpose(x, perm)
+    shape_t = tf.shape(x_t)
+    T = shape_t[1]
+    C = shape_t[-1]
+    others = tf.reduce_prod(shape_t[2:-1])
+    x_flat = tf.reshape(x_t, [shape_t[0] * others, T, C])
+    y_flat = conv(x_flat)
+    T_out = tf.shape(y_flat)[1]
+    C_out = tf.shape(y_flat)[-1]
+    y_t = tf.reshape(y_flat, tf.concat([shape_t[:1], [T_out], shape_t[2:-1], [C_out]], axis=0))
+    inv_perm = [0] * (dims + 2)
+    for i, p in enumerate(perm):
+        inv_perm[p] = i
+    return tf.transpose(y_t, inv_perm)
+
+
+@tf.keras.utils.register_keras_serializable(package='waveletvit')
+class BandDecoder(keras.layers.Layer):
+    """Decode compressed 3D grids back to full-resolution wavelet bands.
+
+    Reverse of WaveletViTBlockND's separable strided conv tokenization.
+    Per-level decoder: shared ConvTranspose1D across 7 HP bands at same level.
+
+    Args:
+        levels: number of DWT levels.
+        groups: number of HP bands per level (7 for 3D DWT).
+        token_dim: channel dim F of compressed grids.
+        stride: spatial stride factor used in tokenization.
+        target_spatials: list of target spatial sizes per level, e.g. [32, 16, 8].
+        out_channels: number of channels in reconstructed wavelet bands.
+    """
+
+    def __init__(self, levels: int, groups: int, token_dim: int, stride: int,
+                 target_spatials: List[int], out_channels: int = 1,
+                 name: str = 'band_decoder', **kwargs):
+        super().__init__(name=name, **kwargs)
+        self.levels = int(levels)
+        self.groups = int(groups)
+        self.F = int(token_dim)
+        self.stride = int(stride)
+        self.target_spatials = [int(s) for s in target_spatials]
+        self.out_channels = int(out_channels)
+        self.dims = 3
+
+        # Per-level: 3 ConvTranspose1D (one per axis) + final Conv1D to raw band channels
+        self._up_convs = []
+        self._out_convs = []
+        Conv1DT = layers.Conv1DTranspose
+        for l in range(self.levels):
+            s = min(self.stride, self.target_spatials[l])
+            axis_convs = [
+                Conv1DT(self.F, s, strides=s, padding='valid', activation='relu',
+                        name=f'{name}_l{l}_ax{ax}')
+                for ax in range(self.dims)
+            ]
+            self._up_convs.append(axis_convs)
+            self._out_convs.append(layers.Conv1D(self.out_channels, 1,
+                kernel_initializer='zeros', bias_initializer='zeros',
+                name=f'{name}_l{l}_out'))
+
+        # LP decoder: simple Conv to raw band channels; zero-init for stable start
+        self._lp_out = layers.Conv1D(self.out_channels, 1,
+            kernel_initializer='zeros', bias_initializer='zeros',
+            name=f'{name}_lp_out')
+
+    def call(self, compressed_per_level: List[tf.Tensor], lp_compressed: tf.Tensor):
+        """Decode compressed features to raw wavelet bands.
+
+        Args:
+            compressed_per_level: list of (B, comp^3, groups*F) stacked features per level.
+            lp_compressed: (B, comp_lp^3, F_lp); deepest LP compressed features.
+
+        Returns:
+            bands: list[level] of {'lp': Tensor, 'hp': List[Tensor]}; raw wavelet bands.
+        """
+        bands = []
+        for l in range(self.levels):
+            stacked = compressed_per_level[l]  # (B, comp^3, groups*F)
+            # Split into per-band features
+            hp_feats = tf.split(stacked, self.groups, axis=-1)  # list of (B, comp^3, F)
+
+            hp_bands = []
+            target = self.target_spatials[l]
+            for g, feat in enumerate(hp_feats):
+                t = feat
+                for ax, conv in enumerate(self._up_convs[l]):
+                    t = _apply_convtranspose1d_along_axis(t, conv, ax, self.dims)
+                # Trim to target spatial size
+                t = t[:, :target, :target, :target, :]
+                # Project to raw band channels
+                shape = tf.shape(t)
+                t_flat = tf.reshape(t, [shape[0] * shape[1] * shape[2] * shape[3], 1, self.F])
+                t_out = self._out_convs[l](t_flat)
+                t_out = tf.reshape(t_out, [shape[0], shape[1], shape[2], shape[3], self.out_channels])
+                hp_bands.append(t_out)
+
+            # LP: only at deepest level
+            if l == self.levels - 1:
+                lp_shape = tf.shape(lp_compressed)
+                lp_flat = tf.reshape(lp_compressed, [lp_shape[0] * lp_shape[1] * lp_shape[2] * lp_shape[3], 1, lp_shape[-1]])
+                lp_out = self._lp_out(lp_flat)
+                lp_out = tf.reshape(lp_out, [lp_shape[0], lp_shape[1], lp_shape[2], lp_shape[3], self.out_channels])
+            else:
+                lp_out = None  # non-deepest LP is passthrough in IDWT
+
+            bands.append({'lp': lp_out, 'hp': hp_bands})
+
+        # For IDWT, non-deepest LP comes from the reconstruction cascade; set to zeros placeholder
+        for l in range(self.levels - 1):
+            if bands[l]['lp'] is None:
+                hp_shape = tf.shape(bands[l]['hp'][0])
+                bands[l]['lp'] = tf.zeros_like(bands[l]['hp'][0])
+
+        return bands
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({
+            'levels': self.levels, 'groups': self.groups,
+            'token_dim': self.F, 'stride': self.stride,
+            'target_spatials': self.target_spatials,
+            'out_channels': self.out_channels,
+        })
+        return config

waveletvit/UWaveletViTNet3D.py

@@ -0,0 +1,213 @@
+"""UWaveletViTNet3D: 3D U-Net with a multilevel wavelet-band transformer bottleneck.
+
+Copyright 2026 Kishore Kumar Tarafdar.
+Licensed under the Apache License, Version 2.0. See LICENSE for details.
+"""
+from __future__ import annotations
+
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras.layers import (
+    Activation,
+    Concatenate,
+    Conv3D,
+    Conv3DTranspose,
+    Input,
+    MaxPool3D,
+)
+
+from waveletvit.WaveletViT import (
+    WaveletAssembler3D,
+    WaveletTokenizer3D,
+    WaveletViTBlock3D,
+)
+
+__all__ = [
+    "waveletvit_bottleneck_factory",
+    "UWaveletViTNet3D",
+]
+
+
+def _default_conv(filters, x, activation="relu"):
+    x = Conv3D(filters, 3, padding="same")(x)
+    x = Conv3D(filters, 3, padding="same")(x)
+    return Activation(activation)(x)
+
+
+def waveletvit_bottleneck_factory(
+    *,
+    levels: int = 1,
+    wave: str = "haar",
+    token_dim: int = 128,
+    stride: int = 4,
+    heads: int = 4,
+    key_dim: int = 16,
+    strength: float = 0.25,
+    level_decay: float = 1.0,
+    use_orient_mix: bool = True,
+    lp_global_if_small: bool = True,
+    max_tokens_global: int = 8192,
+    qkv_bias: bool = True,
+    use_ffn: bool = True,
+    ffn_hidden_mult: float = 2.0,
+    ffn_dropout: float = 0.0,
+    kl_scale: float = 0.0,
+    sample_at_inference: bool = False,
+    tokenizer_norm: bool = True,
+    dwt_compute_dtype: str = "float32",
+    prebuild_banks: bool = True,
+    pad_to_wavelet_multiple: bool = False,
+):
+    """Create a multilevel wavelet-band transformer bottleneck callable for a 3D U-Net.
+
+    The WaveletViT block modulates raw DWT coefficient bands and the
+    assembler reconstructs a feature map with the same channel width as the
+    incoming U-Net bottleneck tensor. ``token_dim`` controls the internal
+    compressed attention width, not the output channel count.
+    """
+    tokenizer = WaveletTokenizer3D(
+        levels=int(levels),
+        wave=wave,
+        norm=bool(tokenizer_norm),
+        pad_to_wavelet_multiple=bool(pad_to_wavelet_multiple),
+        dwt_compute_dtype=dwt_compute_dtype,
+        prebuild_banks=bool(prebuild_banks),
+    )
+    block = WaveletViTBlock3D(
+        heads=int(heads),
+        key_dim=int(key_dim),
+        strength=float(strength),
+        level_decay=float(level_decay),
+        use_orient_mix=bool(use_orient_mix),
+        lp_global_if_small=bool(lp_global_if_small),
+        max_tokens_global=int(max_tokens_global),
+        qkv_bias=bool(qkv_bias),
+        use_ffn=bool(use_ffn),
+        ffn_hidden_mult=float(ffn_hidden_mult),
+        ffn_dropout=float(ffn_dropout),
+        kl_scale=float(kl_scale),
+        sample_at_inference=bool(sample_at_inference),
+        token_dim=int(token_dim),
+        stride=int(stride),
+    )
+    assembler = WaveletAssembler3D(
+        levels=int(levels),
+        wave=wave,
+        mode="feature",
+        dwt_compute_dtype=dwt_compute_dtype,
+        prebuild_banks=bool(prebuild_banks),
+    )
+
+    def bottleneck(x: tf.Tensor, training=None) -> tf.Tensor:
+        bands = tokenizer(x)
+        bands = block(bands, training=training)
+        y = assembler(bands)
+        return tf.cast(y, x.dtype) if y.dtype != x.dtype else y
+
+    return bottleneck
+
+
+def UWaveletViTNet3D(
+    input_shape=(32, 32, 32, 1),
+    config=(16, 32, 64, 128),
+    n_classes=4,
+    *,
+    token_dim: int = 128,
+    levels: int = 1,
+    wave: str = "haar",
+    stride: int = 4,
+    heads: int = 4,
+    key_dim: int = 16,
+    strength: float = 0.25,
+    level_decay: float = 1.0,
+    use_orient_mix: bool = True,
+    lp_global_if_small: bool = True,
+    max_tokens_global: int = 8192,
+    qkv_bias: bool = True,
+    use_ffn: bool = True,
+    ffn_hidden_mult: float = 2.0,
+    ffn_dropout: float = 0.0,
+    kl_scale: float = 0.0,
+    sample_at_inference: bool = False,
+    tokenizer_norm: bool = True,
+    dwt_compute_dtype: str = "float32",
+    prebuild_banks: bool = True,
+    output_kernel_regularizer=None,
+    one_hot_encode=True,
+    residual=False,
+):
+    """Build a 3D U-Net with one multilevel wavelet-band transformer module at the bottleneck."""
+    bottleneck = waveletvit_bottleneck_factory(
+        levels=levels,
+        wave=wave,
+        token_dim=token_dim,
+        stride=stride,
+        heads=heads,
+        key_dim=key_dim,
+        strength=strength,
+        level_decay=level_decay,
+        use_orient_mix=use_orient_mix,
+        lp_global_if_small=lp_global_if_small,
+        max_tokens_global=max_tokens_global,
+        qkv_bias=qkv_bias,
+        use_ffn=use_ffn,
+        ffn_hidden_mult=ffn_hidden_mult,
+        ffn_dropout=ffn_dropout,
+        kl_scale=kl_scale,
+        sample_at_inference=sample_at_inference,
+        tokenizer_norm=tokenizer_norm,
+        dwt_compute_dtype=dwt_compute_dtype,
+        prebuild_banks=prebuild_banks,
+        pad_to_wavelet_multiple=False,
+    )
+
+    inputs = Input(shape=input_shape, name="input")
+    x = inputs
+    skips = []
+    for filters in config:
+        x = _default_conv(int(filters), x)
+        skips.append(x)
+        x = MaxPool3D(pool_size=2)(x)
+
+    x = bottleneck(x)
+    if x.shape[-1] is None or int(x.shape[-1]) != int(config[-1]):
+        x = Conv3D(int(config[-1]), 1, padding="same", name="bottleneck_channel_align")(x)
+
+    for i, filters in reversed(list(enumerate(config))):
+        filters = int(filters)
+        x = Conv3DTranspose(filters, 2, strides=2, padding="same")(x)
+        x = Concatenate()([x, skips[i]])
+        x = _default_conv(filters, x)
+
+    if residual:
+        x = Concatenate()([x, inputs])
+
+    outputs = Conv3D(
+        n_classes,
+        1,
+        kernel_regularizer=output_kernel_regularizer,
+        activation="softmax" if one_hot_encode else None,
+        name="segmentation_logits",
+    )(x)
+    model_name = "CVUWaveletViTNet3D" if float(kl_scale) > 0.0 else "UWaveletViTNet3D"
+    return keras.Model(inputs, outputs, name=model_name)
+
+if __name__ == "__main__":
+    tf.get_logger().setLevel("ERROR")
+    tf.random.set_seed(0)
+    n = 32
+    model = UWaveletViTNet3D(
+        input_shape=(n, n, n, 1),
+        config=(16, 32, 64, 128),
+        n_classes=4,
+        token_dim=64,
+        levels=1,
+        wave="haar",
+        heads=2,
+        key_dim=16,
+        strength=0.5,
+        level_decay=1.0,
+    )
+    y = model(tf.random.normal((1, n, n, n, 1)))
+    model.summary()
+    print("out:", y.shape)