SAM3UNeXTV6.py

"""
SAM3-UNeXT V6: Complete SAM3 Backbone + DINOv2 with DRFU Dual-Branch Upsampling

Improvements over V2:
1. DRFU (Dual-Branch Receptive Field Upsampling) decoder
   - Bilinear Upsample branch + ConvTranspose branch
   - Gaussian smoothing for feature refinement
   - Better feature preservation during upsampling
2. No Deep Supervision (removed based on ablation study)

Architecture:
- SAM3 Complete Backbone (frozen) with FPN -> 4 levels [256, 256, 256, 256]
- Adapter (bottleneck=64) on ViT blocks
- DINOv2 ViT-L Encoder (frozen) -> 1024 channels
- DRFU Dual-Branch Upsampling Decoder
- Single output head (deep supervision removed for better performance)

Expected Performance: mIoU ~0.719 (based on ablation study)
"""

import sys
import math
from pathlib import Path
from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
from timm.models.layers import trunc_normal_

PROJECT_ROOT = Path(__file__).resolve().parent
VENDORED_SAM3 = PROJECT_ROOT / "sam3"

if VENDORED_SAM3.exists():
    candidate_str = VENDORED_SAM3.as_posix()
    if candidate_str not in sys.path:
        sys.path.insert(0, candidate_str)

from sam3.model_builder import build_sam3_image_model



class Adapter(nn.Module):
    """
    Adapter module for parameter-efficient fine-tuning.
    Same as V2: bottleneck=64 for sufficient capacity.
    """
    def __init__(self, blk, bottleneck: int = 64):
        super().__init__()
        self.block = blk
        dim = blk.attn.qkv.in_features
        self.adapter = nn.Sequential(
            nn.Linear(dim, bottleneck),
            nn.GELU(),
            nn.Linear(bottleneck, dim),
            nn.GELU()
        )
        self._init_weights()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        adapted = x + self.adapter(x)
        return self.block(adapted)

    def _init_weights(self):
        for m in self.adapter.modules():
            if isinstance(m, nn.Linear):
                trunc_normal_(m.weight, std=.02)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)


# ==================== Gaussian Filter (from LEGNet) ====================

