DFKDE refactor to accept any layer name like other models (#482)

* DFKDE refactor to accept any layer name like other models * Removed unused variables
openvinotoolkit · Aug 4, 2022 · 381360e · 381360e
1 parent e19428f
commit 381360e
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 1 deletion.
diff --git a/anomalib/models/dfkde/config.yaml b/anomalib/models/dfkde/config.yaml
@@ -24,6 +24,8 @@ model:
   threshold_steepness: 0.05
   threshold_offset: 12
   normalization_method: min_max # options: [null, min_max, cdf]
+  layers:
+    - avgpool
 
 metrics:
   image:

diff --git a/anomalib/models/dfkde/lightning_model.py b/anomalib/models/dfkde/lightning_model.py
@@ -47,6 +47,7 @@ class Dfkde(AnomalyModule):
 
     def __init__(
         self,
+        layers: List[str],
         backbone: str,
         pre_trained: bool = True,
         max_training_points: int = 40000,
@@ -58,6 +59,7 @@ def __init__(
         super().__init__()
 
         self.model = DfkdeModel(
+            layers=layers,
             backbone=backbone,
             pre_trained=pre_trained,
             n_comps=n_components,
@@ -126,6 +128,7 @@ class DfkdeLightning(Dfkde):
 
     def __init__(self, hparams: Union[DictConfig, ListConfig]) -> None:
         super().__init__(
+            layers=hparams.model.layers,
             backbone=hparams.model.backbone,
             max_training_points=hparams.model.max_training_points,
             pre_processing=hparams.model.pre_processing,

diff --git a/anomalib/models/dfkde/torch_model.py b/anomalib/models/dfkde/torch_model.py
@@ -43,6 +43,7 @@ class DfkdeModel(nn.Module):
 
     def __init__(
         self,
+        layers: List[str],
         backbone: str,
         pre_trained: bool = True,
         n_comps: int = 16,
@@ -59,7 +60,7 @@ def __init__(
         self.threshold_offset = threshold_offset
 
         _backbone = getattr(torchvision.models, backbone)
-        self.feature_extractor = FeatureExtractor(backbone=_backbone(pretrained=pre_trained), layers=["avgpool"]).eval()
+        self.feature_extractor = FeatureExtractor(backbone=_backbone(pretrained=pre_trained), layers=layers).eval()
 
         self.pca_model = PCA(n_components=self.n_components)
         self.kde_model = GaussianKDE()

diff --git a/configs/model/dfkde.yaml b/configs/model/dfkde.yaml
@@ -24,6 +24,8 @@ model:
   init_args:
     backbone: resnet18
     pre_trained: true
+    layers:
+      - avgpool
     max_training_points: 40000
     pre_processing: scale
     n_components: 16