File size: 2,957 Bytes
b8fae22 | 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 | """Swin-Unet wrapper.
Instantiates SwinTransformerSys from sota/Swin-Unet directly (bypassing the repo's
yacs config + .npz/.h5 Synapse pipeline). Handles grayscale by repeating 1->3.
Constraint: the windowed attention requires img_size = 224 (patches resolutions
56/28/14/7 are each divisible by window_size=7). Other sizes will assert; we keep
img_size=224 for this backbone.
"""
from __future__ import annotations
import os
import sys
import torch
import torch.nn as nn
_REPO = os.path.join(os.path.dirname(__file__), "..", "..", "sota", "Swin-Unet")
_REPO = os.path.abspath(_REPO)
def _ensure_path():
if _REPO not in sys.path:
sys.path.insert(0, _REPO)
class SwinUnetWrapper(nn.Module):
def __init__(self, in_channels: int, num_classes: int, img_size: int = 224,
pretrained_ckpt: str = ""):
super().__init__()
_ensure_path()
from networks.swin_transformer_unet_skip_expand_decoder_sys import SwinTransformerSys
if img_size != 224:
raise ValueError("Swin-Unet backbone requires img_size=224.")
self.net = SwinTransformerSys(
img_size=img_size, patch_size=4, in_chans=3, num_classes=num_classes,
embed_dim=96, depths=[2, 2, 2, 2], num_heads=[3, 6, 12, 24],
window_size=7, mlp_ratio=4.0, qkv_bias=True, drop_path_rate=0.1,
ape=False, patch_norm=True, use_checkpoint=False,
)
if pretrained_ckpt and os.path.isfile(pretrained_ckpt):
self._load_pretrained(pretrained_ckpt)
def _load_pretrained(self, path):
"""Port of Swin-Unet's load_from: load the ImageNet Swin-T encoder AND
mirror its `layers.X` weights into the decoder `layers_up.(3-X)` (the
scheme the paper uses to initialize the symmetric decoder)."""
import copy
ckpt = torch.load(path, map_location="cpu", weights_only=False)
if "model" not in ckpt:
self.net.load_state_dict(ckpt, strict=False)
return
pretrained = ckpt["model"]
model_dict = self.net.state_dict()
full = copy.deepcopy(pretrained)
for k, v in pretrained.items():
if "layers." in k:
n = 3 - int(k[7:8])
full["layers_up." + str(n) + k[8:]] = v
for k in list(full.keys()):
if k in model_dict and full[k].shape != model_dict[k].shape:
del full[k]
msg = self.net.load_state_dict(full, strict=False)
print(f"[swinunet] loaded pretrained {path}: "
f"missing={len(msg.missing_keys)} unexpected={len(msg.unexpected_keys)}")
def forward(self, x):
if x.size(1) == 1:
x = x.repeat(1, 3, 1, 1)
return self.net(x)
def build_swinunet(in_channels: int, num_classes: int, img_size: int = 224,
pretrained_ckpt: str = "", **_):
return SwinUnetWrapper(in_channels, num_classes, img_size, pretrained_ckpt)
|