WinCLIP improvements

src/anomalib/models/__init__.py

@@ -29,6 +29,7 @@
 from anomalib.models.reverse_distillation import ReverseDistillation
 from anomalib.models.rkde import Rkde
 from anomalib.models.stfpm import Stfpm
+from anomalib.models.win_clip import WinClip
 
 __all__ = [
     "Cfa",
@@ -46,6 +47,7 @@
     "Stfpm",
     "AiVad",
     "EfficientAd",
+    "WinClip"
 ]
 
 logger = logging.getLogger(__name__)

src/anomalib/models/win_clip/__init__.py

@@ -0,0 +1,2 @@
+from .torch_model import *
+from .lightning_model import WinClip, WinClipLightning

src/anomalib/models/win_clip/config.yaml

@@ -0,0 +1,105 @@
+dataset:
+  name: mvtec
+  format: mvtec
+  path: ./datasets/MVTec
+  category: hazelnut
+  task: segmentation
+  train_batch_size: 32
+  eval_batch_size: 32
+  num_workers: 8
+  image_size: 240 # dimensions to which images are resized (mandatory)
+  center_crop: null # dimensions to which images are center-cropped after resizing (optional)
+  normalization: imagenet # data distribution to which the images will be normalized: [none, imagenet]
+  transform_config:
+    train: null
+    eval: null
+  test_split_mode: from_dir # options: [from_dir, synthetic]
+  test_split_ratio: 0.2 # fraction of train images held out testing (usage depends on test_split_mode)
+  val_split_mode: same_as_test # options: [same_as_test, from_test, synthetic]
+  val_split_ratio: 0.5 # fraction of train/test images held out for validation (usage depends on val_split_mode)
+  tiling:
+    apply: false
+    tile_size: null
+    stride: null
+    remove_border_count: 0
+    use_random_tiling: False
+    random_tile_count: 16
+
+model:
+  name: win_clip
+  class_name: object
+  n_shot: 0
+  normalization_method: min_max
+
+metrics:
+  image:
+    - F1Score
+    - AUROC
+  pixel:
+    - F1Score
+    - AUROC
+  threshold:
+    method: adaptive #options: [adaptive, manual]
+    manual_image: null
+    manual_pixel: null
+
+visualization:
+  show_images: False # show images on the screen
+  save_images: True # save images to the file system
+  log_images: True # log images to the available loggers (if any)
+  image_save_path: null # path to which images will be saved
+  mode: full # options: ["full", "simple"]
+
+project:
+  seed: 42
+  path: ./results
+
+logging:
+  logger: [] # options: [comet, tensorboard, wandb, csv] or combinations.
+  log_graph: false # Logs the model graph to respective logger.
+
+optimization:
+  export_mode: null # options: torch, onnx, openvino
+
+# PL Trainer Args. Don't add extra parameter here.
+trainer:
+  enable_checkpointing: true
+  default_root_dir: null
+  gradient_clip_val: 0
+  gradient_clip_algorithm: norm
+  num_nodes: 1
+  devices: 1
+  enable_progress_bar: true
+  overfit_batches: 0.0
+  track_grad_norm: -1
+  check_val_every_n_epoch: 1 # Don't validate before extracting features.
+  fast_dev_run: false
+  accumulate_grad_batches: 1
+  max_epochs: 1
+  min_epochs: null
+  max_steps: -1
+  min_steps: null
+  max_time: null
+  limit_train_batches: 0
+  limit_val_batches: 1.0
+  limit_test_batches: 1.0
+  limit_predict_batches: 1.0
+  val_check_interval: 1.0 # Don't validate before extracting features.
+  log_every_n_steps: 50
+  accelerator: auto # <"cpu", "gpu", "tpu", "ipu", "hpu", "auto">
+  strategy: null
+  sync_batchnorm: false
+  precision: 32
+  enable_model_summary: true
+  num_sanity_val_steps: -1
+  profiler: null
+  benchmark: false
+  deterministic: false
+  reload_dataloaders_every_n_epochs: 0
+  auto_lr_find: false
+  replace_sampler_ddp: true
+  detect_anomaly: false
+  auto_scale_batch_size: false
+  plugins: null
+  move_metrics_to_cpu: false
+  multiple_trainloader_mode: max_size_cycle

