File size: 4,555 Bytes
a78ad5d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | import torch
import torch.nn as nn
import torch.nn.functional as F
from .backbones import create_backbone
from .blocks import ConvBlock
class UNet3Plus(nn.Module):
"""UNet 3+ with full-scale skip connections.
Every decoder node aggregates features from ALL encoder levels (not just
the matching scale), giving each node simultaneous access to fine-grained
detail and deep semantic context. Each incoming stream is projected to
`inter_ch` channels before concatenation so total input channels are
predictable regardless of backbone channel widths.
Reference: Huang et al., "UNet 3+: A Full-Scale Connected UNet for
Medical Image Segmentation", ICASSP 2020.
"""
def __init__(
self,
backbone_name: str = "default",
in_channels: int = 3,
out_channels: int = 1,
inter_ch: int = 64,
**_kw,
):
super().__init__()
self.backbone = create_backbone(backbone_name, in_channels)
all_ch = list(self.backbone.out_channels) # [skip0, ..., skip_{D-1}, bottleneck]
D = len(all_ch) - 1 # number of skip levels = decoder levels
self._D = D
self._inter_ch = inter_ch
self._dec_ch = (D + 1) * inter_ch # fixed output channels for every decoder node
# --- projections: each incoming stream is independently projected to inter_ch ---
# enc_projs[k][j]: encoder/bottleneck level j → inter_ch, for use at decoder k
self.enc_projs = nn.ModuleList([
nn.ModuleList([
nn.Sequential(
nn.Conv2d(all_ch[j], inter_ch, 1, bias=False),
nn.BatchNorm2d(inter_ch),
nn.ReLU(inplace=True),
)
for j in range(D + 1)
])
for k in range(D)
])
# dec_projs[k][i]: prior decoder level i → inter_ch, for use at decoder k (i < k)
# All prior decoder outputs have self._dec_ch channels.
self.dec_projs = nn.ModuleList([
nn.ModuleList([
nn.Sequential(
nn.Conv2d(self._dec_ch, inter_ch, 1, bias=False),
nn.BatchNorm2d(inter_ch),
nn.ReLU(inplace=True),
)
for _ in range(k) # one projection per prior decoder level
])
for k in range(D)
])
# Fusion conv at each decoder level: (D+1+k)*inter_ch → dec_ch
self.dec_convs = nn.ModuleList([
ConvBlock((D + 1 + k) * inter_ch, self._dec_ch)
for k in range(D)
])
self.head = nn.Conv2d(self._dec_ch, out_channels, kernel_size=1)
def forward(
self,
x: torch.Tensor | None = None,
pixel_values: torch.Tensor | None = None,
**_kw,
) -> torch.Tensor:
if pixel_values is not None:
x = pixel_values
input_size = x.shape[2:]
skips, bottleneck = self.backbone(x)
D = self._D
# enc_feats: shallowest encoder first, bottleneck last (D+1 tensors)
enc_feats = list(skips) + [bottleneck]
dec_outputs: list[torch.Tensor] = []
# Build decoder from deepest (k=0) to shallowest (k=D-1).
# Decoder k targets the same spatial resolution as skips[D-1-k].
for k in range(D):
target = skips[D - 1 - k].shape[2:]
parts: list[torch.Tensor] = []
# All encoder / bottleneck streams
for j, feat in enumerate(enc_feats):
feat = _resize(feat, target)
parts.append(self.enc_projs[k][j](feat))
# All prior decoder streams (deeper → current scale = always upsample)
for i, df in enumerate(dec_outputs):
feat = F.interpolate(df, size=target, mode="bilinear", align_corners=False)
parts.append(self.dec_projs[k][i](feat))
dec_outputs.append(self.dec_convs[k](torch.cat(parts, dim=1)))
out = self.head(dec_outputs[-1])
if out.shape[2:] != input_size:
out = F.interpolate(out, size=input_size, mode="bilinear", align_corners=False)
return out
def _resize(feat: torch.Tensor, target: tuple[int, int] | torch.Size) -> torch.Tensor:
if feat.shape[2:] == torch.Size(target):
return feat
if feat.shape[2] > target[0]:
return F.adaptive_max_pool2d(feat, target)
return F.interpolate(feat, size=target, mode="bilinear", align_corners=False)
|