Merge pull request #73 from BloodAxe/feature/0.5.1

Feature/0.5.1

Author: BloodAxe

pytorch_toolbelt/inference/functional.py

@@ -7,26 +7,27 @@
 from ..utils.support import pytorch_toolbelt_deprecated
 
 __all__ = [
+    "geometric_mean",
+    "harmonic_mean",
+    "logodd_mean",
+    "pad_image_tensor",
+    "torch_fliplr",
+    "torch_flipud",
     "torch_none",
+    "torch_rot180",
+    "torch_rot270",
     "torch_rot90",
-    "torch_rot90_cw",
     "torch_rot90_ccw",
-    "torch_transpose_rot90_cw",
-    "torch_transpose_rot90_ccw",
     "torch_rot90_ccw_transpose",
+    "torch_rot90_cw",
     "torch_rot90_cw_transpose",
-    "torch_rot180",
-    "torch_rot270",
-    "torch_fliplr",
-    "torch_flipud",
     "torch_transpose",
     "torch_transpose2",
     "torch_transpose_",
-    "pad_image_tensor",
+    "torch_transpose_rot90_ccw",
+    "torch_transpose_rot90_cw",
     "unpad_image_tensor",
     "unpad_xyxy_bboxes",
-    "geometric_mean",
-    "harmonic_mean",
 ]
 
 
@@ -240,3 +241,23 @@ def harmonic_mean(x: Tensor, dim: int, eps: float = 1e-6) -> Tensor:
     x = torch.mean(x, dim=dim)
     x = torch.reciprocal(x.clamp_min(eps))
     return x
+
+
+def logodd_mean(x: Tensor, dim: int, eps: float = 1e-6) -> Tensor:
+    """
+    Compute log-odd mean along given dimension.
+    logodd = log(p / (1 - p))
+
+    This implementation assume values are in range [0, 1] (Probabilities)
+    Args:
+        x: Input tensor of arbitrary shape
+        dim: Dimension to reduce
+
+    Returns:
+        Tensor
+    """
+    x = x.clamp(min=eps, max=1.0 - eps)
+    x = torch.log(x / (1 - x))
+    x = torch.mean(x, dim=dim)
+    x = torch.exp(x) / (1 + torch.exp(x))
+    return x

pytorch_toolbelt/inference/tta.py

@@ -79,6 +79,8 @@ def _deaugment_averaging(x: Tensor, reduction: MaybeStrOrCallable) -> Tensor:
         x = F.geometric_mean(x, dim=0)
     elif reduction in {"hmean", "harmonic_mean"}:
         x = F.harmonic_mean(x, dim=0)
+    elif reduction == "logodd":
+        x = F.logodd_mean(x, dim=0)
     elif callable(reduction):
         x = reduction(x, dim=0)
     elif reduction in {None, "None", "none"}:

pytorch_toolbelt/optimization/functional.py

@@ -46,7 +46,9 @@ def get_optimizable_parameters(model: nn.Module) -> Iterator[nn.Parameter]:
     return filter(lambda x: x.requires_grad, model.parameters())
 
 
-def freeze_model(module: nn.Module, freeze_parameters: Optional[bool] = True, freeze_bn: Optional[bool] = True):
+def freeze_model(
+    module: nn.Module, freeze_parameters: Optional[bool] = True, freeze_bn: Optional[bool] = True
+) -> nn.Module:
     """
     Change 'requires_grad' value for module and it's child modules and
     optionally freeze batchnorm modules.
@@ -70,3 +72,5 @@ def freeze_model(module: nn.Module, freeze_parameters: Optional[bool] = True, fr
         for m in module.modules():
             if isinstance(m, bn_types):
                 module.track_running_stats = not freeze_bn
+
+    return module

pytorch_toolbelt/utils/fs.py

@@ -16,6 +16,7 @@
     "find_images_in_dir",
     "find_in_dir",
     "find_in_dir_glob",
+    "find_subdirectories_in_dir",
     "has_ext",
     "has_image_ext",
     "id_from_fname",
@@ -45,6 +46,19 @@ def find_in_dir(dirname: str) -> List[str]:
     return [os.path.join(dirname, fname) for fname in sorted(os.listdir(dirname))]
 
 
+def find_subdirectories_in_dir(dirname: str) -> List[str]:
+    """
+    Retrieve list of subdirectories (non-recursive) in the given directory.
+    Args:
+        dirname: Target directory name
+
+    Returns:
+        Sorted list of absolute paths to directories
+    """
+    all_entries = find_in_dir(dirname)
+    return [entry for entry in all_entries if os.path.isdir(entry)]
+
+
 def find_in_dir_with_ext(dirname: str, extensions: Union[str, List[str]]) -> List[str]:
     return [os.path.join(dirname, fname) for fname in sorted(os.listdir(dirname)) if has_ext(fname, extensions)]
 

pytorch_toolbelt/utils/visualization.py

@@ -29,12 +29,31 @@ def plot_confusion_matrix(
     title: str = "Confusion matrix",
     cmap=None,
     fname=None,
+    show_scores: bool = True,
     noshow: bool = False,
     backend: str = "Agg",
     format_string: Optional[str] = None,
 ):
-    """Render the confusion matrix and return matplotlib's figure with it.
+    """
+    Render the confusion matrix and return matplotlib's figure with it.
     Normalization can be applied by setting `normalize=True`.
+
+    Args:
+        cm: Numpy array of (N,N) shape - confusion matrix array
+        class_names: List of [N] names of the classes
+        figsize:
+        fontsize:
+        normalize: Whether to apply normalization for each row of CM
+        title: Title of the confusion matrix
+        cmap:
+        fname: Filename of the rendered confusion matrix
+        show_scores: Show scores in each cell
+        noshow:
+        backend:
+        format_string:
+
+    Returns:
+        Matplotlib's figure
     """
     import matplotlib
 
@@ -64,26 +83,12 @@ def plot_confusion_matrix(
     if format_string is None:
         format_string = ".3f" if normalize else "d"
 
-    thresh = (cm.max() + cm.min()) / 2.0
-    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
-        if np.isfinite(cm[i, j]):
-            plt.text(
-                j,
-                i,
-                format(cm[i, j], format_string),
-                horizontalalignment="center",
-                fontsize=fontsize,
-                color="white" if cm[i, j] > thresh else "black",
-            )
-        else:
-            plt.text(
-                j,
-                i,
-                "N/A",
-                horizontalalignment="center",
-                fontsize=fontsize,
-                color="black",
-            )
+    if show_scores:
+        thresh = (cm.max() + cm.min()) / 2.0
+        for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
+            text = format(cm[i, j], format_string) if np.isfinite(cm[i, j]) else "N/A"
+            color = "white" if cm[i, j] > thresh else "black"
+            plt.text(j, i, text, horizontalalignment="center", fontsize=fontsize, color=color)
 
     plt.ylabel("True label")