waveseg_model.py

"""WaveSeg (A1) model definition - a frozen, self-contained copy for this
Hugging Face Space, so app/ can be deployed on its own without the training
repo's src/ tree as a dependency.

This is a verbatim copy of the classes in the training repo's
src/models/{fba,segformer_base,waveseg}.py, unmodified, to guarantee zero
drift from the exact architecture the checkpoints were trained with. The
deployed configuration is always: HaarDWT + gating ON + deep placement
(A1 = WaveSeg, ours) - see MODEL_CARD.md for why the boundary-frequency loss
and the other ablation axes (FFT, shallow placement, no-gate) are not part
of the shipped model. Those extra axes (FFTHighPass, use_gate=False,
placement="shallow") are kept here anyway rather than stripped out, since a
verbatim copy is lower-risk than a hand-simplified one.
"""
from __future__ import annotations

from typing import Tuple

import torch
import torch.nn.functional as F
from torch import nn
from transformers import SegformerForSemanticSegmentation


class HaarDWT(nn.Module):
    """Single-level 2D Haar discrete wavelet transform via a fixed,
    non-learnable depthwise stride-2 convolution.

    Splits each input channel into 4 orthonormal sub-bands (LL, LH, HL, HH),
    each at half spatial resolution. The filters are registered as a buffer
    (not a parameter) since this is a classical transform, not a learned
    filter bank - it contributes 0 trainable parameters.
    """

    def __init__(self, in_channels: int) -> None:
        super().__init__()
        self.in_channels = in_channels

        ll = torch.tensor([[1.0, 1.0], [1.0, 1.0]])
        lh = torch.tensor([[1.0, -1.0], [1.0, -1.0]])
        hl = torch.tensor([[1.0, 1.0], [-1.0, -1.0]])
        hh = torch.tensor([[1.0, -1.0], [-1.0, 1.0]])
        filters = torch.stack([ll, lh, hl, hh], dim=0) * 0.5  # (4, 2, 2), orthonormal

        # Depthwise grouped conv: each input channel gets its own copy of the
        # 4 filters -> weight shape (4*C, 1, 2, 2), consumed with groups=C so
        # group g (= input channel g) produces output channels [4g : 4g+4].
        weight = filters.unsqueeze(1).repeat(in_channels, 1, 1, 1)
        self.register_buffer("weight", weight)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Decompose ``x`` into (LL, LH, HL, HH), each shape ``(B, C, H/2, W/2)``.

        Args:
            x: Input feature map, shape ``(B, C, H, W)``, with ``C`` equal to
                ``in_channels`` and ``H, W`` even.
        """
        out = F.conv2d(x, self.weight, stride=2, groups=self.in_channels)  # (B, 4*C, H/2, W/2)
        b, _, h, w = out.shape
        out = out.view(b, self.in_channels, 4, h, w)
        ll, lh, hl, hh = out[:, :, 0], out[:, :, 1], out[:, :, 2], out[:, :, 3]
        return ll, lh, hl, hh


class FFTHighPass(nn.Module):
    """Single-map FFT high-pass filter - the A4 ablation's alternative to the
    Haar DWT. Not used by A1 (WaveSeg, ours uses ``freq_transform="dwt"``);
    kept here only because this file is a verbatim copy of the training repo.
    """

    def __init__(self, cutoff_ratio: float = 0.25) -> None:
        super().__init__()
        if not 0.0 < cutoff_ratio < 1.0:
            raise ValueError(f"cutoff_ratio must be in (0, 1), got {cutoff_ratio}")
        self.cutoff_ratio = cutoff_ratio

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Return the high-pass-filtered map, same shape as ``x``.

        Args:
            x: Input feature map, shape ``(B, C, H, W)``.
        """
        h, w = x.shape[-2:]
        cy, cx = h // 2, w // 2
        ry = max(1, int(round(h * self.cutoff_ratio / 2)))
        rx = max(1, int(round(w * self.cutoff_ratio / 2)))

        with torch.autocast(device_type=x.device.type, enabled=False):
            spectrum = torch.fft.fftshift(torch.fft.fft2(x.float()), dim=(-2, -1))
            low_freq_mask = torch.ones(h, w, device=x.device, dtype=torch.float32)
            low_freq_mask[cy - ry : cy + ry, cx - rx : cx + rx] = 0.0
            spectrum = spectrum * low_freq_mask
            high_pass = torch.fft.ifft2(torch.fft.ifftshift(spectrum, dim=(-2, -1))).real

        return high_pass.to(x.dtype)


