Initial Implementation of GLASS Model

Signed-off-by: Devansh Agarwal <[email protected]>

Author: code-dev05

src/anomalib/models/components/__init__.py

@@ -38,7 +38,7 @@
 
 from .base import AnomalibModule, BufferListMixin, DynamicBufferMixin, MemoryBankMixin
 from .dimensionality_reduction import PCA, SparseRandomProjection
-from .feature_extractors import TimmFeatureExtractor
+from .feature_extractors import TimmFeatureExtractor, NetworkFeatureAggregator
 from .filters import GaussianBlur2d
 from .sampling import KCenterGreedy
 from .stats import GaussianKDE, MultiVariateGaussian

src/anomalib/models/components/feature_extractors/__init__.py

@@ -28,7 +28,7 @@
 
 from .timm import TimmFeatureExtractor
 from .utils import dryrun_find_featuremap_dims
-
+from .network_feature_extractor import NetworkFeatureAggregator
 __all__ = [
     "dryrun_find_featuremap_dims",
     "TimmFeatureExtractor",

src/anomalib/models/components/feature_extractors/network_feature_extractor.py

@@ -0,0 +1,91 @@
+import torch
+from torch import nn
+import copy
+
+
+class NetworkFeatureAggregator(torch.nn.Module):
+    """Efficient extraction of network features."""
+
+    def __init__(self, backbone, layers_to_extract_from, train_backbone=False):
+        super(NetworkFeatureAggregator, self).__init__()
+        """Extraction of network features.
+
+        Runs a network only to the last layer of the list of layers where
+        network features should be extracted from.
+
+        Args:
+            backbone: torchvision.model
+            layers_to_extract_from: [list of str]
+        """
+        self.layers_to_extract_from = layers_to_extract_from
+        self.backbone = backbone
+        self.train_backbone = train_backbone
+        if not hasattr(backbone, "hook_handles"):
+            self.backbone.hook_handles = []
+        for handle in self.backbone.hook_handles:
+            handle.remove()
+        self.outputs = {}
+
+        for extract_layer in layers_to_extract_from:
+            self.register_hook(extract_layer)
+
+        self.to(self.device)
+
+    def forward(self, images, eval=True):
+        self.outputs.clear()
+        if self.train_backbone and not eval:
+            self.backbone(images)
+        else:
+            with torch.no_grad():
+                try:
+                    _ = self.backbone(images)
+                except LastLayerToExtractReachedException:
+                    pass
+        return self.outputs
+
+    def feature_dimensions(self, input_shape):
+        """Computes the feature dimensions for all layers given input_shape."""
+        _input = torch.ones([1] + list(input_shape)).to(self.device)
+        _output = self(_input)
+        return [_output[layer].shape[1] for layer in self.layers_to_extract_from]
+
+    def register_hook(self, layer_name):
+        module = self.find_module(self.backbone, layer_name)
+        if module is not None:
+            forward_hook = ForwardHook(
+                self.outputs, layer_name, self.layers_to_extract_from[-1]
+            )
+            if isinstance(module, torch.nn.Sequential):
+                hook = module[-1].register_forward_hook(forward_hook)
+            else:
+                hook = module.register_forward_hook(forward_hook)
+            self.backbone.hook_handles.append(hook)
+        else:
+            raise ValueError(f"Module {layer_name} not found in the model")
+
+    def find_module(self, model, module_name):
+        for name, module in model.named_modules():
+            if name == module_name:
+                return module
+            elif "." in module_name:
+                father, child = module_name.split(".", 1)
+                if name == father:
+                    return self.find_module(module, child)
+        return None
+
+
+class ForwardHook:
+    def __init__(self, hook_dict, layer_name: str, last_layer_to_extract: str):
+        self.hook_dict = hook_dict
+        self.layer_name = layer_name
+        self.raise_exception_to_break = copy.deepcopy(
+            layer_name == last_layer_to_extract
+        )
+
+    def __call__(self, module, input, output):
+        self.hook_dict[self.layer_name] = output
+        return None
+
+
+class LastLayerToExtractReachedException(Exception):
+    pass

src/anomalib/models/image/glass/__init__.py

@@ -0,0 +1 @@
+from .torch_model import GlassModel

src/anomalib/models/image/glass/torch_model.py

@@ -0,0 +1,237 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from anomalib.models.components import NetworkFeatureAggregator
+import math
+
+def init_weight(m):
+    if isinstance(m, torch.nn.Linear):
+        torch.nn.init.xavier_normal_(m.weight)
+    if isinstance(m, torch.nn.BatchNorm2d):
+        m.weight.data.normal_(1.0, 0.02)
+        m.bias.data.fill_(0)
+    elif isinstance(m, torch.nn.Conv2d):
+        m.weight.data.normal_(0.0, 0.02)
+
+class Preprocessing(torch.nn.Module):
+    def __init__(self, input_dims, output_dim):
+        super(Preprocessing, self).__init__()
+        self.input_dims = input_dims
+        self.output_dim = output_dim
+
+        self.preprocessing_modules = torch.nn.ModuleList()
+        for _ in input_dims:
+            module = MeanMapper(output_dim)
+            self.preprocessing_modules.append(module)
+
+    def forward(self, features):
+        _features = []
+        for module, feature in zip(self.preprocessing_modules, features):
+            _features.append(module(feature))
+        return torch.stack(_features, dim=1)
+
+
+class MeanMapper(torch.nn.Module):
+    def __init__(self, preprocessing_dim):
+        super(MeanMapper, self).__init__()
+        self.preprocessing_dim = preprocessing_dim
+
+    def forward(self, features):
+        features = features.reshape(len(features), 1, -1)
+        return F.adaptive_avg_pool1d(features, self.preprocessing_dim).squeeze(1)
+
+
+class Aggregator(torch.nn.Module):
+    def __init__(self, target_dim):
+        super(Aggregator, self).__init__()
+        self.target_dim = target_dim
+
+    def forward(self, features):
+        """Returns reshaped and average pooled features."""
+        features = features.reshape(len(features), 1, -1)
+        features = F.adaptive_avg_pool1d(features, self.target_dim)
+        return features.reshape(len(features), -1)
+
+class Projection(torch.nn.Module):
+    def __init__(self, in_planes, out_planes=None, n_layers=1, layer_type=0):
+        super(Projection, self).__init__()
+
+        if out_planes is None:
+            out_planes = in_planes
+        self.layers = torch.nn.Sequential()
+        _in = None
+        _out = None
+        for i in range(n_layers):
+            _in = in_planes if i == 0 else _out
+            _out = out_planes
+            self.layers.add_module(f"{i}fc", torch.nn.Linear(_in, _out))
+            if i < n_layers - 1:
+                if layer_type > 1:
+                    self.layers.add_module(f"{i}relu", torch.nn.LeakyReLU(.2))
+        self.apply(init_weight)
+
+    def forward(self, x):
+
+        x = self.layers(x)
+        return x
+    
+class Discriminator(torch.nn.Module):
+    def __init__(self, in_planes, n_layers=2, hidden=None):
+        super(Discriminator, self).__init__()
+
+        _hidden = in_planes if hidden is None else hidden
+        self.body = torch.nn.Sequential()
+        for i in range(n_layers - 1):
+            _in = in_planes if i == 0 else _hidden
+            _hidden = int(_hidden // 1.5) if hidden is None else hidden
+            self.body.add_module('block%d' % (i + 1),
+                                 torch.nn.Sequential(
+                                     torch.nn.Linear(_in, _hidden),
+                                     torch.nn.BatchNorm1d(_hidden),
+                                     torch.nn.LeakyReLU(0.2)
+                                 ))
+        self.tail = torch.nn.Sequential(torch.nn.Linear(_hidden, 1, bias=False),
+                                        torch.nn.Sigmoid())
+        self.apply(init_weight)
+
+    def forward(self, x):
+        x = self.body(x)
+        x = self.tail(x)
+        return x
+
+class PatchMaker:
+    def __init__(self, patchsize, top_k=0, stride=None):
+        self.patchsize = patchsize
+        self.stride = stride
+        self.top_k = top_k
+
+    def patchify(self, features, return_spatial_info=False):
+        padding = int((self.patchsize - 1) / 2)
+        unfolder = torch.nn.Unfold(kernel_size=self.patchsize, stride=self.stride, padding=padding, dilation=1)
+        unfolded_features = unfolder(features)
+        number_of_total_patches = []
+        for s in features.shape[-2:]:
+            n_patches = (s + 2 * padding - 1 * (self.patchsize - 1) - 1) / self.stride + 1
+            number_of_total_patches.append(int(n_patches))
+        unfolded_features = unfolded_features.reshape(
+            *features.shape[:2], self.patchsize, self.patchsize, -1
+        )
+        unfolded_features = unfolded_features.permute(0, 4, 1, 2, 3)
+
+        if return_spatial_info:
+            return unfolded_features, number_of_total_patches
+        return unfolded_features
+
+    def unpatch_scores(self, x, batchsize):
+        return x.reshape(batchsize, -1, *x.shape[1:])
+
+    def score(self, x):
+        x = x[:, :, 0]
+        x = torch.max(x, dim=1).values
+        return x
+
+class GlassModel(nn.Module):
+    def __init__(
+        self,
+        input_shape,
+        pretrain_embed_dim,
+        target_embed_dim,
+        backbone: nn.Module,
+        patchsize: int =3,
+        patchstride: int =1,
+        pre_trained: bool =True,
+        layers: list[str] = ["layer1", "layer2", "layer3"],
+        pre_proj: int = 1,
+        dsc_layers=2,
+        dsc_hidden=1024
+    ) -> None:
+        super().__init__()
+        self.backbone = backbone
+        self.layers = layers
+        self.input_shape = input_shape
+
+        self.forward_modules = torch.ModuleDict({})
+        feature_aggregator = NetworkFeatureAggregator(
+            self.backbone, self.layers, pre_trained
+        )
+        feature_dimensions = feature_aggregator.feature_dimensions(input_shape)
+        self.forward_modules["feature_aggregator"] = feature_aggregator
+
+        preprocessing = Preprocessing(feature_dimensions, pretrain_embed_dim)
+        self.forward_modules["preprocessing"] = preprocessing
+        self.target_embed_dimension = target_embed_dim
+        preadapt_aggregator = Aggregator(target_dim=target_embed_dim)
+        self.forward_modules["preadapt_aggregator"] = preadapt_aggregator
+
+        self.pre_trained = pre_trained
+
+        self.pre_proj = pre_proj
+        if self.pre_proj > 0:
+            self.pre_projection = Projection(self.target_embed_dimension, self.target_embed_dimension, pre_proj)
+
+        self.discriminator = Discriminator(self.target_embed_dimension, n_layers=dsc_layers, hidden=dsc_hidden)
+
+        self.patch_maker = PatchMaker(patchsize, stride=patchstride)
+
+    def generate_embeddings(self, images, provide_patch_shapes=False, eval=False):
+        if not eval and not self.pre_trained:
+            self.forward_modules["feature_aggregator"].train()
+            features = self.forward_modules["feature_aggregator"](images, eval=eval)
+        else:
+            self.forward_modules["feature_aggregator"].eval()
+            with torch.no_grad():
+                features = self.forward_modules["feature_aggregator"](images)
+        
+        features = [features[layer] for layer in self.layers_to_extract_from]
+        for i, feat in enumerate(features):
+            if len(feat.shape) == 3:
+                B, L, C = feat.shape
+                features[i] = feat.reshape(B, int(math.sqrt(L)), int(math.sqrt(L)), C).permute(0, 3, 1, 2)
+
+        features = [self.patch_maker.patchify(x, return_spatial_info=True) for x in features]
+        patch_shapes = [x[1] for x in features]
+        patch_features = [x[0] for x in features]
+        ref_num_patches = patch_shapes[0]
+
+        for i in range(1, len(patch_features)):
+            _features = patch_features[i]
+            patch_dims = patch_shapes[i]
+
+            _features = _features.reshape(
+                _features.shape[0], patch_dims[0], patch_dims[1], *_features.shape[2:]
+            )
+            _features = _features.permute(0, 3, 4, 5, 1, 2)
+            perm_base_shape = _features.shape
+            _features = _features.reshape(-1, *_features.shape[-2:])
+            _features = F.interpolate(
+                _features.unsqueeze(1),
+                size=(ref_num_patches[0], ref_num_patches[1]),
+                mode="bilinear",
+                align_corners=False,
+            )
+            _features = _features.squeeze(1)
+            _features = _features.reshape(
+                *perm_base_shape[:-2], ref_num_patches[0], ref_num_patches[1]
+            )
+            _features = _features.permute(0, 4, 5, 1, 2, 3)
+            _features = _features.reshape(len(_features), -1, *_features.shape[-3:])
+            patch_features[i] = _features
+
+        patch_features = [x.reshape(-1, *x.shape[-3:]) for x in patch_features]
+        patch_features = self.forward_modules["preprocessing"](patch_features)
+        patch_features = self.forward_modules["preadapt_aggregator"](patch_features)
+
+        return patch_features, patch_shapes
+    
+    def forward(self, images, eval=False):
+        self.forward_modules.eval()
+        with torch.no_grad():
+            if self.pre_proj > 0:
+                outputs = self.pre_proj(self.generate_embeddings(images, eval))
+                outputs = outputs[0] if len(outputs) == 2 else outputs
+            else:
+                outputs = self.generate_embeddings(images, eval)[0]
+            outputs = outputs[0] if len(outputs) == 2 else outputs
+            outputs = outputs.reshape(images.shape[0], -1, outputs.shape[-1])
+        return outputs
+