Added support for gpu

Signed-off-by: Devansh Agarwal <[email protected]>

Author: code-dev05

src/anomalib/models/image/glass/__init__.py

@@ -1,4 +1,23 @@
+"""GLASS - Unsupervised anomaly detection via Gradient Ascent for Industrial Anomaly detection and localization.
+
+This module implements the GLASS model for unsupervised anomaly detection and localization. GLASS synthesizes both
+global and local anomalies using Gaussian noise guided by gradient ascent to enhance weak defect detection in
+industrial settings.
+
+The model consists of:
+    - A feature extractor and feature adaptor to obtain robust normal representations
+    - A Global Anomaly Synthesis (GAS) module that perturbs features using Gaussian noise and gradient ascent with
+      truncated projection
+    - A Local Anomaly Synthesis (LAS) module that overlays augmented textures onto images using Perlin noise masks
+    - A shared discriminator trained with features from normal, global, and local synthetic samples
+
+Paper: `A Unified Anomaly Synthesis Strategy with Gradient Ascent for Industrial Anomaly Detection and Localization
+<https://arxiv.org/pdf/2407.09359>`
+"""
+
 # Copyright (C) 2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from .lightning_model import Glass as Glass
+from .lightning_model import Glass
+
+__all__ = ["Glass"]

src/anomalib/models/image/glass/lightning_model.py

@@ -1,10 +1,13 @@
 """GLASS - Unsupervised anomaly detection via Gradient Ascent for Industrial Anomaly detection and localization.
 
-This module implements the GLASS model for unsupervised anomaly detection and localization. GLASS synthesizes both global and local anomalies using Gaussian noise guided by gradient ascent to enhance weak defect detection in industrial settings.
+This module implements the GLASS model for unsupervised anomaly detection and localization. GLASS synthesizes both
+global and local anomalies using Gaussian noise guided by gradient ascent to enhance weak defect detection in
+industrial settings.
 
 The model consists of:
     - A feature extractor and feature adaptor to obtain robust normal representations
-    - A Global Anomaly Synthesis (GAS) module that perturbs features using Gaussian noise and gradient ascent with truncated projection
+    - A Global Anomaly Synthesis (GAS) module that perturbs features using Gaussian noise and gradient ascent with
+      truncated projection
     - A Local Anomaly Synthesis (LAS) module that overlays augmented textures onto images using Perlin noise masks
     - A shared discriminator trained with features from normal, global, and local synthetic samples
 
@@ -21,7 +24,7 @@
 import torch
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from torch import nn, optim
-from torch.nn import functional as F
+from torch.nn import functional as f
 from torchvision.transforms.v2 import CenterCrop, Compose, Normalize, Resize
 
 from anomalib import LearningType
@@ -40,17 +43,21 @@
 class Glass(AnomalibModule):
     """PyTorch Lightning Implementation of the GLASS Model.
 
-    The model uses a pre-trained feature extractor to extract features and a feature adaptor to mitigate latent domain bias.
+    The model uses a pre-trained feature extractor to extract features and a feature adaptor to mitigate latent domain
+      bias.
     Global anomaly features are synthesized from adapted normal features using gradient ascent.
-    Local anomaly images are synthesized using texture overlay datasets like dtd which are then processed by feature extractor and feature adaptor.
+    Local anomaly images are synthesized using texture overlay datasets like dtd which are then processed by feature
+      extractor and feature adaptor.
     All three different features are passed to the discriminator trained using loss functions.
 
     Args:
-        input_shape (tuple[int, int]): Input image dimensions as a tuple of (height, width). Required for shaping the input pipeline.
-        anomaly_source_path (str): Path to the dataset or source directory containing normal images and anomaly textures.
+        input_shape (tuple[int, int]): Input image dimensions as a tuple of (height, width). Required for shaping the
+          input pipeline.
+        anomaly_source_path (str): Path to the dataset or source directory containing normal images and anomaly textures
         backbone (str, optional): Name of the CNN backbone used for feature extraction.
             Defaults to `"resnet18"`.
-        pretrain_embed_dim (int, optional): Dimensionality of features extracted by the pre-trained backbone before adaptation.
+        pretrain_embed_dim (int, optional): Dimensionality of features extracted by the pre-trained backbone before
+          adaptation.
             Defaults to `1024`.
         target_embed_dim (int, optional): Dimensionality of the target adapted features after projection.
             Defaults to `1024`.
@@ -62,31 +69,37 @@ class Glass(AnomalibModule):
             Defaults to `True`.
         layers (list[str], optional): List of backbone layers to extract features from.
             Defaults to `["layer1", "layer2", "layer3"]`.