class FrequencyBoundaryAdapter(nn.Module):
    """FBA: gates a decoder feature map using a boundary-attention map derived
    from a high-frequency decomposition of that feature map.

    ``out = feat * (1 + alpha * attn)`` when ``use_gate=True`` (A1/WaveSeg -
    the deployed configuration). ``use_gate=False`` and ``freq_transform="fft"``
    are other ablations, not used here; kept only for a verbatim copy.
    """

    def __init__(
        self,
        in_channels: int,
        alpha_init: float = 0.1,
        use_gate: bool = True,
        freq_transform: str = "dwt",
        fft_cutoff_ratio: float = 0.25,
    ) -> None:
        super().__init__()
        if freq_transform not in ("dwt", "fft"):
            raise ValueError(f"freq_transform must be 'dwt' or 'fft', got {freq_transform!r}")
        self.freq_transform = freq_transform
        self.use_gate = use_gate

        if freq_transform == "dwt":
            self.dwt = HaarDWT(in_channels)
            self.gate = nn.Conv2d(in_channels * 3, 1, kernel_size=1)
        else:
            self.fft_highpass = FFTHighPass(cutoff_ratio=fft_cutoff_ratio)
            self.gate = nn.Conv2d(in_channels, 1, kernel_size=1)

        self.alpha = nn.Parameter(torch.tensor(float(alpha_init)))

    def forward(self, feat: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Gate a decoder feature map with a frequency-derived boundary attention map.

        Args:
            feat: Decoder feature map, shape ``(B, C, H, W)``. ``H, W`` must
                be even when ``freq_transform="dwt"``.

        Returns:
            ``(out, attn)``: ``out`` has the same shape as ``feat``; ``attn``
            is the boundary attention map, shape ``(B, 1, H, W)``, at
            ``feat``'s resolution.
        """
        h, w = feat.shape[-2:]
        if self.freq_transform == "dwt":
            _, lh, hl, hh = self.dwt(feat)
            high_freq = torch.cat([lh, hl, hh], dim=1)  # (B, 3C, H/2, W/2)
            attn = torch.sigmoid(self.gate(high_freq))  # (B, 1, H/2, W/2)
            attn = F.interpolate(attn, size=(h, w), mode="bilinear", align_corners=False)
        else:
            high_freq = self.fft_highpass(feat)  # (B, C, H, W), already full resolution
            attn = torch.sigmoid(self.gate(high_freq))  # (B, 1, H, W)

        out = feat * (1.0 + self.alpha * attn) if self.use_gate else feat
        return out, attn


class SegFormerBaseline(nn.Module):
    """SegFormer-B0 wrapped for binary segmentation at full input resolution.

    HuggingFace's SegFormer decode head predicts at 1/4 input resolution, so
    this wrapper bilinearly upsamples logits back to the input size.
    """

    def __init__(self, pretrained_name: str = "nvidia/mit-b0", num_classes: int = 1) -> None:
        super().__init__()
        # token=False: this is a public checkpoint that never needs auth - a
        # stale cached HF token can otherwise turn a no-auth-needed request
        # into a 401 (see the training repo's CLAUDE.md).
        self.model = SegformerForSemanticSegmentation.from_pretrained(
            pretrained_name,
            num_labels=num_classes,
            ignore_mismatched_sizes=True,
            token=False,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Run the encoder+decoder and upsample logits to the input resolution."""
        h, w = x.shape[-2:]
        logits = self.model(pixel_values=x).logits
        return F.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False)


class WaveSeg(nn.Module):
    """SegFormer-B0 baseline with an FBA gating a decoder/encoder feature map.

    Deployed configuration (A1, WaveSeg ours): ``use_gate=True``,
    ``freq_transform="dwt"``, ``placement="deep"`` - see
    ``build_waveseg_a1()`` below. HuggingFace's ``SegformerDecodeHead``
    always ends with ``self.classifier``, a plain 1x1 conv applied last; we
    swap it for ``nn.Identity``, gate the pre-classifier features with the
    FBA, then apply the real classifier ourselves.
    """

    def __init__(
        self,
        pretrained_name: str = "nvidia/mit-b0",
        num_classes: int = 1,
        alpha_init: float = 0.1,
        use_gate: bool = True,
        freq_transform: str = "dwt",
        fft_cutoff_ratio: float = 0.25,
        placement: str = "deep",
    ) -> None:
        super().__init__()
        if placement not in ("deep", "shallow"):
            raise ValueError(f"placement must be 'deep' or 'shallow', got {placement!r}")
        self.placement = placement
        self.baseline = SegFormerBaseline(pretrained_name=pretrained_name, num_classes=num_classes)

        if placement == "deep":
            decode_head = self.baseline.model.decode_head
            fba_channels = decode_head.classifier.in_channels
            self.classifier = decode_head.classifier
            decode_head.classifier = nn.Identity()
        else:  # shallow: not used by A1, kept for a verbatim copy
            fba_channels = self.baseline.model.config.hidden_sizes[0]

        self.fba = FrequencyBoundaryAdapter(
            in_channels=fba_channels,
            alpha_init=alpha_init,
            use_gate=use_gate,
            freq_transform=freq_transform,
            fft_cutoff_ratio=fft_cutoff_ratio,
        )

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Run the encoder+decoder, gate with the FBA, then classify.

        Args:
            x: Input batch, shape ``(B, 3, H, W)``.

        Returns:
            ``(logits, attn)``: segmentation logits ``(B, num_classes, H, W)``
            and the FBA's boundary attention map ``(B, 1, H, W)``, both
            upsampled to input resolution.
        """
        h, w = x.shape[-2:]

        if self.placement == "deep":
            decoder_feat = self.baseline.model(pixel_values=x).logits  # pre-classifier features
            gated_feat, attn = self.fba(decoder_feat)
            logits = self.classifier(gated_feat)
        else:  # shallow
            outputs = self.baseline.model.segformer(pixel_values=x, output_hidden_states=True, return_dict=True)
            hidden_states = list(outputs.hidden_states)
            gated_stage0, attn = self.fba(hidden_states[0])
            hidden_states[0] = gated_stage0
            logits = self.baseline.model.decode_head(tuple(hidden_states))

        logits = F.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False)
        attn = F.interpolate(attn, size=(h, w), mode="bilinear", align_corners=False)
        return logits, attn


def build_waveseg_a1(pretrained_name: str = "nvidia/mit-b0", num_classes: int = 1) -> WaveSeg:
    """Build WaveSeg with the locked A1 configuration: FBA gating (DWT, deep
    placement), no boundary loss involved (that's a training-time concern -
    architecturally A1 and A3 are identical; only the loss used to train
    them differs, which is why this one class serves the deployed model).
    """
    return WaveSeg(
        pretrained_name=pretrained_name,
        num_classes=num_classes,
        alpha_init=0.1,
        use_gate=True,
        freq_transform="dwt",
        placement="deep",
    )