Skip to content

Commit

Permalink
Fix transforms for draem, dsr and rkde
Browse files Browse the repository at this point in the history
Signed-off-by: Blaz Rolih <blaz.rolih@gmail.com>
  • Loading branch information
blaz-r committed Sep 25, 2024
1 parent 3768bb3 commit e8cf720
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 0 deletions.
11 changes: 11 additions & 0 deletions src/anomalib/models/image/draem/lightning_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import torch
from lightning.pytorch.utilities.types import STEP_OUTPUT
from torch import nn
from torchvision.transforms.v2 import Compose, Resize, Transform

from anomalib import LearningType
from anomalib.data.utils import Augmenter
Expand Down Expand Up @@ -150,3 +151,13 @@ def learning_type(self) -> LearningType:
LearningType: Learning type of the model.
"""
return LearningType.ONE_CLASS

@staticmethod
def configure_transforms(image_size: tuple[int, int] | None = None) -> Transform:
"""Default transform for DRAEM. Normalization is not needed as the images are scaled to [0, 1] in Dataset."""
image_size = image_size or (256, 256)
return Compose(
[
Resize(image_size, antialias=True),
],
)
11 changes: 11 additions & 0 deletions src/anomalib/models/image/dsr/lightning_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import torch
from lightning.pytorch.utilities.types import STEP_OUTPUT, OptimizerLRScheduler
from torch import Tensor
from torchvision.transforms.v2 import Compose, Resize, Transform

from anomalib import LearningType
from anomalib.data.utils import DownloadInfo, download_and_extract
Expand Down Expand Up @@ -191,3 +192,13 @@ def learning_type(self) -> LearningType:
LearningType: Learning type of the model.
"""
return LearningType.ONE_CLASS

@staticmethod
def configure_transforms(image_size: tuple[int, int] | None = None) -> Transform:
"""Default transform for DSR. Normalization is not needed as the images are scaled to [0, 1] in Dataset."""
image_size = image_size or (256, 256)
return Compose(
[
Resize(image_size, antialias=True),
],
)
11 changes: 11 additions & 0 deletions src/anomalib/models/image/rkde/lightning_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@

import torch
from lightning.pytorch.utilities.types import STEP_OUTPUT
from torchvision.transforms.v2 import Compose, Resize, Transform

from anomalib import LearningType
from anomalib.models.components import AnomalyModule, MemoryBankMixin
Expand Down Expand Up @@ -143,3 +144,13 @@ def learning_type(self) -> LearningType:
LearningType: Learning type of the model.
"""
return LearningType.ONE_CLASS

@staticmethod
def configure_transforms(image_size: tuple[int, int] | None = None) -> Transform:
"""Default transform for RKDE."""
image_size = image_size or (240, 360)
return Compose(
[
Resize(image_size, antialias=True),
],
)

0 comments on commit e8cf720

Please sign in to comment.