-        pre_proj (int, optional): Number of projection layers used in the feature adaptor (e.g., MLP before discriminator).
+        pre_proj (int, optional): Number of projection layers used in the feature adaptor (e.g., MLP before
+          discriminator).
             Defaults to `1`.
         dsc_layers (int, optional): Number of layers in the discriminator network.
             Defaults to `2`.
         dsc_hidden (int, optional): Number of hidden units in each discriminator layer.
             Defaults to `1024`.
-        dsc_margin (float, optional): Margin used for contrastive or binary classification loss in discriminator training.
+        dsc_margin (float, optional): Margin used for contrastive or binary classification loss in discriminator
+          training.
             Defaults to `0.5`.
         pre_processor (PreProcessor | bool, optional): reprocessing module or flag to enable default preprocessing.
             Set to `True` to apply default normalization and resizing.
             Defaults to `True`.
-        post_processor (PostProcessor | bool, optional): Postprocessing module or flag to enable default output smoothing or thresholding.
+        post_processor (PostProcessor | bool, optional): Postprocessing module or flag to enable default output
+          smoothing or thresholding.
             Defaults to `True`.
         evaluator (Evaluator | bool, optional): Evaluation module for calculating metrics such as AUROC and PRO.
             Defaults to `True`.
-        visualizer (Visualizer | bool, optional): Visualization module to generate heatmaps, segmentation overlays, and anomaly scores.
+        visualizer (Visualizer | bool, optional): Visualization module to generate heatmaps, segmentation overlays, and
+          anomaly scores.
             Defaults to `True`.
-        mining (int, optional): Number of iterations or difficulty level for Online Hard Example Mining (OHEM) during training.
+        mining (int, optional): Number of iterations or difficulty level for Online Hard Example Mining (OHEM) during
+          training.
             Defaults to `1`.
         noise (float, optional): Standard deviation of Gaussian noise used in feature-level anomaly synthesis.
             Defaults to `0.015`.
         radius (float, optional): Radius parameter used for truncated projection in the anomaly synthesis strategy.
             Determines the range for valid synthetic anomalies in the hypersphere or manifold.
             Defaults to `0.75`.
-        p (float, optional): Probability used in random selection logic, such as anomaly mask generation or augmentation choice.
+        p (float, optional): Probability used in random selection logic, such as anomaly mask generation or augmentation
+          choice.
             Defaults to `0.5`.
         lr (float, optional): Learning rate for training the feature adaptor and discriminator networks.
             Defaults to `0.0001`.
@@ -106,7 +119,7 @@ def __init__(
         patchsize: int = 3,
         patchstride: int = 1,
         pre_trained: bool = True,
-        layers: list[str] = ["layer1", "layer2", "layer3"],
+        layers: list[str] | None = None,
         pre_proj: int = 1,
         dsc_layers: int = 2,
         dsc_hidden: int = 1024,
@@ -122,14 +135,17 @@ def __init__(
         lr: float = 0.0001,
         step: int = 20,
         svd: int = 0,
-    ):
+    ) -> None:
         super().__init__(
             pre_processor=pre_processor,
             post_processor=post_processor,
             evaluator=evaluator,
             visualizer=visualizer,
         )
 
