Unet

docs/source/index.rst

@@ -54,6 +54,7 @@ Dana-Farber Cancer Institute.
  examples/link_stain_normalization
  examples/link_nucleus_detection
  examples/link_preprocessing_pipeline
+ examples/link_train_hovernet
 
 .. toctree::
  :maxdepth: 2

docs/source/models.rst

@@ -3,20 +3,28 @@ Models
 
 ``PathML`` comes with several model architectures ready to use out of the box.
 
-.. table::
- :widths: 20, 20, 60
+.. list-table:: Models included in PathML
+ :widths: 15, 70, 15
+ :header-rows: 1
 
- ===================================== ============ =============
- Model Reference Description
- ===================================== ============ =============
- U-net (in progress) [Unet]_ A model for segmentation in biomedical images
- :class:`~pathml.ml.hovernet.HoVerNet` [HoVerNet]_ A model for nucleus segmentation and classification in H&E images
- ===================================== ============ =============
+ * - Model
+ - Description
+ - Reference
+ * - :class:`~pathml.ml.unet.UNet`
+ - A standard general-purpose model designed for segmentation in biomedical images.
+ Architecture consists of an 4 encoder blocks followed by 4 decoder blocks.
+ Skip connections propagate information from each layer of the encoder to the corresponding layer in
+ the decoder.
+ - [Unet]_
+ * - :class:`~pathml.ml.hovernet.HoVerNet`
+ - A model for simultaneous nucleus segmentation and classification in H&E images.
+ Architecture consists of a single encoder with three separate decoder branches: one to perform binary
+ classification of nuclear pixels (NP), one to compute horizontal and vertical nucleus maps (HV), and one which
+ is used in the classification setting to perform classification of nuclear pixels (NC).
+ - [HoVerNet]_
 
 You can also use models from `torchvision.models <https://pytorch.org/docs/stable/torchvision/models.html>`_, or create your own!
 
-In many cases, model parameters (weights) for pretrained networks may be available for use through the Model Repository.
-
 References
 ----------
 

environment.yml

@@ -13,6 +13,7 @@ dependencies:
  - python-spams>=2.6.1
  - pip>=20.0.2
  - pytorch>=1.7.0
+ - torchvision>=0.8.0
  - openjdk==8.0.152
  - h5py>=3.1.0
  - pip:

pathml/ml/unet.py

