Spaces:
Build error
Build error
File size: 6,161 Bytes
c8c12e9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 |
"""PyTorch model for DFM model implementation."""
# Copyright (C) 2020 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions
# and limitations under the License.
import math
import torch
import torch.nn.functional as F
import torchvision
from torch import Tensor, nn
from anomalib.models.components import PCA, DynamicBufferModule, FeatureExtractor
class SingleClassGaussian(DynamicBufferModule):
"""Model Gaussian distribution over a set of points."""
def __init__(self):
super().__init__()
self.register_buffer("mean_vec", Tensor())
self.register_buffer("u_mat", Tensor())
self.register_buffer("sigma_mat", Tensor())
self.mean_vec: Tensor
self.u_mat: Tensor
self.sigma_mat: Tensor
def fit(self, dataset: Tensor) -> None:
"""Fit a Gaussian model to dataset X.
Covariance matrix is not calculated directly using:
``C = X.X^T``
Instead, it is represented in terms of the Singular Value Decomposition of X:
``X = U.S.V^T``
Hence,
``C = U.S^2.U^T``
This simplifies the calculation of the log-likelihood without requiring full matrix inversion.
Args:
dataset (Tensor): Input dataset to fit the model.
"""
num_samples = dataset.shape[1]
self.mean_vec = torch.mean(dataset, dim=1)
data_centered = (dataset - self.mean_vec.reshape(-1, 1)) / math.sqrt(num_samples)
self.u_mat, self.sigma_mat, _ = torch.linalg.svd(data_centered, full_matrices=False)
def score_samples(self, features: Tensor) -> Tensor:
"""Compute the NLL (negative log likelihood) scores.
Args:
features (Tensor): semantic features on which density modeling is performed.
Returns:
nll (Tensor): Torch tensor of scores
"""
features_transformed = torch.matmul(features - self.mean_vec, self.u_mat / self.sigma_mat)
nll = torch.sum(features_transformed * features_transformed, dim=1) + 2 * torch.sum(torch.log(self.sigma_mat))
return nll
def forward(self, dataset: Tensor) -> None:
"""Provides the same functionality as `fit`.
Transforms the input dataset based on singular values calculated earlier.
Args:
dataset (Tensor): Input dataset
"""
self.fit(dataset)
class DFMModel(nn.Module):
"""Model for the DFM algorithm.
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, 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=[layer]).eval()
def fit(self, dataset: Tensor) -> None:
"""Fit a pca transformation and a Gaussian model to dataset.
Args:
dataset (Tensor): Input dataset to fit the model.
"""
self.pca_model.fit(dataset)
features_reduced = self.pca_model.transform(dataset)
self.gaussian_model.fit(features_reduced.T)
def score(self, features: Tensor) -> Tensor:
"""Compute scores.
Scores are either PCA-based feature reconstruction error (FRE) scores or
the Gaussian density-based NLL scores
Args:
features (torch.Tensor): semantic features on which PCA and density modeling is performed.
Returns:
score (Tensor): numpy array of scores
"""
feats_projected = self.pca_model.transform(features)
if self.score_type == "nll":
score = self.gaussian_model.score_samples(feats_projected)
elif self.score_type == "fre":
feats_reconstructed = self.pca_model.inverse_transform(feats_projected)
score = torch.sum(torch.square(features - feats_reconstructed), dim=1)
else:
raise ValueError(f"unsupported score type: {self.score_type}")
return score
def get_features(self, batch: Tensor) -> Tensor:
"""Extract features from the pretrained network.
Args:
batch (Tensor): Image batch.
Returns:
Tensor: Tensor containing extracted features.
"""
self.feature_extractor.eval()
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.
Args:
batch (Tensor): Input images
Returns:
Tensor: Scores
"""
feature_vector = self.get_features(batch)
return self.score(feature_vector.view(feature_vector.shape[:2]))
|