From 3d3712455767b07a814214821b6ed124a8473f99 Mon Sep 17 00:00:00 2001 From: blaz-r Date: Sun, 7 Apr 2024 21:31:38 +0200 Subject: [PATCH] Fix anomaly map shapes to work with tiling Signed-off-by: blaz-r --- src/anomalib/models/image/padim/torch_model.py | 3 ++- src/anomalib/models/image/patchcore/torch_model.py | 3 ++- src/anomalib/models/image/stfpm/torch_model.py | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/anomalib/models/image/padim/torch_model.py b/src/anomalib/models/image/padim/torch_model.py index e24e8a03a2..f45dde1f79 100644 --- a/src/anomalib/models/image/padim/torch_model.py +++ b/src/anomalib/models/image/padim/torch_model.py @@ -126,6 +126,7 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor: torch.Size([32, 128, 28, 28]), torch.Size([32, 256, 14, 14])] """ + output_size = input_tensor.shape[-2:] if self.tiler: input_tensor = self.tiler.tile(input_tensor) @@ -143,7 +144,7 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor: embedding=embeddings, mean=self.gaussian.mean, inv_covariance=self.gaussian.inv_covariance, - image_size=input_tensor.shape[-2:], + image_size=output_size, ) return output diff --git a/src/anomalib/models/image/patchcore/torch_model.py b/src/anomalib/models/image/patchcore/torch_model.py index 5b26a12ee3..a2ceb32b91 100644 --- a/src/anomalib/models/image/patchcore/torch_model.py +++ b/src/anomalib/models/image/patchcore/torch_model.py @@ -70,6 +70,7 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor | dict[str, torch. Returns: Tensor | dict[str, torch.Tensor]: Embedding for training, anomaly map and anomaly score for testing. """ + output_size = input_tensor.shape[-2:] if self.tiler: input_tensor = self.tiler.tile(input_tensor) @@ -98,7 +99,7 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor | dict[str, torch. # reshape to w, h patch_scores = patch_scores.reshape((batch_size, 1, width, height)) # get anomaly map - anomaly_map = self.anomaly_map_generator(patch_scores, input_tensor.shape[-2:]) + anomaly_map = self.anomaly_map_generator(patch_scores, output_size) output = {"anomaly_map": anomaly_map, "pred_score": pred_score} diff --git a/src/anomalib/models/image/stfpm/torch_model.py b/src/anomalib/models/image/stfpm/torch_model.py index 58beabea0e..5b80a6ec7a 100644 --- a/src/anomalib/models/image/stfpm/torch_model.py +++ b/src/anomalib/models/image/stfpm/torch_model.py @@ -61,6 +61,7 @@ def forward(self, images: torch.Tensor) -> torch.Tensor | dict[str, torch.Tensor Returns: Teacher and student features when in training mode, otherwise the predicted anomaly maps. """ + output_size = images.shape[-2:] if self.tiler: images = self.tiler.tile(images) teacher_features: dict[str, torch.Tensor] = self.teacher_model(images) @@ -78,7 +79,7 @@ def forward(self, images: torch.Tensor) -> torch.Tensor | dict[str, torch.Tensor output = self.anomaly_map_generator( teacher_features=teacher_features, student_features=student_features, - image_size=images.shape[-2:], + image_size=output_size, ) return output