Merge pull request #69 from BloodAxe/develop

Release of pytorch-toolbelt 0.5

.gitignore

@@ -19,3 +19,4 @@ var/
 .idea/
 .pytest_cache/
 /tests/tta_eval.csv
+/tests/tmp.onnx

README.md

@@ -1,6 +1,29 @@
-# Pytorch-toolbelt
+# Important Update
+
+![ukraine-flag](docs/480px-Flag_of_Ukraine.jpg)
+
+On February 24th, 2022, Russia declared war and invaded peaceful Ukraine. 
+After the annexation of Crimea and the occupation of the Donbas region, Putin's regime decided to destroy Ukrainian nationality.
+Ukrainians show fierce resistance and demonstrate to the entire world what it's like to fight for the nation's independence.
+
+Ukraine's government launched a website to help russian mothers, wives & sisters find their beloved ones killed or captured in Ukraine - https://200rf.com & https://t.me/rf200_now (Telegram channel).
+Our goal is to inform those still in Russia & Belarus, so they refuse to assault Ukraine. 
+
+Help us get maximum exposure to what is happening in Ukraine, violence, and inhuman acts of terror that the "Russian World" has brought to Ukraine. 
+This is a comprehensive Wiki on how you can help end this war: https://how-to-help-ukraine-now.super.site/ 
 