@@ -0,0 +1,124 @@
+import torch
+from torch import nn
+from torchvision.transforms import CenterCrop
+from torch.nn.functional import interpolate
+
+
+class _UNetConvBlock(nn.Module):
+ """
+ Convolution block for U-Net
+
+ From the paper:
+ The contracting path follows the typical architecture of a convolutional network.
+ It consists of the repeated application of two 3x3 convolutions (unpadded convolutions), each followed
+ by a rectified linear unit (ReLU)...
+ """
+ def __init__(self, in_c, out_c):
+ super().__init__()
+ self.conv1 = nn.Conv2d(in_channels = in_c, out_channels = out_c, kernel_size = 3)
+ self.relu = nn.ReLU()
+ self.conv2 = nn.Conv2d(in_channels = out_c, out_channels = out_c, kernel_size = 3)
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.relu(x)
+ x = self.conv2(x)
+ x = self.relu(x)
+ return x
+
+
+class _UNetUpConvBlock(nn.Module):
+ """
+ Up-Convolution block for U-Net
+
+ From the paper:
+ Every step in the expansive path consists of an upsampling of the feature map followed by a 2x2 convolution
+ (“up-convolution”) that halves the number of feature channels, a concatenation with the correspondingly
+ cropped feature map from the contracting path, and two 3x3 convolutions, each followed by a ReLU.
+ """
+ def __init__(self, in_c, out_c):
+ super().__init__()
+ self.up = nn.ConvTranspose2d(in_channels = in_c, out_channels = out_c, kernel_size = 2, stride = 2)
+ self.conv = _UNetConvBlock(in_c = in_c, out_c = out_c)
+
+ def forward(self, x, x_skip):
+ """
+ x is the input
+ x_skip is the skip connection from the encoder block
+ """
+ x = self.up(x)
+ # crop tensor from skip connection to match H and W of x
+ x_skip = CenterCrop((x.shape[2], x.shape[3]))(x_skip)
+ x = torch.cat([x, x_skip], dim = 1)
+ x = self.conv(x)
+ return x
+
+
+class UNet(nn.Module):
+ """
+ U-Net is a convolutional network for biomedical image segmentation.
+ The architecture consists of a contracting path to capture context and a symmetric expanding
+ path that enables precise localization.
+
+ As described in the original paper, by default no padding is used, so the dimensions get smaller each layer.
+ Input of size 572px will lead to output of size 388px (See Fig. 1 in the paper).
+ The ``keep_dim`` parameter can be used to enfore the output to be the same shape as the input.
+
+ Code is based on:
+ https://amaarora.github.io/2020/09/13/unet.html
+ https://github.com/LeeJunHyun/Image_Segmentation
+
+ Args:
+ in_channels (int): Number of channels in input. E.g. 3 for RGB image
+ out_channels (int): Number of channels in output. E.g. 1 for a binary classification setting.
+ keep_dim (bool): Whether to enforce output to match the dimensions of input. If ``True``, a final interpolation
+ step will be applied. Defaults to ``False``.
+
+ References:
+ Ronneberger, O., Fischer, P. and Brox, T., 2015, October.
+ U-net: Convolutional networks for biomedical image segmentation.
+ In International Conference on Medical image computing and computer-assisted intervention (pp. 234-241).
+ Springer, Cham.
+ """
+
+ def __init__(self, in_channels=3, out_channels=1, keep_dim=False):
+ super().__init__()
+ self.keep_dim = keep_dim
+ self.pool = nn.MaxPool2d(2)
+
+ self.conv1 = _UNetConvBlock(in_c = in_channels, out_c = 64)
+ self.conv2 = _UNetConvBlock(in_c = 64, out_c = 128)
+ self.conv3 = _UNetConvBlock(in_c = 128, out_c = 256)
+ self.conv4 = _UNetConvBlock(in_c = 256, out_c = 512)
+ self.conv5 = _UNetConvBlock(in_c = 512, out_c = 1024)
+
+ self.upconv1 = _UNetUpConvBlock(in_c = 1024, out_c = 512)
+ self.upconv2 = _UNetUpConvBlock(in_c = 512, out_c = 256)
+ self.upconv3 = _UNetUpConvBlock(in_c = 256, out_c = 128)
+ self.upconv4 = _UNetUpConvBlock(in_c = 128, out_c = 64)
+
+ self.head = nn.Conv2d(in_channels = 64, out_channels = out_channels, kernel_size = 1)
+
+ def forward(self, x):
+ # encoder
+ x1 = self.conv1(x)
+ x2 = self.pool(x1)
+ x2 = self.conv2(x2)
+ x3 = self.pool(x2)
+ x3 = self.conv3(x3)
+ x4 = self.pool(x3)
+ x4 = self.conv4(x4)
+ x5 = self.pool(x4)
+ x5 = self.conv5(x5)
+
+ # decoder
+ up1 = self.upconv1(x = x5, x_skip = x4)
+ up2 = self.upconv2(x = up1, x_skip = x3)
+ up3 = self.upconv3(x = up2, x_skip = x2)
+ up4 = self.upconv4(x = up3, x_skip = x1)
+ out = self.head(up4)
+
+ if self.keep_dim:
+ out = interpolate(out, size = (x.shape[2], x.shape[3]))
+
+ return out

tests/ml_tests/test_unet.py

@@ -0,0 +1,35 @@
+import pytest
+import torch
+
+from pathml.ml.unet import UNet
+
+
+@pytest.mark.parametrize("keepdim", [True, False])
+@pytest.mark.parametrize("out_c", [1, 3])
+def test_unet_shapes(out_c, keepdim):
+ batch_size = 1
+ channels_in = 3
+ im_size_in = 572
+
+ x = torch.randn(batch_size, channels_in, im_size_in, im_size_in)
+
+ net = UNet(out_channels = out_c, keep_dim = keepdim)
+ out = net(x)
+
+ if keepdim:
+ assert out.shape == (batch_size, out_c, im_size_in, im_size_in)
+ else:
+
+ # compute output size, if keep_dim is false
+ nlayers = 4
+ im_size_out = im_size_in
+ # at each layer in encoder, two conv layers without padding loses 4px total, followed by a downsample by 2
+ for _ in range(nlayers):
+ im_size_out = (im_size_out - 4) / 2
+ # two more conv layers at the bottom layer
+ im_size_out = im_size_out - 4
+ # at each layer in decoder, upsample by 2 followed by two conv layers without padding for a loss of 4
+ for _ in range(nlayers):
+ im_size_out = (im_size_out * 2) - 4
+
+ assert out.shape == (batch_size, out_c, im_size_out, im_size_out)