Remove freia as dependency and include it in anomalib/models/components

.pre-commit-config.yaml

@@ -37,7 +37,7 @@ repos:
     hooks:
       - id: flake8
         args: [--config=tox.ini]
-        exclude: "tests/"
+        exclude: "tests|anomalib/models/components/freia"
 
   # python linting
   - repo: local
@@ -47,7 +47,7 @@ repos:
         entry: pylint --score=no --rcfile=tox.ini
         language: system
         types: [python]
-        exclude: "tests|docs"
+        exclude: "tests|docs|anomalib/models/components/freia"
 
   # python static type checking
   - repo: https://github.com/pre-commit/mirrors-mypy
@@ -56,7 +56,7 @@ repos:
       - id: mypy
         args: [--config-file=tox.ini]
         additional_dependencies: [types-PyYAML]
-        exclude: "tests/"
+        exclude: "tests|anomalib/models/components/freia"
 
   - repo: https://github.com/PyCQA/pydocstyle
     rev: 6.1.1
@@ -67,4 +67,4 @@ repos:
         language: python
         types: [python]
         args: [--config=tox.ini]
-        exclude: "tests|docs"
+        exclude: "tests|docs|anomalib/models/components/freia"

anomalib/models/cflow/backbone.py

@@ -16,12 +16,12 @@
 
 import math
 
-import FrEIA.framework as Ff
-import FrEIA.modules as Fm
 import torch
-from FrEIA.framework.sequence_inn import SequenceINN
 from torch import nn
 
+from anomalib.models.components.freia.framework import SequenceINN
+from anomalib.models.components.freia.modules import AllInOneBlock
+
 
 def positional_encoding_2d(condition_vector: int, height: int, width: int) -> torch.Tensor:
     """Creates embedding to store relative position of the feature vector using sine and cosine functions.
@@ -90,11 +90,11 @@ def cflow_head(
     Returns:
         SequenceINN: decoder network block
     """