class GaussianFilter(nn.Module):
    """Gaussian smoothing filter for feature refinement."""
    def __init__(self, channels: int, kernel_size: int = 5, sigma: float = 1.0):
        super().__init__()
        self.channels = channels

        # Generate Gaussian kernel
        kernel = self._gaussian_kernel(kernel_size, sigma)
        kernel = kernel.repeat(channels, 1, 1, 1)

        self.register_buffer('weight', kernel)
        self.padding = kernel_size // 2

    def _gaussian_kernel(self, size: int, sigma: float) -> torch.Tensor:
        coords = torch.arange(size, dtype=torch.float32) - size // 2
        g = torch.exp(-(coords ** 2) / (2 * sigma ** 2))
        g = g / g.sum()
        kernel = g.unsqueeze(0) * g.unsqueeze(1)
        return kernel.unsqueeze(0).unsqueeze(0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.conv2d(x, self.weight, padding=self.padding, groups=self.channels)


# ==================== DRFU: Dual-Branch Receptive Field Upsampling (NEW) ====================

class DRFU(nn.Module):
    """
    Dual-Branch Receptive Field Upsampling (inspired by LEGNet's DRFD).

    Instead of downsampling, this module performs upsampling with dual branches:
    - Branch 1: Bilinear interpolation (preserves spatial smoothness)
    - Branch 2: ConvTranspose (learnable upsampling)

    Features:
    - Gaussian smoothing for feature refinement
    - Dual-branch fusion for better feature preservation
    """
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()

        # Branch 1: Bilinear upsample + Conv
        self.up_bilinear = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.conv_bilinear = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.GELU()
        )

        # Branch 2: ConvTranspose (learnable upsample)
        self.up_convt = nn.Sequential(
            nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.GELU()
        )

        # Gaussian smoothing for refinement
        self.gaussian = GaussianFilter(out_channels, kernel_size=5, sigma=1.0)
        self.norm_g = nn.BatchNorm2d(out_channels)

        # Fusion: concat two branches then reduce
        self.fusion = nn.Sequential(
            nn.Conv2d(out_channels * 2, out_channels, kernel_size=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.GELU()
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Branch 1: Bilinear
        x_bilinear = self.conv_bilinear(self.up_bilinear(x))

        # Branch 2: ConvTranspose
        x_convt = self.up_convt(x)

        # Ensure same size (in case of odd dimensions)
        if x_bilinear.shape[-2:] != x_convt.shape[-2:]:
            x_convt = F.interpolate(x_convt, size=x_bilinear.shape[-2:], mode='bilinear', align_corners=False)

        # Gaussian smoothing on bilinear branch
        gaussian = self.gaussian(x_bilinear)
        x_bilinear = self.norm_g(x_bilinear + gaussian)

        # Fusion
        x = torch.cat([x_bilinear, x_convt], dim=1)
        x = self.fusion(x)

        return x


# ==================== Decoder Blocks ====================

class LightweightBlock(nn.Module):
    """
    Lightweight decoder block from SAM3-UNet paper.
    Design: Bottleneck + Feature Splitting + DWConv
    """
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()

        mid_channels = max(in_channels // 4, 32)

        self.squeeze = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=1, bias=False),
            nn.BatchNorm2d(mid_channels),
            nn.GELU()
        )

        part_channels = mid_channels // 2

        self.dwconv1 = nn.Sequential(
            nn.Conv2d(part_channels, part_channels, kernel_size=3, padding=1,
                      groups=part_channels, bias=False),
            nn.BatchNorm2d(part_channels),
            nn.GELU()
        )

        self.dwconv2 = nn.Sequential(
            nn.Conv2d(part_channels, part_channels, kernel_size=3, padding=1,
                      groups=part_channels, bias=False),
            nn.BatchNorm2d(part_channels),
            nn.GELU()
        )

        self.expand = nn.Sequential(
            nn.Conv2d(mid_channels * 2, out_channels, kernel_size=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.GELU()
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.squeeze(x)
        part1, part2 = torch.chunk(x, 2, dim=1)
        part3 = self.dwconv1(part2)
        part4 = self.dwconv2(part3)
        x = torch.cat([part1, part2, part3, part4], dim=1)
        return self.expand(x)


class UpDRFU(nn.Module):
    """
    Upsampling block using DRFU (Dual-Branch Receptive Field Upsampling).
    Replaces the simple bilinear upsample in V2.
    """
    def __init__(self, in_channels: int, skip_channels: int, out_channels: int):
        super().__init__()

        # DRFU for upsampling (replaces simple bilinear)
        self.drfu = DRFU(in_channels, in_channels)

        # After concat with skip connection
        self.conv = LightweightBlock(in_channels + skip_channels, out_channels)

    def forward(self, x: torch.Tensor, skip: Optional[torch.Tensor] = None) -> torch.Tensor:
        # DRFU upsample
        x = self.drfu(x)

        if skip is not None:
            # Handle size mismatch
            if x.shape[-2:] != skip.shape[-2:]:
                diffY = skip.size(2) - x.size(2)
                diffX = skip.size(3) - x.size(3)
                x = F.pad(x, [diffX // 2, diffX - diffX // 2,
                              diffY // 2, diffY - diffY // 2])
            x = torch.cat([x, skip], dim=1)

        return self.conv(x)


# ==================== SAM3 Backbone Wrapper ====================

class SAM3BackboneEncoder(nn.Module):
    """Wrapper for SAM3 backbone to extract FPN features."""

    def __init__(self, backbone: nn.Module):
        super().__init__()
        self.backbone = backbone

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
        backbone_out = self.backbone.forward_image(x)
        features = backbone_out.get("backbone_fpn")

        if not isinstance(features, (list, tuple)) or len(features) < 4:
            raise RuntimeError("SAM3 backbone did not return enough feature maps")

        return tuple(features[:4])


# ==================== Main Model ====================

def _resolve_default_checkpoint():
    candidates = [
        VENDORED_SAM3 / "sam3.pt",
        VENDORED_SAM3 / "checkpoints" / "sam3.pt",
    ]
    for candidate in candidates:
        if candidate and candidate.exists():
            return candidate.as_posix()
    return None


class SAM3UNeXTV6(nn.Module):
    """
    SAM3-UNeXT V6: Complete SAM3 + DINOv2 with DRFU Dual-Branch Upsampling

    Improvements over V2:
    - DRFU (Dual-Branch Receptive Field Upsampling) in decoder
      - Bilinear branch + ConvTranspose branch
      - Gaussian smoothing for refinement
      - Better feature preservation
    - No Deep Supervision (removed based on ablation study showing better performance)

    Configuration:
    - Adapter bottleneck=64
    - Complete SAM3 FPN features
    - Single output head

    Expected Performance: mIoU ~0.719

    Args:
        checkpoint_path: Path to SAM3 checkpoint
        dinov2_path: Path to DINOv2 checkpoint (optional)
        device: Device to use
        num_classes: Number of output classes
        adapter_bottleneck: Adapter bottleneck dimension (default: 64)
    """
    def __init__(
        self,
        checkpoint_path: Optional[str] = None,
        dinov2_path: Optional[str] = None,
        device: Optional[str] = None,
        num_classes: int = 1,
        adapter_bottleneck: int = 64
    ):
        super().__init__()

        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        if checkpoint_path is None:
            checkpoint_path = _resolve_default_checkpoint()
        if checkpoint_path is None:
            raise ValueError(
                "SAM3 checkpoint not found. Provide checkpoint_path or place sam3.pt under ./sam3/."
            )

        # ===== SAM3 Complete Backbone =====
        build_kwargs = {
            "enable_segmentation": False,
            "enable_inst_interactivity": False,
            "device": device,
            "eval_mode": True,
            "load_from_HF": False,
            "checkpoint_path": checkpoint_path,
        }
        model = build_sam3_image_model(**build_kwargs)
        backbone = model.backbone
        backbone.scalp = 0  # Get all FPN levels
        del model

        # Freeze SAM3 backbone
        for param in backbone.parameters():
            param.requires_grad = False

        # Insert Adapters (same as V2, bottleneck=64)
        trunk = backbone.vision_backbone.trunk
        adapted_blocks = []
        for block in trunk.blocks:
            adapted_block = Adapter(block, bottleneck=adapter_bottleneck).to(device)
            adapted_blocks.append(adapted_block)
        trunk.blocks = nn.ModuleList(adapted_blocks)

        self.sam3_encoder = SAM3BackboneEncoder(backbone)

        # ===== DINOv2 Encoder =====
        if dinov2_path:
            self.dino = timm.create_model(
                'vit_large_patch14_dinov2',
                features_only=True,
                img_size=(448, 448),
                pretrained=True,
                pretrained_cfg_overlay=dict(file=dinov2_path)
            )
        else:
            self.dino = timm.create_model(
                'vit_large_patch14_dinov2',
                features_only=True,
                img_size=(448, 448),
                pretrained=True
            )

        # Freeze DINOv2
        for param in self.dino.parameters():
            param.requires_grad = False

        # ===== Feature Alignment (same as V2) =====
        # DINOv2: 1024 -> 256 to match SAM3 FPN channels
        self.align1 = nn.Sequential(
            nn.Conv2d(1024, 256, kernel_size=1, bias=False),
            nn.BatchNorm2d(256),
            nn.GELU()
        )
        self.align2 = nn.Sequential(
            nn.Conv2d(1024, 256, kernel_size=1, bias=False),
            nn.BatchNorm2d(256),
            nn.GELU()
        )
        self.align3 = nn.Sequential(
            nn.Conv2d(1024, 256, kernel_size=1, bias=False),
            nn.BatchNorm2d(256),
            nn.GELU()
        )
        self.align4 = nn.Sequential(
            nn.Conv2d(1024, 256, kernel_size=1, bias=False),
            nn.BatchNorm2d(256),
            nn.GELU()
        )

        # ===== Feature Reduction (same as V2) =====
        # SAM3 (256) + DINOv2 aligned (256) = 512 -> 128
        self.reduce1 = nn.Sequential(
            nn.Conv2d(512, 128, kernel_size=1, bias=False),
            nn.BatchNorm2d(128),
            nn.GELU()
        )
        self.reduce2 = nn.Sequential(
            nn.Conv2d(512, 128, kernel_size=1, bias=False),
            nn.BatchNorm2d(128),
            nn.GELU()
        )
        self.reduce3 = nn.Sequential(
            nn.Conv2d(512, 128, kernel_size=1, bias=False),
            nn.BatchNorm2d(128),
            nn.GELU()
        )
        self.reduce4 = nn.Sequential(
            nn.Conv2d(512, 128, kernel_size=1, bias=False),
            nn.BatchNorm2d(128),
            nn.GELU()
        )

        # ===== DRFU Decoder (NEW in V6) =====
        # x4 -> upsample (no skip)
        self.up4 = UpDRFU(128, 0, 128)
        # concat with x3
        self.up3 = UpDRFU(128, 128, 128)
        # concat with x2
        self.up2 = UpDRFU(128, 128, 128)
        # concat with x1
        self.up1 = UpDRFU(128, 128, 128)

        # Final upsample to original resolution
        self.final_up = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            LightweightBlock(128, 64)
        )

        # ===== Output Head =====
        self.head = nn.Conv2d(64, num_classes, kernel_size=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_size = x.shape[-2:]

        # SAM3 requires input size to be multiple of 14
        h, w = x.shape[-2:]
        sam3_size = 1008

        if h != sam3_size or w != sam3_size:
            x_sam3 = F.interpolate(x, size=(sam3_size, sam3_size),
                                    mode='bilinear', align_corners=False)
        else:
            x_sam3 = x

        # ===== SAM3 Encoding (Real FPN features) =====
        x1_s, x2_s, x3_s, x4_s = self.sam3_encoder(x_sam3)

        # ===== DINOv2 Encoding =====
        x_dino = F.interpolate(x, size=(448, 448), mode='bilinear', align_corners=False)
        x_d = self.dino(x_dino)[-1]  # [B, 1024, 32, 32]

        # ===== Align DINOv2 features to SAM3 spatial sizes =====
        x1_d = F.interpolate(self.align1(x_d), size=x1_s.shape[-2:], mode='bilinear', align_corners=False)
        x2_d = F.interpolate(self.align2(x_d), size=x2_s.shape[-2:], mode='bilinear', align_corners=False)
        x3_d = F.interpolate(self.align3(x_d), size=x3_s.shape[-2:], mode='bilinear', align_corners=False)
        x4_d = F.interpolate(self.align4(x_d), size=x4_s.shape[-2:], mode='bilinear', align_corners=False)

        # ===== Feature Fusion (same as V2) =====
        x1 = self.reduce1(torch.cat([x1_s, x1_d], dim=1))
        x2 = self.reduce2(torch.cat([x2_s, x2_d], dim=1))
        x3 = self.reduce3(torch.cat([x3_s, x3_d], dim=1))
        x4 = self.reduce4(torch.cat([x4_s, x4_d], dim=1))

        # ===== DRFU Decoder =====
        d4 = self.up4(x4, None)  # No skip for deepest level
        d3 = self.up3(d4, x3)
        d2 = self.up2(d3, x2)
        d1 = self.up1(d2, x1)
        d_final = self.final_up(d1)

        # ===== Output =====
        out = self.head(d_final)
        out = F.interpolate(out, size=orig_size, mode='bilinear', align_corners=False)

        return out


# ==================== Test ====================

if __name__ == "__main__":
    checkpoint = _resolve_default_checkpoint()
    if checkpoint is None:
        raise SystemExit("Please place sam3.pt under ./sam3/ (or ./sam3/checkpoints/) to run this demo.")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Create model
    model = SAM3UNeXTV6(
        checkpoint_path=checkpoint,
        device=str(device),
        num_classes=1,
        adapter_bottleneck=64
    ).to(device)

    # Count parameters
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"\nTotal parameters: {total_params / 1e6:.2f}M")
    print(f"Trainable parameters: {trainable_params / 1e6:.2f}M")
    print(f"Frozen parameters: {(total_params - trainable_params) / 1e6:.2f}M")

    # Test forward pass
    print("\n--- Forward Pass Test ---")
    model.train()
    x = torch.randn(1, 3, 1024, 1024, device=device)

    with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
        out = model(x)

    print(f"Input shape: {x.shape}")
    print(f"Output shape: {out.shape}")

    # Test inference mode
    print("\n--- Inference Mode ---")
    model.eval()
    with torch.no_grad():
        out = model(x)
        print(f"Input shape: {x.shape}")
        print(f"Output shape: {out.shape}")

    # Test with different input size
    print("\n--- Testing with 512x512 input ---")
    with torch.no_grad():
        x_small = torch.randn(1, 3, 512, 512, device=device)
        out_small = model(x_small)
        print(f"Input shape: {x_small.shape}")
        print(f"Output shape: {out_small.shape}")

    print("\n✓ All tests passed!")