Allow specifying feature layer and pool factor in DFM (#215)

* Allow specifying feature layer and pool factor in DFM

* 🏷 Rename pool to pooling_kernel_size to make it more descriptive.

* 🏷 renamed variable names in `get_features` method

* 🏷 Change the keyword argument in Dfm Lightning Module to properly import the Dfm Model.

Co-authored-by: Samet Akcay <samet.akcay@intel.com>

anomalib/models/dfm/config.yaml

@@ -16,6 +16,8 @@ dataset:
 model:
   name: dfm
   backbone: resnet18
+  layer: layer3
+  pooling_kernel_size: 4
   pca_level: 0.97
   score_type: fre # nll: for Gaussian modeling, fre: pca feature reconstruction error
   project_path: ./results

anomalib/models/dfm/dfm_model.py

@@ -17,6 +17,7 @@
 import math
 
 import torch
+import torch.nn.functional as F
 import torchvision
 from torch import Tensor, nn
 
@@ -85,18 +86,23 @@ class DFMModel(nn.Module):
 
     Args:
         backbone (str): Pre-trained model backbone.
+        layer (str): Layer from which to extract features.
+        pool (int): _description_
         n_comps (float, optional): Ratio from which number of components for PCA are calculated. Defaults to 0.97.
         score_type (str, optional): Scoring type. Options are `fre` and `nll`. Defaults to "fre".
     """
 
-    def __init__(self, backbone: str, n_comps: float = 0.97, score_type: str = "fre"):
+    def __init__(
+        self, backbone: str, layer: str, pooling_kernel_size: int, n_comps: float = 0.97, score_type: str = "fre"
+    ):
         super().__init__()
         self.backbone = getattr(torchvision.models, backbone)
+        self.pooling_kernel_size = pooling_kernel_size
         self.n_components = n_comps
         self.pca_model = PCA(n_components=self.n_components)
         self.gaussian_model = SingleClassGaussian()
         self.score_type = score_type
-        self.feature_extractor = FeatureExtractor(backbone=self.backbone(pretrained=True), layers=["avgpool"]).eval()
+        self.feature_extractor = FeatureExtractor(backbone=self.backbone(pretrained=True), layers=[layer]).eval()
 
     def fit(self, dataset: Tensor) -> None:
         """Fit a pca transformation and a Gaussian model to dataset.
@@ -142,9 +148,15 @@ def get_features(self, batch: Tensor) -> Tensor:
             Tensor: Tensor containing extracted features.
         """
         self.feature_extractor.eval()
-        layer_outputs = self.feature_extractor(batch)
-        layer_outputs = torch.cat(list(layer_outputs.values())).detach()
-        return layer_outputs
+        features = self.feature_extractor(batch)
+        for layer in features:
+            batch_size = len(features[layer])
+            if self.pooling_kernel_size > 1:
+                features[layer] = F.avg_pool2d(input=features[layer], kernel_size=self.pooling_kernel_size)
+            features[layer] = features[layer].view(batch_size, -1)
+
+        features = torch.cat(list(features.values())).detach()
+        return features
 
     def forward(self, batch: Tensor) -> Tensor:
         """Computer score from input images.

anomalib/models/dfm/model.py

@@ -32,7 +32,11 @@ def __init__(self, hparams: Union[DictConfig, ListConfig]):
         super().__init__(hparams)
 
         self.model: DFMModel = DFMModel(
-            backbone=hparams.model.backbone, n_comps=hparams.model.pca_level, score_type=hparams.model.score_type
+            backbone=hparams.model.backbone,
+            layer=hparams.model.layer,
+            pooling_kernel_size=hparams.model.pooling_kernel_size,
+            n_comps=hparams.model.pca_level,
+            score_type=hparams.model.score_type,
         )
         self.embeddings: List[Tensor] = []