Added 10-Crop TTA

Author: BloodAxe

pytorch_toolbelt/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import absolute_import
 
-__version__ = '0.0.4'
+__version__ = '0.0.5'

pytorch_toolbelt/inference/functional.py

@@ -25,10 +25,20 @@ def torch_rot270(x: Tensor):
 
 
 def torch_flipud(x: Tensor):
+    """
+    Flip image tensor vertically
+    :param x:
+    :return:
+    """
     return x.flip(2)
 
 
-def torch_fliplp(x: Tensor):
+def torch_fliplr(x: Tensor):
+    """
+    Flip image tensor horizontally
+    :param x:
+    :return:
+    """
     return x.flip(3)
 
 
@@ -85,7 +95,7 @@ def pad_image_tensor(image_tensor: Tensor, pad_size: int = 32):
     return image_tensor, pad
 
 
-def unpad_tensor(image_tensor, pad):
+def unpad_image_tensor(image_tensor, pad):
     pad_left, pad_right, pad_top, pad_btm = pad
     rows, cols = image_tensor.size(2), image_tensor.size(3)
     return image_tensor[..., pad_top:rows - pad_btm, pad_left: cols - pad_right]

pytorch_toolbelt/inference/tta.py

@@ -21,7 +21,7 @@ def fliplr_image2label(model: nn.Module, image: Tensor) -> Tensor:
     :param image:
     :return:
     """
-    output = model(image) + model(F.torch_fliplp(image))
+    output = model(image) + model(F.torch_fliplr(image))
     one_over_2 = float(1.0 / 2.0)
     return output * one_over_2
 
@@ -30,10 +30,10 @@ def fivecrop_image2label(model: nn.Module, image: Tensor, crop_size: Tuple) -> T
     """Test-time augmentation for image classification that takes five crops out of input tensor (4 on corners and central)
     and averages predictions from them.
 
-    :param model:
-    :param image:
-    :param crop_size:
-    :return:
+    :param model: Classification model
+    :param image: Input image tensor
+    :param crop_size: Crop size. Must be smaller than image size
+    :return: Averaged logits
     """
     image_height, image_width = int(image.size(2)), int(image.size(3))
     crop_height, crop_width = crop_size
@@ -70,6 +70,55 @@ def fivecrop_image2label(model: nn.Module, image: Tensor, crop_size: Tuple) -> T
     return output * one_over_5
 
 
+def tencrop_image2label(model: nn.Module, image: Tensor, crop_size: Tuple) -> Tensor:
+    """Test-time augmentation for image classification that takes five crops out of input tensor (4 on corners and central)
+    and averages predictions from them and from their horisontally-flipped versions (10-Crop TTA).
+
+    :param model: Classification model
+    :param image: Input image tensor
+    :param crop_size: Crop size. Must be smaller than image size
+    :return: Averaged logits
+    """
+    image_height, image_width = int(image.size(2)), int(image.size(3))
+    crop_height, crop_width = crop_size
+
+    assert crop_height <= image_height
+    assert crop_width <= image_width
+
+    bottom_crop_start = image_height - crop_height
+    right_crop_start = image_width - crop_width
+    crop_tl = image[..., :crop_height, :crop_width]
+    crop_tr = image[..., :crop_height, right_crop_start:]
+    crop_bl = image[..., bottom_crop_start:, :crop_width]
+    crop_br = image[..., bottom_crop_start:, right_crop_start:]
+
+    assert crop_tl.size(2) == crop_height
+    assert crop_tr.size(2) == crop_height
+    assert crop_bl.size(2) == crop_height
+    assert crop_br.size(2) == crop_height
+
+    assert crop_tl.size(3) == crop_width
+    assert crop_tr.size(3) == crop_width
+    assert crop_bl.size(3) == crop_width
+    assert crop_br.size(3) == crop_width
+
+    center_crop_y = (image_height - crop_height) // 2
+    center_crop_x = (image_width - crop_width) // 2
+
+    crop_cc = image[..., center_crop_y:center_crop_y + crop_height, center_crop_x:center_crop_x + crop_width]
+    assert crop_cc.size(2) == crop_height
+    assert crop_cc.size(3) == crop_width
+
+    output = model(crop_tl) + model(F.torch_fliplr(crop_tl)) + \
+             model(crop_tr) + model(F.torch_fliplr(crop_tr)) + \
+             model(crop_bl) + model(F.torch_fliplr(crop_bl)) + \
+             model(crop_br) + model(F.torch_fliplr(crop_br)) + \
+             model(crop_cc) + model(F.torch_fliplr(crop_cc))
+
+    one_over_10 = float(1.0 / 10.0)
+    return output * one_over_10
+
+
 def fliplr_image2mask(model: nn.Module, image: Tensor) -> Tensor:
     """Test-time augmentation for image segmentation that averages predictions
     for input image and vertically flipped one.
