[models] Add Little W-Net architecture

bob/ip/binseg/models/lwnet.py

+#!/usr/bin/env python
+# coding=utf-8
+
+"""Little W-Net
+
+Code was originally developed by Adrian Galdran
+(https://github.com/agaldran/lwnet), loosely inspired on
+https://github.com/jvanvugt/pytorch-unet
+
+It is based on two simple U-Nets with 3 layers concatenated to each other.  The
+first U-Net produces a segmentation map that is used by the second to better
+guide segmentation.
+"""
+
+
+import torch
+import torch.nn
+
+
+def _conv1x1(in_planes, out_planes, stride=1):
+    return torch.nn.Conv2d(
+        in_planes, out_planes, kernel_size=1, stride=stride, bias=False
+    )
+
+
+class ConvBlock(torch.nn.Module):
+    def __init__(self, in_c, out_c, k_sz=3, shortcut=False, pool=True):
+        """
+        pool_mode can be False (no pooling) or True ('maxpool')
+        """
+
+        super(ConvBlock, self).__init__()
+        if shortcut == True:
+            self.shortcut = torch.nn.Sequential(
+                _conv1x1(in_c, out_c), torch.nn.BatchNorm2d(out_c)
+            )
+        else:
+            self.shortcut = False
+        pad = (k_sz - 1) // 2
+
+        block = []
+        if pool:
+            self.pool = torch.nn.MaxPool2d(kernel_size=2)
+        else:
+            self.pool = False
+
+        block.append(
+            torch.nn.Conv2d(in_c, out_c, kernel_size=k_sz, padding=pad)
+        )
+        block.append(torch.nn.ReLU())
+        block.append(torch.nn.BatchNorm2d(out_c))
+
+        block.append(
+            torch.nn.Conv2d(out_c, out_c, kernel_size=k_sz, padding=pad)
+        )
+        block.append(torch.nn.ReLU())
+        block.append(torch.nn.BatchNorm2d(out_c))
+
+        self.block = torch.nn.Sequential(*block)
+
+    def forward(self, x):
+        if self.pool:
+            x = self.pool(x)
+        out = self.block(x)
+        if self.shortcut:
+            return out + self.shortcut(x)
+        else:
+            return out
+
+
+class UpsampleBlock(torch.nn.Module):
+    def __init__(self, in_c, out_c, up_mode="transp_conv"):
+        super(UpsampleBlock, self).__init__()
+        block = []
+        if up_mode == "transp_conv":
+            block.append(
+                torch.nn.ConvTranspose2d(in_c, out_c, kernel_size=2, stride=2)
+            )
+        elif up_mode == "up_conv":
+            block.append(
+                torch.nn.Upsample(
+                    mode="bilinear", scale_factor=2, align_corners=False
+                )
+            )
+            block.append(torch.nn.Conv2d(in_c, out_c, kernel_size=1))
+        else:
+            raise Exception("Upsampling mode not supported")
+
+        self.block = torch.nn.Sequential(*block)
+
+    def forward(self, x):
+        out = self.block(x)
+        return out
+
+
+class ConvBridgeBlock(torch.nn.Module):
+    def __init__(self, channels, k_sz=3):
+        super(ConvBridgeBlock, self).__init__()
+        pad = (k_sz - 1) // 2
+        block = []
+
+        block.append(
+            torch.nn.Conv2d(channels, channels, kernel_size=k_sz, padding=pad)
+        )
+        block.append(torch.nn.ReLU())
+        block.append(torch.nn.BatchNorm2d(channels))
+
+        self.block = torch.nn.Sequential(*block)
+
+    def forward(self, x):
+        out = self.block(x)
+        return out
+
+
+class UpConvBlock(torch.nn.Module):
+    def __init__(
+        self,
+        in_c,
+        out_c,
+        k_sz=3,
+        up_mode="up_conv",
+        conv_bridge=False,
+        shortcut=False,
+    ):
+        super(UpConvBlock, self).__init__()
+        self.conv_bridge = conv_bridge
+
+        self.up_layer = UpsampleBlock(in_c, out_c, up_mode=up_mode)
+        self.conv_layer = ConvBlock(
+            2 * out_c, out_c, k_sz=k_sz, shortcut=shortcut, pool=False
+        )
+        if self.conv_bridge:
+            self.conv_bridge_layer = ConvBridgeBlock(out_c, k_sz=k_sz)
+
+    def forward(self, x, skip):
+        up = self.up_layer(x)
+        if self.conv_bridge:
+            out = torch.cat([up, self.conv_bridge_layer(skip)], dim=1)
+        else:
+            out = torch.cat([up, skip], dim=1)
+        out = self.conv_layer(out)
+        return out
+
+
+class LittleUNet(torch.nn.Module):
+    """Little U-Net model"""
+
+    def __init__(
+        self,
+        in_c,
+        n_classes,
+        layers,
+        k_sz=3,
+        up_mode="transp_conv",
+        conv_bridge=True,
+        shortcut=True,
+    ):
+        super(UNet, self).__init__()
+        self.n_classes = n_classes
+        self.first = ConvBlock(
+            in_c=in_c, out_c=layers[0], k_sz=k_sz, shortcut=shortcut, pool=False
+        )
+
+        self.down_path = torch.nn.ModuleList()
+        for i in range(len(layers) - 1):
+            block = ConvBlock(
+                in_c=layers[i],
+                out_c=layers[i + 1],
+                k_sz=k_sz,
+                shortcut=shortcut,
+                pool=True,
+            )
+            self.down_path.append(block)
+
+        self.up_path = torch.nn.ModuleList()
+        reversed_layers = list(reversed(layers))
+        for i in range(len(layers) - 1):
+            block = UpConvBlock(
+                in_c=reversed_layers[i],
+                out_c=reversed_layers[i + 1],
+                k_sz=k_sz,
+                up_mode=up_mode,
+                conv_bridge=conv_bridge,
+                shortcut=shortcut,
+            )
+            self.up_path.append(block)
+
+        # init, shamelessly lifted from torchvision/models/resnet.py
+        for m in self.modules():
+            if isinstance(m, torch.nn.Conv2d):
+                torch.nn.init.kaiming_normal_(
+                    m.weight, mode="fan_out", nonlinearity="relu"
+                )
+            elif isinstance(m, (torch.nn.BatchNorm2d, torch.nn.GroupNorm)):
+                torch.nn.init.constant_(m.weight, 1)
+                torch.nn.init.constant_(m.bias, 0)
+
+        self.final = torch.nn.Conv2d(layers[0], n_classes, kernel_size=1)
+
+    def forward(self, x):
+        x = self.first(x)
+        down_activations = []
+        for i, down in enumerate(self.down_path):
+            down_activations.append(x)
+            x = down(x)
+        down_activations.reverse()
+        for i, up in enumerate(self.up_path):
+            x = up(x, down_activations[i])
+        return self.final(x)
+
+
+class LittleWNet(torch.nn.Module):
+    """Little W-Net model, concatenating two Little U-Net models"""
+
+    def __init__(
+        self,
+        n_classes=1,
+        in_c=3,
+        layers=(8, 16, 32),
+        conv_bridge=True,
+        shortcut=True,
+        mode="train",
+    ):
+
+        super(wnet, self).__init__()
+        self.unet1 = LittleUNet(
+            in_c=in_c,
+            n_classes=n_classes,
+            layers=layers,
+            conv_bridge=conv_bridge,
+            shortcut=shortcut,
+        )
+        self.unet2 = LittleUNet(
+            in_c=in_c + n_classes,
+            n_classes=n_classes,
+            layers=layers,
+            conv_bridge=conv_bridge,
+            shortcut=shortcut,
+        )
+        self.n_classes = n_classes
+        self.mode = mode
+
+    def forward(self, x):
+        x1 = self.unet1(x)
+        x2 = self.unet2(torch.cat([x, x1], dim=1))
+        if self.mode != "train":
+            return x2
+        return x1, x2
+
+
+def lunet(in_c=3, n_classes=1):
+    """Builds Little U-Net segmentation network (uninitialized)
+
+
+    Parameters
+    ----------
+
+    input_channels : :py:class:`int`, Optional
+        Number of input channels the network should operate with
+
+    output_classes : :py:class:`int`, Optional
+        Number of output classes
+
+
+    Returns
+    -------
+
+    module : :py:class:`torch.nn.Module`
+        Network model for Little U-Net
+
+    """
+
+    return LittleUNet(
+        in_c=input_channels,
+        n_classes=output_classes,
+        layers=[8, 16, 32],
+        conv_bridge=True,
+        shortcut=True,
+    )
+
+
+def lwnet(input_channels=3, output_classes=1):
+    """Builds Little W-Net segmentation network (uninitialized)
+
+
+    Parameters
+    ----------
+
+    input_channels : :py:class:`int`, Optional
+        Number of input channels the network should operate with
+
+    output_classes : :py:class:`int`, Optional
+        Number of output classes
+
+
+    Returns
+    -------
+
+    module : :py:class:`torch.nn.Module`
+        Network model for Little W-Net
+
+    """
+
+    return LittleWNet(
+        in_c=input_channels,
+        n_classes=output_classes,
+        layers=[8, 16, 32],
+        conv_bridge=True,
+        shortcut=True,
+    )