-[![Documentation Status](https://readthedocs.org/projects/pytorch-toolbelt/badge/?version=latest)](https://pytorch-toolbelt.readthedocs.io/en/latest/?badge=latest)
+Official channels
+* [Official account of the Parliament of Ukraine](https://t.me/verkhovnaradaofukraine)
+* [Ministry of Defence](https://www.facebook.com/MinistryofDefence.UA)
+* [Office of the president](https://www.facebook.com/president.gov.ua)
+* [Cabinet of Ministers of Ukraine](https://www.facebook.com/KabminUA)
+* [Center of strategic communications](https://www.facebook.com/StratcomCentreUA)
+* [Minister of Foreign Affairs of Ukraine](https://twitter.com/DmytroKuleba)
+
+Glory to Ukraine!
+
+
+# Pytorch-toolbelt
 
 A `pytorch-toolbelt` is a Python library with a set of bells and whistles for PyTorch for fast R&D prototyping and Kaggle farming:
 

pytorch_toolbelt/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import absolute_import
 
-__version__ = "0.4.4"
+__version__ = "0.5.0"

pytorch_toolbelt/datasets/__init__.py

@@ -2,3 +2,4 @@
 from .classification import *
 from .segmentation import *
 from .wrappers import *
+from .mean_std import *

pytorch_toolbelt/datasets/mean_std.py

@@ -0,0 +1,76 @@
+import numpy as np
+from typing import Optional, Tuple
+
+__all__ = ["DatasetMeanStdCalculator"]
+
+
+class DatasetMeanStdCalculator:
+    __slots__ = ["global_mean", "global_var", "n_items", "num_channels", "global_max", "global_min"]
+
+    """
+    Class to calculate running mean and std of the dataset. It helps when whole dataset does not fit entirely in RAM.
+    """
+
+    def __init__(self, num_channels: int = 3):
+        """
+        Create a new instance of DatasetMeanStdCalculator
+
+        Args:
+            num_channels: Number of channels in the image. Default value is 3
+        """
+        super(DatasetMeanStdCalculator, self).__init__()
+        self.num_channels = num_channels
+        self.global_mean = None
+        self.global_var = None
+        self.global_max = None
+        self.global_min = None
+        self.n_items = 0
+        self.reset()
+
+    def reset(self):
+        self.global_mean = np.zeros(self.num_channels, dtype=np.float64)
+        self.global_var = np.zeros(self.num_channels, dtype=np.float64)
+        self.global_max = np.ones_like(self.global_mean) * float("-inf")
+        self.global_min = np.ones_like(self.global_mean) * float("+inf")
+        self.n_items = 0
+
+    def accumulate(self, image: np.ndarray, mask: Optional[np.ndarray] = None) -> None:
+        """
+        Compute mean and std of a single image and integrates it into global statistics
+        Args:
+            image: Input image (Must be HWC, with number of channels C equal to self.num_channels)
+            mask: Optional mask to include only certain parts of image from statistics computation.
+            Only non-zero elements will be included,
+        """
+        if len(image.shape) == 2:
+            image = np.expand_dims(image, axis=-1)
+
+        if self.num_channels != image.shape[2]:
+            raise RuntimeError(f"Number of channels in image must be {self.num_channels}, got {image.shape[2]}.")
+        image = image.reshape((-1, self.num_channels))
+
+        if mask is not None:
+            mask = mask.reshape((mask.shape[0] * mask.shape[1], 1))
+            image = image[mask]
+
+            # In case the whole image is masked out, we exclude it entirely
+            if len(image) == 0:
+                return
+
+        mean = np.mean(image, axis=0)
+        std = np.std(image, axis=0)
+
+        self.global_mean += np.squeeze(mean)
+        self.global_var += np.squeeze(std) ** 2
+        self.global_max = np.maximum(self.global_max, np.max(image, axis=0))
+        self.global_min = np.minimum(self.global_min, np.min(image, axis=0))
+        self.n_items += 1
+
+    def compute(self) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Compute dataset-level mean & std
+
+        Returns:
+            Tuple of global [mean, std] per channel
+        """
+        return self.global_mean / self.n_items, np.sqrt(self.global_var / self.n_items)

pytorch_toolbelt/inference/ensembling.py

@@ -1,14 +1,20 @@
 import torch
 from torch import nn, Tensor
-from typing import List, Union, Iterable, Optional, Dict
+from typing import List, Union, Iterable, Optional, Dict, Tuple
 
 __all__ = ["ApplySoftmaxTo", "ApplySigmoidTo", "Ensembler", "PickModelOutput", "SelectByIndex"]
 
 from pytorch_toolbelt.inference.tta import _deaugment_averaging
 
 
 class ApplySoftmaxTo(nn.Module):
-    def __init__(self, model: nn.Module, output_key: Union[str, List[str]] = "logits", dim=1, temperature=1):
+    output_keys: Tuple
+    temperature: float
+    dim: int
+
+    def __init__(
+        self, model: nn.Module, output_key: Union[str, Iterable[str]] = "logits", dim: int = 1, temperature: float = 1
+    ):
         """
         Apply softmax activation on given output(s) of the model
         :param model: Model to wrap
@@ -17,39 +23,42 @@ def __init__(self, model: nn.Module, output_key: Union[str, List[str]] = "logits
         :param temperature: Temperature scaling coefficient. Values > 1 will make logits sharper.
         """
         super().__init__()
-        output_key = output_key if isinstance(output_key, (list, tuple)) else [output_key]
         # By converting to set, we prevent double-activation by passing output_key=["logits", "logits"]
-        self.output_keys = set(output_key)
+        output_key = tuple(set(output_key)) if isinstance(output_key, Iterable) else tuple([output_key])
+        self.output_keys = output_key
         self.model = model
         self.dim = dim
         self.temperature = temperature
 
     def forward(self, *input, **kwargs):
         output = self.model(*input, **kwargs)
         for key in self.output_keys:
-            output[key] = output[key].mul(self.temperature).softmax(dim=1)
+            output[key] = output[key].mul(self.temperature).softmax(dim=self.dim)
         return output
 
 
 class ApplySigmoidTo(nn.Module):
-    def __init__(self, model: nn.Module, output_key: Union[str, List[str]] = "logits", temperature=1):
+    output_keys: Tuple
+    temperature: float
+
+    def __init__(self, model: nn.Module, output_key: Union[str, Iterable[str]] = "logits", temperature=1):
         """
         Apply sigmoid activation on given output(s) of the model
         :param model: Model to wrap
         :param output_key: string or list of strings, indicating to what outputs sigmoid activation should be applied.
         :param temperature: Temperature scaling coefficient. Values > 1 will make logits sharper.
         """
         super().__init__()
-        output_key = output_key if isinstance(output_key, (list, tuple)) else [output_key]
         # By converting to set, we prevent double-activation by passing output_key=["logits", "logits"]
-        self.output_keys = set(output_key)
+        output_key = tuple(set(output_key)) if isinstance(output_key, Iterable) else tuple([output_key])
+        self.output_keys = output_key
         self.model = model
         self.temperature = temperature
 
     def forward(self, *input, **kwargs):  # skipcq: PYL-W0221
         output = self.model(*input, **kwargs)
         for key in self.output_keys:
-            output[key] = output[key].mul(self.temperature).sigmoid()
+            output[key] = output[key].mul(self.temperature).sigmoid_()
         return output
 
 

pytorch_toolbelt/inference/functional.py

@@ -140,7 +140,9 @@ def torch_transpose2(x: Tensor):
     return x.transpose(3, 2)
 
 
-def pad_image_tensor(image_tensor: Tensor, pad_size: Union[int, Tuple[int, int]] = 32):
+def pad_image_tensor(
+    image_tensor: Tensor, pad_size: Union[int, Tuple[int, int]] = 32
+) -> Tuple[Tensor, Tuple[int, int, int, int]]:
     """Pad input tensor to make it's height and width dividable by @pad_size
 
     :param image_tensor: 4D image tensor of shape NCHW

pytorch_toolbelt/inference/tiles.py

@@ -334,14 +334,24 @@ def integrate_batch(self, batch: torch.Tensor, crop_coords):
         if len(batch) != len(crop_coords):
             raise ValueError("Number of images in batch does not correspond to number of coordinates")
 
-        batch = batch.to(device=self.image.device)
+        if batch.device != self.image.device:
+            batch = batch.to(device=self.image.device)
+
         for tile, (x, y, tile_width, tile_height) in zip(batch, crop_coords):
             self.image[:, y : y + tile_height, x : x + tile_width] += tile * self.weight
             self.norm_mask[:, y : y + tile_height, x : x + tile_width] += self.weight
 
+    @property
+    def device(self) -> torch.device:
+        return self.image.device
+
     def merge(self) -> torch.Tensor:
         return self.image / self.norm_mask
 
+    def merge_(self) -> torch.Tensor:
+        self.image /= self.norm_mask
+        return self.image
+
 
 @pytorch_toolbelt_deprecated("This class is deprecated and will be removed in 0.5.0. Please use TileMerger instead.")
 class CudaTileMerger(TileMerger):