+        if layers is None:
+            layers = ["layer1", "layer2", "layer3"]
+
         self.augmentor = PerlinAnomalyGenerator(anomaly_source_path)
 
         self.model = GlassModel(
@@ -157,6 +173,8 @@ def __init__(
         self.step = step
         self.svd = svd
 
+        self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
         self.focal_loss = FocalLoss()
 
         if pre_proj > 0:
@@ -170,7 +188,7 @@ def __init__(
 
         if not pre_trained:
             self.backbone_opt = optim.AdamW(
-                self.model.foward_modules["feature_aggregator"].backbone.parameters(),
+                self.mosdel.forward_modules["feature_aggregator"].backbone.parameters(),
                 self.lr,
             )
         else:
@@ -182,6 +200,30 @@ def configure_pre_processor(
         image_size: tuple[int, int] | None = None,
         center_crop_size: tuple[int, int] | None = None,
     ) -> PreProcessor:
+        """Configure the default pre-processor for GLASS.
+
+        If valid center_crop_size is provided, the pre-processor will
+        also perform center cropping, according to the paper.
+
+        Args:
+            image_size (tuple[int, int] | None, optional): Target size for
+                resizing. Defaults to ``(256, 256)``.
+            center_crop_size (tuple[int, int] | None, optional): Size for center
+                cropping. Defaults to ``None``.
+
+        Returns:
+            PreProcessor: Configured pre-processor instance.
+
+        Raises:
+            ValueError: If at least one dimension of ``center_crop_size`` is larger
+                than correspondent ``image_size`` dimension.
+
+        Example:
+            >>> pre_processor = Glass.configure_pre_processor(
+            ...     image_size=(256, 256)
+            ... )
+            >>> transformed_image = pre_processor(image)
+        """
         image_size = image_size or (256, 256)
 
         if center_crop_size is not None:
@@ -201,10 +243,13 @@ def configure_pre_processor(
 
         return PreProcessor(transform=transform)
 
-    def configure_optimizers(self) -> list[optim.Optimizer]:
-        dsc_opt = optim.AdamW(self.model.discriminator.parameters(), lr=self.lr * 2)
+    def configure_optimizers(self) -> optim.Optimizer:
+        """Configure optimizer for the discriminator.
 
-        return dsc_opt
+        Returns:
+            Optimizer: AdamW Optimizer for the discriminator.
+        """
+        return optim.AdamW(self.model.discriminator.parameters(), lr=self.lr * 2)
 
     def training_step(
         self,
@@ -220,6 +265,7 @@ def training_step(
         Returns:
             STEP_OUTPUT: Dictionary containing loss values and metrics
         """
+        del batch_idx
         dsc_opt = self.optimizers()
 
         self.model.forward_modules.eval()
@@ -235,21 +281,22 @@ def training_step(
 
         img = batch.image
         aug, mask_s = self.augmentor(img)
-        batch_size = img.shape[0]
+        if img is not None:
+            batch_size = img.shape[0]
 
         true_feats, fake_feats = self.model(img, aug)
 
         h_ratio = mask_s.shape[2] // int(math.sqrt(fake_feats.shape[0] // batch_size))
         w_ratio = mask_s.shape[3] // int(math.sqrt(fake_feats.shape[0] // batch_size))
 
-        mask_s_resized = F.interpolate(
+        mask_s_resized = f.interpolate(
             mask_s.float(),
             size=(mask_s.shape[2] // h_ratio, mask_s.shape[3] // w_ratio),
             mode="nearest",
         )
         mask_s_gt = mask_s_resized.reshape(-1, 1)
 
-        noise = torch.normal(0, self.noise, true_feats.shape)
+        noise = torch.normal(0, self.noise, true_feats.shape).to(self.dev)
         gaus_feats = true_feats + noise
 
         center = self.c.repeat(img.shape[0], 1, 1)
@@ -260,7 +307,7 @@ def training_step(
         )
         c_t_points = torch.concat([center[mask_s_gt[:, 0] == 0], center], dim=0)
         dist_t = torch.norm(true_points - c_t_points, dim=1)
-        r_t = torch.tensor([torch.quantile(dist_t, q=self.radius)]).to(self.device)
+        r_t = torch.tensor([torch.quantile(dist_t, q=self.radius)]).to(self.dev)
 
         for step in range(self.step + 1):
             scores = self.model.discriminator(torch.cat([true_feats, gaus_feats]))
@@ -272,10 +319,6 @@ def training_step(
 
             if step == self.step:
                 break
-            if self.mining == 0:
-                dist_g = torch.norm(gaus_feats - center, dim=1)
-                r_g = torch.tensor([torch.quantile(dist_g, q=self.radius)])
-                break
 
             grad = torch.autograd.grad(gaus_loss, [gaus_feats])[0]
             grad_norm = torch.norm(grad, dim=1)
@@ -326,7 +369,7 @@ def training_step(
         self.log("true_loss", true_loss, prog_bar=True)
         self.log("gaus_loss", gaus_loss, prog_bar=True)
         self.log("bce_loss", bce_loss, prog_bar=True)
-        self.log("focal_losss", focal_loss, prog_bar=True)
+        self.log("focal_loss", focal_loss, prog_bar=True)
         self.log("loss", loss, prog_bar=True)
 
     def on_train_start(self) -> None:
@@ -340,9 +383,9 @@ def on_train_start(self) -> None:
         with torch.no_grad():
             for i, batch in enumerate(dataloader):
                 if i == 0:
-                    self.c = self.model.calculate_mean(batch.image)
+                    self.c = self.model.calculate_mean(batch.image.to(self.dev))
                 else:
-                    self.c += self.model.calculate_mean(batch.image)
+                    self.c += self.model.calculate_mean(batch.image.to(self.dev))
 
             self.c /= len(dataloader)