Skip to content

Commit

Permalink
fix: Replace kornia gaussian filter which fails when exporting to ONNX
Browse files Browse the repository at this point in the history
  • Loading branch information
Carlos Nácher committed Apr 16, 2024
1 parent 75d7cea commit ef6b3ac
Showing 1 changed file with 10 additions and 5 deletions.
15 changes: 10 additions & 5 deletions src/anomalib/models/reverse_distillation/anomaly_map.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@

import torch
import torch.nn.functional as F
from kornia.filters import gaussian_blur2d
# from kornia.filters import gaussian_blur2d
from anomalib.models.components import GaussianBlur2d
from omegaconf import ListConfig
from torch import Tensor, nn

Expand Down Expand Up @@ -84,8 +85,12 @@ def forward(self, student_features: list[Tensor], teacher_features: list[Tensor]
elif self.mode == AnomalyMapGenerationMode.ADD:
anomaly_map += distance_map

anomaly_map = gaussian_blur2d(
anomaly_map, kernel_size=(self.kernel_size, self.kernel_size), sigma=(self.sigma, self.sigma)
)

# anomaly_map = gaussian_blur2d(
# anomaly_map, kernel_size=(self.kernel_size, self.kernel_size), sigma=(self.sigma, self.sigma)
# )
anomaly_map = GaussianBlur2d(
kernel_size=(self.kernel_size, self.kernel_size),
sigma=(self.sigma, self.sigma),
)(anomaly_map)

return anomaly_map

0 comments on commit ef6b3ac

Please sign in to comment.