@@ -80,7 +129,7 @@ def fliplr_image2mask(model: nn.Module, image: Tensor) -> Tensor:
     :param image: Model input.
     :return: Arithmetically averaged predictions
     """
-    output = model(image) + F.torch_fliplp(model(F.torch_fliplp(image)))
+    output = model(image) + F.torch_fliplr(model(F.torch_fliplr(image)))
     one_over_2 = float(1.0 / 2.0)
     return output * one_over_2
 
@@ -129,7 +178,7 @@ def d4_image2mask(model: nn.Module, image: Tensor) -> Tensor:
 
     for aug, deaug in zip([F.torch_none, F.torch_rot90, F.torch_rot180, F.torch_rot270], [F.torch_none, F.torch_rot270, F.torch_rot180, F.torch_rot90]):
         x = deaug(model(aug(image)))
-        output = output + x
+        output = output + F.torch_transpose(x)
 
     one_over_8 = float(1.0 / 8.0)
     return output * one_over_8

pytorch_toolbelt/utils/torch_utils.py

@@ -9,6 +9,17 @@
 from torch import nn
 
 
+def freeze_bn(module: nn.Module):
+    """Freezes BatchNorm
+    """
+    if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
+        module.track_running_stats = False
+
+    for m in module.modules():
+        if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
+            module.track_running_stats = False
+
+
 def logit(x: torch.Tensor, eps=1e-5):
     x = torch.clamp(x.float(), eps, 1.0 - eps)
     return torch.log(x / (1.0 - x))

tests/test_tta.py

@@ -0,0 +1,89 @@
+import torch
+import numpy as np
+from pytorch_toolbelt.inference import tta
+from pytorch_toolbelt.utils.torch_utils import to_numpy
+from torch import nn
+
+
+class NoOp(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, input):
+        return input
+
+
+class SumAll(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, input):
+        return input.sum(dim=[1, 2, 3])
+
+
+def test_d4_image2mask():
+    input = torch.rand((4, 3, 224, 224))
+    model = NoOp()
+
+    output = tta.d4_image2mask(model, input)
+    np.testing.assert_allclose(to_numpy(output), to_numpy(input), atol=1e-6, rtol=1e-6)
+
+
+def test_fliplr_image2mask():
+    input = torch.rand((4, 3, 224, 224))
+    model = NoOp()
+
+    output = tta.fliplr_image2mask(model, input)
+    np.testing.assert_allclose(to_numpy(output), to_numpy(input), atol=1e-6, rtol=1e-6)
+
+
+def test_d4_image2label():
+    input = torch.tensor([[1, 2, 3, 4],
+                          [5, 6, 7, 8],
+                          [9, 0, 1, 2],
+                          [3, 4, 5, 6]]).unsqueeze(0).unsqueeze(0).float()
+    model = SumAll()
+
+    output = tta.d4_image2label(model, input)
+    expected = int(input.sum())
+
+    assert int(output) == expected
+
+
+def test_fliplr_image2label():
+    input = torch.tensor([[1, 2, 3, 4],
+                          [5, 6, 7, 8],
+                          [9, 0, 1, 2],
+                          [3, 4, 5, 6]]).unsqueeze(0).unsqueeze(0).float()
+    model = SumAll()
+
+    output = tta.fliplr_image2label(model, input)
+    expected = int(input.sum())
+
+    assert int(output) == expected
+
+
+def test_fivecrop_image2label():
+    input = torch.tensor([[1, 2, 3, 4],
+                          [5, 6, 7, 8],
+                          [9, 0, 1, 2],
+                          [3, 4, 5, 6]]).unsqueeze(0).unsqueeze(0).float()
+    model = SumAll()
+
+    output = tta.fivecrop_image2label(model, input, (2, 2))
+    expected = ((1 + 2 + 5 + 6) + (3 + 4 + 7 + 8) + (9 + 0 + 3 + 4) + (1 + 2 + 5 + 6) + (6 + 7 + 0 + 1)) / 5
+
+    assert int(output) == expected
+
+
+def test_tencrop_image2label():
+    input = torch.tensor([[1, 2, 3, 4],
+                          [5, 6, 7, 8],
+                          [9, 0, 1, 2],
+                          [3, 4, 5, 6]]).unsqueeze(0).unsqueeze(0).float()
+    model = SumAll()
+
+    output = tta.tencrop_image2label(model, input, (2, 2))
+    expected = (2 * ((1 + 2 + 5 + 6) + (3 + 4 + 7 + 8) + (9 + 0 + 3 + 4) + (1 + 2 + 5 + 6) + (6 + 7 + 0 + 1))) / 10
+
+    assert int(output) == expected