Allow specifying feature layer and pool factor in DFM

anomalib/models/dfm/config.yaml

@@ -14,6 +14,8 @@ dataset:
 model:
   name: dfm
   backbone: resnet18
+  layer: layer3
+  pool: 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

@@ -19,6 +19,7 @@
 import torch
 import torchvision
 from torch import Tensor, nn
+import torch.nn.functional as F
 
 from anomalib.models.components import PCA, DynamicBufferModule, FeatureExtractor
 
@@ -89,14 +90,15 @@ class DFMModel(nn.Module):
         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, pool: int, n_comps: float = 0.97, score_type: str = "fre"):
         super().__init__()
         self.backbone = getattr(torchvision.models, backbone)
+        self.pool = pool
         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.
@@ -143,6 +145,12 @@ def get_features(self, batch: Tensor) -> Tensor:
         """
         self.feature_extractor.eval()
         layer_outputs = self.feature_extractor(batch)
+        for k in layer_outputs:
+            s0 = len(layer_outputs[k])
+            if self.pool > 1:
+                layer_outputs[k] = F.avg_pool2d(layer_outputs[k], self.pool)
+            layer_outputs[k] = layer_outputs[k].view(s0, -1)
+
         layer_outputs = torch.cat(list(layer_outputs.values())).detach()
         return layer_outputs
 

anomalib/models/dfm/model.py

@@ -32,7 +32,8 @@ 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, pool=hparams.model.pool, 
+            n_comps=hparams.model.pca_level, score_type=hparams.model.score_type
         )
         self.automatic_optimization = False
         self.embeddings: List[Tensor] = []