-    coder = Ff.SequenceINN(n_features)
+    coder = SequenceINN(n_features)
     print("CNF coder:", n_features)
     for _ in range(coupling_blocks):
         coder.append(
-            Fm.AllInOneBlock,
+            AllInOneBlock,
             cond=0,
             cond_shape=(condition_vector,),
             subnet_constructor=subnet_fc,

anomalib/models/components/freia/README.md

@@ -0,0 +1,5 @@
+## FrEIA
+This sub-package contains freia packages to use within flow-based algorithms such as Cflow.
+
+## Description
+[FrEIA](https://github.com/VLL-HD/FrEIA) package is currently not available in pypi to install via pip. The only way to install it is `pip install git+https://github.com/VLL-HD/FrEIA.git`. PyPI, however, does not support installing packages from git links. Due to this limitation, anomalib cannot be updated on PyPI. To avoid this, `anomalib` contains some of the [FrEIA](https://github.com/VLL-HD/FrEIA) modules to facilitate CFlow training/inference.

anomalib/models/components/freia/__init__.py

@@ -0,0 +1,11 @@
+"""Framework for Easily Invertible Architectures.
+
+Module to construct invertible networks with pytorch, based on a graph
+structure of operations.
+
+Link to the original repo: https://github.com/VLL-HD/FrEIA
+"""
+
+# Copyright (c) 2018-2022 Lynton Ardizzone, Visual Learning Lab Heidelberg.
+# SPDX-License-Identifier: MIT
+#

anomalib/models/components/freia/framework/__init__.py

@@ -0,0 +1,9 @@
+"""Framework."""
+
+# Copyright (c) 2018-2022 Lynton Ardizzone, Visual Learning Lab Heidelberg.
+# SPDX-License-Identifier: MIT
+#
+
+from .sequence_inn import SequenceINN
+
+__all__ = ["SequenceINN"]

anomalib/models/components/freia/framework/sequence_inn.py

@@ -0,0 +1,120 @@
+"""Sequence INN."""
+
+# Copyright (c) 2018-2022 Lynton Ardizzone, Visual Learning Lab Heidelberg.
+# SPDX-License-Identifier: MIT
+#
+
+# pylint: disable=invalid-name
+# flake8: noqa
+# pylint: skip-file
+# type: ignore
+# pydocstyle: noqa
+
+from typing import Iterable, List, Tuple
+
+import torch
+from torch import Tensor, nn
+
+from anomalib.models.components.freia.modules.base import InvertibleModule
+
+
+class SequenceINN(InvertibleModule):
+    """Simpler than FrEIA.framework.GraphINN.
+
+    Only supports a sequential series of modules (no splitting, merging,
+    branching off).
+    Has an append() method, to add new blocks in a more simple way than the
+    computation-graph based approach of GraphINN. For example:
+    .. code-block:: python
+       inn = SequenceINN(channels, dims_H, dims_W)
+       for i in range(n_blocks):
+           inn.append(FrEIA.modules.AllInOneBlock, clamp=2.0, permute_soft=True)
+       inn.append(FrEIA.modules.HaarDownsampling)
+       # and so on
+    """
+
+    def __init__(self, *dims: int, force_tuple_output=False):
+        super().__init__([dims])
+
+        self.shapes = [tuple(dims)]
+        self.conditions = []
+        self.module_list = nn.ModuleList()
+
+        self.force_tuple_output = force_tuple_output
+
+    def append(self, module_class, cond=None, cond_shape=None, **kwargs):
+        """Append a reversible block from FrEIA.modules to the network.
+
+        Args:
+          module_class: Class from FrEIA.modules.
+          cond (int): index of which condition to use (conditions will be passed as list to forward()).
+            Conditioning nodes are not needed for SequenceINN.
+          cond_shape (tuple[int]): the shape of the condition tensor.
+          **kwargs: Further keyword arguments that are passed to the constructor of module_class (see example).
+        """
+
+        dims_in = [self.shapes[-1]]
+        self.conditions.append(cond)
+
+        if cond is not None:
+            kwargs["dims_c"] = [cond_shape]
+
+        module = module_class(dims_in, **kwargs)
+        self.module_list.append(module)
+        ouput_dims = module.output_dims(dims_in)
+        assert len(ouput_dims) == 1, "Module has more than one output"
+        self.shapes.append(ouput_dims[0])
+
+    def __getitem__(self, item):
+        """Get item."""
+        return self.module_list.__getitem__(item)
+
+    def __len__(self):
+        """Get length."""
+        return self.module_list.__len__()
+
+    def __iter__(self):
+        """Iter."""
+        return self.module_list.__iter__()
+
+    def output_dims(self, input_dims: List[Tuple[int]]) -> List[Tuple[int]]:
+        """Output Dims."""
+        if not self.force_tuple_output:
+            raise ValueError(
+                "You can only call output_dims on a SequentialINN " "when setting force_tuple_output=True."
+            )
+        return input_dims
+
+    def forward(
+        self, x_or_z: Tensor, c: Iterable[Tensor] = None, rev: bool = False, jac: bool = True
+    ) -> Tuple[Tensor, Tensor]:
+        """Execute the sequential INN in forward or inverse (rev=True) direction.
+
+        Args:
+            x_or_z: input tensor (in contrast to GraphINN, a list of
+                    tensors is not supported, as SequenceINN only has
+                    one input).
+            c: list of conditions.
+            rev: whether to compute the network forward or reversed.
+            jac: whether to compute the log jacobian
+        Returns:
+            z_or_x (Tensor): network output.
+            jac (Tensor): log-jacobian-determinant.
+        """
+
+        iterator = range(len(self.module_list))
+        log_det_jac = 0
+
+        if rev:
+            iterator = reversed(iterator)
+
+        if torch.is_tensor(x_or_z):
+            x_or_z = (x_or_z,)
+        for i in iterator:
+            if self.conditions[i] is None:
+                x_or_z, j = self.module_list[i](x_or_z, jac=jac, rev=rev)
+            else:
+                x_or_z, j = self.module_list[i](x_or_z, c=[c[self.conditions[i]]], jac=jac, rev=rev)
+            log_det_jac = j + log_det_jac
+
+        return x_or_z if self.force_tuple_output else x_or_z[0], log_det_jac

anomalib/models/components/freia/modules/__init__.py

@@ -0,0 +1,10 @@
+"""Modules."""
+
+# Copyright (c) 2018-2022 Lynton Ardizzone, Visual Learning Lab Heidelberg.
+# SPDX-License-Identifier: MIT
+#
+
+from .all_in_one_block import AllInOneBlock
+from .base import InvertibleModule
+
+__all__ = ["AllInOneBlock", "InvertibleModule"]