implemented dinounet

src/model/dinounet.py

@@ -0,0 +1,201 @@
+from dataclasses import dataclass
+from typing import List
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import DINOv3ViTModel
+
+@dataclass
+class DinoUNETConfig:
+    model_name: str = "facebook/dinov2-small"
+    num_classes: int = 2
+    features_per_stage: List[int] = [32, 64, 128, 256]
+    n_conv_per_stage_decoder: List[int]  = [2, 2, 2]
+    deep_supervision: bool = False
+    rank: int = 256
+
+class SqueezeExcitation(nn.Module):
+    def __init__(self, channels: int, reduction: int=16):
+        super().__init__()
+
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Conv2d(channels, max(1, channels // reduction), 1),
+            nn.ReLU(True),
+            nn.Conv2d(max(1, channels // reduction), channels, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x: torch.Tensor):
+        return x * self.fc(self.pool(x))
+
+class DepthwiseSeparableConv(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, kernel_size: int=3, stride: int=1, padding: int=1):
+        super().__init__()
+
+        self.depthwise = nn.Conv2d(in_ch, in_ch, kernel_size, stride, padding, groups=in_ch)
+        self.pointwise = nn.Conv2d(in_ch, out_ch, 1)
+        self.bn = nn.BatchNorm2d(out_ch)
+        self.act = nn.ReLU(True)
+
+    def forward(self, x: torch.Tensor):
+        return self.act(self.bn(self.pointwise(self.depthwise(x))))
+
+class LearnableUpsample(nn.Module):
+    def __init__(self, channels: int):
+        super().__init__()
+
+        self.up2 = nn.ConvTranspose2d(channels, channels, 2, 2)
+
+    def forward(self, x: torch.Tensor, target_size: tuple[int, int]):
+        h, w = x.shape[2:]
+        out = x
+        while h * 2 <= target_size[0] and w * 2 <= target_size[1]:
+            out = self.up2(out)
+            h, w = out.shape[2:]
+        if (h, w) != target_size:
+            out = F.interpolate(out, target_size, mode='bilinear', align_corners=False)
+        return out
+
+class FAPM(nn.Module):
+    def __init__(self, in_ch: int, rank: int, out_ch_list: list[int]):
+        super().__init__()
+
+        self.shared_basis = nn.Conv2d(in_ch, rank, 1)
+        self.specific_bases = nn.ModuleList([nn.Conv2d(in_ch, rank, 1) for _ in out_ch_list])
+        self.film_generators = nn.ModuleList([nn.Conv2d(rank, rank * 2, 1) for _ in out_ch_list])
+        self.refinement_blocks = nn.ModuleList()
+        self.shortcuts = nn.ModuleList()
+
+        for oc in out_ch_list:
+            self.refinement_blocks.append(nn.Sequential(
+                nn.Conv2d(rank, oc, 1),
+                nn.BatchNorm2d(oc),
+                nn.ReLU(True),
+                DepthwiseSeparableConv(oc, oc),
+                nn.Conv2d(oc, oc, 1),
+                SqueezeExcitation(oc)
+            ))
+            self.shortcuts.append(nn.Conv2d(rank, oc, 1) if rank != oc else nn.Identity())
+
+    def forward(self, x_list: list[torch.Tensor]):
+        out = []
+        for i, x in enumerate(x_list):
+            z_shared = self.shared_basis(x)
+            z_specific = self.specific_bases[i](x)
+            gamma, beta = torch.chunk(self.film_generators[i](z_shared), 2, dim=1)
+            z_modulated = gamma * z_specific + beta
+            refined = self.refinement_blocks[i](z_modulated)
+            shortcut = self.shortcuts[i](z_modulated)
+            out.append(refined + shortcut)
+        return out
+
+class DINOv3Encoder(nn.Module):
+    def __init__(self, config: DinoUNETConfig):
+        super().__init__()
+
+        self.output_channels = config.features_per_stage
+        self.strides = [(2, 2)] * len(config.features_per_stage)
+
+        self.backbone = DINOv3ViTModel.from_pretrained(config.model_name)
+        self.fapm = FAPM(self.backbone.config.hidden_size, config.rank, config.features_per_stage)
+        self.ups = nn.ModuleList([LearnableUpsample(oc) for oc in config.features_per_stage])
+
+
+    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
+        B, C, H, W = x.shape
+
+        if C == 1:
+            x = x.repeat(1, 3, 1, 1)
+        elif C != 3:
+            x = x[:, :3] if C > 3 else F.pad(x, (0, 0, 0, 0, 0, 3 - C))[:, :3]
+
+        outputs = self.backbone(x, output_hidden_states=True, return_dict=True)
+        hidden_states = outputs.hidden_states
+
+        h, w = H // self.backbone.config.patch_size, W // self.backbone.config.patch_size
+        features = []
+        indices = [3, 6, 9, 12] if len(hidden_states) > 12 else [2, 4, 6, -1]
+
+        for idx in indices:
+            feat = hidden_states[idx][:, 1:].transpose(1, 2).reshape(B, -1, h, w)
+            features.append(feat)
+
+        features = self.fapm(features)
+        skips = []
+        for i, feat in enumerate(features):
+            target_size = (H // (2 ** i), W // (2 ** i))
+            skips.append(self.ups[i](feat, target_size))
+
+        return skips
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, kernel_size: int=3, stride: int=1):
+        super().__init__()
+
+        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, stride, kernel_size // 2)
+        self.bn = nn.BatchNorm2d(out_ch)
+        self.act = nn.ReLU(True)
+
+    def forward(self, x):
+        return self.act(self.bn(self.conv(x)))
+
+class StackedConvBlocks(nn.Module):
+    def __init__(self, n_blocks: int, in_ch: int, out_ch: int, kernel_size: int=3):
+        super().__init__()
+
+        self.blocks = nn.Sequential(*(
+            [ConvBlock(in_ch, out_ch, kernel_size)] +
+            [ConvBlock(out_ch, out_ch, kernel_size) for _ in range(n_blocks - 1)]
+        ))
+
+    def forward(self, x):
+        return self.blocks(x)
+
+class UNetDecoder(nn.Module):
+    def __init__(self, encoder: DINOv3Encoder, config: DinoUNETConfig):
+        super().__init__()
+        self.deep_supervision = config.deep_supervision
+        self.encoder = encoder
+
+        self.stages = nn.ModuleList()
+        self.transpconvs = nn.ModuleList()
+        self.seg_layers = nn.ModuleList()
+
+        for s in range(1, len(encoder.output_channels)):
+            in_below = encoder.output_channels[-s]
+            in_skip = encoder.output_channels[-(s + 1)]
+            stride = encoder.strides[-s]
+
+            self.transpconvs.append(nn.ConvTranspose2d(in_below, in_skip, stride, stride))
+            self.stages.append(StackedConvBlocks(
+                config.n_conv_per_stage_decoder[s-1],
+                2 * in_skip,
+                in_skip
+            ))
+            self.seg_layers.append(nn.Conv2d(in_skip, config.num_classes, 1))
+
+    def forward(self, skips: list[torch.Tensor]) -> torch.Tensor | list[torch.Tensor]: # only list if doing deep supervision
+        lres_input = skips[-1]
+        seg_outputs = []
+
+        for s in range(len(self.stages)):
+            x = self.transpconvs[s](lres_input)
+            x = torch.cat((x, skips[-(s+2)]), 1)
+            x = self.stages[s](x)
+            if self.deep_supervision or s == len(self.stages) - 1:
+                seg_outputs.append(self.seg_layers[s](x))
+            lres_input = x
+
+        seg_outputs = seg_outputs[::-1]
+        return seg_outputs if self.deep_supervision else seg_outputs[0]
+
+class DinoUNet(nn.Module):
+    def __init__(self, config: DinoUNETConfig):
+        super().__init__()
+        self.encoder = DINOv3Encoder(config)
+        self.decoder = UNetDecoder(self.encoder, config)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        skips = self.encoder(x)
+        return self.decoder(skips)