Skip to content

Commit

Permalink
Fix anomaly map shape to also work with tiling (#1959)
Browse files Browse the repository at this point in the history
* Fix error message in Tiler

Signed-off-by: blaz-r <blaz.rolih@gmail.com>

* Fix anomaly map shapes to work with tiling

Signed-off-by: blaz-r <blaz.rolih@gmail.com>

---------

Signed-off-by: blaz-r <blaz.rolih@gmail.com>
  • Loading branch information
blaz-r authored Apr 9, 2024
1 parent 96ec6a4 commit e077b1e
Show file tree
Hide file tree
Showing 4 changed files with 7 additions and 4 deletions.
2 changes: 1 addition & 1 deletion src/anomalib/data/utils/tiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -378,7 +378,7 @@ def tile(self, image: torch.Tensor, use_random_tiling: bool = False) -> torch.Te
if self.input_h < self.tile_size_h or self.input_w < self.tile_size_w:
msg = (
f"One of the edges of the tile size {self.tile_size_h, self.tile_size_w} is larger than "
f"that of the image {{self.input_h, self.input_w}}."
f"that of the image {self.input_h, self.input_w}."
)
raise ValueError(
msg,
Expand Down
3 changes: 2 additions & 1 deletion src/anomalib/models/image/padim/torch_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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

Expand Down
3 changes: 2 additions & 1 deletion src/anomalib/models/image/patchcore/torch_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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}

Expand Down
3 changes: 2 additions & 1 deletion src/anomalib/models/image/stfpm/torch_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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

0 comments on commit e077b1e

Please sign in to comment.