src/anomalib/models/win_clip/lightning_model.py

@@ -0,0 +1,79 @@
+"""WinCLIP: Zero-/Few-Shot Anomaly Classification and Segmentation
+
+Paper https://arxiv.org/abs/2303.14814
+"""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import torch
+from omegaconf import DictConfig, ListConfig
+from pytorch_lightning.utilities.types import STEP_OUTPUT
+from torch import Tensor
+from torchvision.transforms.functional import resize
+
+from anomalib.models.components import AnomalyModule
+from anomalib.models.win_clip.torch_model import WinClipAD
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["WinClip", "WinClipLightning"]
+
+
+class WinClip(AnomalyModule):
+    """WinCLIP"""
+
+    def __init__(
+        self,
+        class_name: str,
+        n_shot: int,
+    ) -> None:
+        super().__init__()
+        self.model = WinClipAD()
+        self.class_name = class_name
+        self.n_shot = n_shot
+
+    def setup(self, stage) -> None:
+        del stage
+        self.model.build_text_feature_gallery(self.class_name)
+
+    def on_fit_end(self) -> None:
+        model_path = Path(self.trainer.default_root_dir) / "weights/lightning/model.cpkt"
+        self.trainer.save_checkpoint(model_path, weights_only=True)
+
+    @staticmethod
+    def configure_optimizers() -> None:  # pylint: disable=arguments-differ
+        """WinCLIP doesn't require optimization, therefore returns no optimizers."""
+        return None
+
+    def training_step(self, batch: dict[str, str | Tensor], *args, **kwargs) -> None:
+        """Training Step of WinCLIP"""
+        del args, kwargs  # These variables are not used.
+        if self.n_shot:
+            self.model.build_image_feature_gallery(batch["image"][: self.n_shot])
+        return
+
+    def validation_step(self, batch: dict[str, str | Tensor], *args, **kwargs) -> STEP_OUTPUT:
+        """Validation Step of WinCLIP"""
+        del args, kwargs  # These variables are not used.
+        scores = self.model(batch["image"])
+        scores = [resize(torch.tensor(score).unsqueeze(0), batch["image"].shape[-2:]).squeeze() for score in scores]
+        batch["anomaly_maps"] = torch.stack(scores).to(self.device)
+        return batch
+
+
+class WinClipLightning(WinClip):
+    """WinCLIP"""
+
+    def __init__(self, hparams: DictConfig | ListConfig) -> None:
+        super().__init__(
+            class_name=hparams.model.class_name,
+            n_shot=hparams.model.n_shot,
+        )
+        self.hparams: DictConfig | ListConfig  # type: ignore
+        self.save_hyperparameters(hparams)

src/anomalib/models/win_clip/third_party/CLIPAD/__init__.py

@@ -0,0 +1,13 @@
+from .coca_model import CoCa
+from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from .factory import create_model, create_model_and_transforms, create_model_from_pretrained, get_tokenizer, create_loss
+from .factory import list_models, add_model_config, get_model_config, load_checkpoint
+from .loss import ClipLoss, DistillClipLoss, CoCaLoss
+from .model import CLIP, CustomTextCLIP, CLIPTextCfg, CLIPVisionCfg, \
+    convert_weights_to_lp, convert_weights_to_fp16, trace_model, get_cast_dtype
+from .openai import load_openai_model, list_openai_models
+from .pretrained import list_pretrained, list_pretrained_models_by_tag, list_pretrained_tags_by_model, \
+    get_pretrained_url, download_pretrained_from_url, is_pretrained_cfg, get_pretrained_cfg, download_pretrained
+from .push_to_hf_hub import push_pretrained_to_hf_hub, push_to_hf_hub
+from .tokenizer import SimpleTokenizer, tokenize, decode
+from .transform import image_transform, AugmentationCfg