Upload 48 files

AnomalyCLIP_lib/AnomalyCLIP.py

@@ -0,0 +1,531 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+from torch import nn
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu1(self.bn1(self.conv1(x)))
+        out = self.relu2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu3(out)
+        return out
+
+
+# implement attention module for v-v self-attention
+class Attention(nn.Module):
+    def __init__(self, out_dim, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., settings=''):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(out_dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.settings = settings
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        # original self-attention for the original path
+        attn_ori = (q @ k.transpose(-2, -1)) * self.scale
+        attn_ori = attn_ori.softmax(dim=-1)
+        attn_ori = self.attn_drop(attn_ori)
+
+        # replace k & q by v
+        k = v
+        q = k
+
+        # self-attention, higher temperate for resnets performs better
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = (attn).softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x_ori = (attn_ori @ v).transpose(1, 2).reshape(B, N, C)
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj_drop(self.proj(x))
+        x_ori = self.proj_drop(self.proj(x_ori))
+        return [x, x_ori]
+
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None, design_details = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        if isinstance(self.attn, Attention):
+            x = x.transpose(0, 1)
+            x, x_ori = self.attn(x)
+            return [x.transpose(0, 1), x_ori.transpose(0, 1)]
+        else:
+            return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x, whole = False, ffn = False):
+        # print("xxxxx",x.shape)
+        # dual paths for blocks deeper than "d"
+        
+        if isinstance(self.attn, Attention):
+            if isinstance(x, list):
+                if not ffn:
+                    x, x_ori = x
+                    x_res = self.attention(self.ln_1(x_ori))
+                    x_res, x_ori_res = x_res
+                    x_ori += x_ori_res
+                    x_ori = x_ori + self.mlp(self.ln_2(x_ori))
+                    x += x_res # skip ffn for the new path
+                    # print('hellloooo')
+                    return [x, x_ori]
+                else:
+                    x, x_ori_1 = x
+                    x_res = self.attention(self.ln_1(x_ori_1))
+                    x_res, x_ori_res = x_res
+                    x_ori = x_ori_1 +  x_ori_res
+                    x_ori = x_ori + self.mlp(self.ln_2(x_ori))
+                    x += x_res # skip ffn for the new path
+                    x = x_res + x_ori_1
+                    x = x + self.mlp(self.ln_2(x))
+                    return [x, x_ori]
+            # start of dual path
+            else:
+                x_res = self.attention(self.ln_1(x))
+                if isinstance(x_res, list):
+                    x_res, x_ori_res = x_res
+                    x_ori = x + x_ori_res
+                    x_ori = x_ori + self.mlp(self.ln_2(x_ori))
+                    x += x_res
+                    return [x, x_ori]
+
+        # singl path before "d"
+        else:
+            x = x + self.attention(self.ln_1(x))
+            x = x + self.mlp(self.ln_2(x))
+        return x
+
+class ResidualAttentionBlock_learnable_token(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None, design_details=None,
+            text_layer=False, i = 0):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+        
+        self.i = i
+        self.compound_prompt_nctx = design_details['learnabel_text_embedding_length']
+        self.text_layer = text_layer
+        if i == 0:
+            self.first_layer = True
+        else:
+            self.first_layer = False
+    
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        if isinstance(self.attn, Attention):
+            x = x.transpose(0, 1)
+            x, x_ori = self.attn(x)
+            return [x.transpose(0, 1), x_ori.transpose(0, 1)]
+        else:
+            return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, inputs):
+
+        # dual paths for blocks deeper than "d"
+        if isinstance(self.attn, Attention):
+            x = inputs[0]
+            if isinstance(x, list):
+                x, x_ori = x
+                x_res = self.attention(self.ln_1(x_ori))
+                x_res, x_ori_res = x_res
+                x_ori += x_ori_res
+                x_ori = x_ori + self.mlp(self.ln_2(x_ori))
+                x += x_res # skip ffn for the new path
+                return [x, x_ori]
+
+            # start of dual path
+            else:
+                x_res = self.attention(self.ln_1(x))
+                if isinstance(x_res, list):
+                    x_res, x_ori_res = x_res
+                    x_ori = x + x_ori_res
+                    x_ori = x_ori + self.mlp(self.ln_2(x_ori))
+                    x += x_res
+                    return [x, x_ori]
+
+        # singl path before "d"
+        else:
+            x = inputs[0]
+            compound_prompts_deeper = inputs[1]
+            counter = inputs[2]
+            if not self.first_layer:
+                # First check if the ith layer needs compound prompts or not
+                if not (counter > len(compound_prompts_deeper) - 1):
+                    # Appending the learnable tokens in different way
+                    # x -> [77, NCLS, DIM]
+                    # First remove the learnable tokens from previous layer
+                    prefix = x[:1, :, :]
+                    suffix = x[1 + self.compound_prompt_nctx:, :, :]
+                    textual_context = compound_prompts_deeper[counter]
+                    textual_context = textual_context.expand(x.shape[1], -1, -1).permute(1, 0, 2).half()
+                    # Add the learnable tokens of this layer with the input, replaced by previous
+                    # layer learnable tokens
+                    x = torch.cat([prefix, textual_context, suffix], dim=0)
+                    # Once done, update the counter, so that the next time, it does not use same learnable tokens
+                    counter += 1
+            x = x + self.attention(self.ln_1(x))
+            x = x + self.mlp(self.ln_2(x))
+        return [x, compound_prompts_deeper, counter]
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None, need_weights: bool = False, design_details = None ,text_layer = False):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.text_layer = text_layer
+        self.design_deatails = design_details
+        print("text_layer", self.text_layer)
+        if self.text_layer and (design_details is not None):
+            self.resblocks = nn.ModuleList([ResidualAttentionBlock_learnable_token(width, heads, attn_mask, design_details, text_layer, i=i) for i in range(layers)])
+        else:
+            self.resblocks = nn.ModuleList([ResidualAttentionBlock(width, heads, attn_mask,) for i in range(layers)])
+
+    def ori_CLIP_with_patch_forward(self, x, out_layers):
+        idx = 0
+        out_tokens = []
+        for r in self.resblocks:
+            idx += 1
+            x = r(x)
+            if idx in out_layers:
+                if isinstance(x, list):
+                    out_tokens.append(x[1])
+                else:
+                    out_tokens.append(x)
+
+        return [x, x], out_tokens
+
+    def AnomalyCLIP_forward(self, x, out_layers, ffn):
+        idx = 0
+        out_tokens = []
+        for r in self.resblocks:
+            idx += 1
+            x = r(x, ffn = ffn)
+            # print("out_layers", out_layers, idx)
+            if idx in out_layers:
+                if isinstance(x, list):
+                    out_tokens.append(x[0])
+                else:
+                    out_tokens.append(x)
+        return x, out_tokens
+
+    def forward(self, x: torch.Tensor, out_layers = [6, 12, 18, 24], DPAM_layer = None, ffn = False):
+        # visual encoder forward
+        if not self.text_layer:
+            out_tokens = []
+
+            if DPAM_layer is None:
+                [x, x], out_tokens = self.ori_CLIP_with_patch_forward(x, out_layers)
+                return [x, x], out_tokens
+            else:
+                x, out_tokens = self.AnomalyCLIP_forward(x, out_layers, ffn)
+                return x, out_tokens
+        # text encoder forward
+        # ori text embedding
+        elif self.design_deatails is None:
+            for idx, r in enumerate(self.resblocks):
+                x = r(x)
+            return x
+        # insert learnable text embedding
+        elif self.design_deatails is not None:
+            for idx, r in enumerate(self.resblocks):
+                x = r(x)
+            return x[0]
+    def get_cast_dtype(self) -> torch.dtype:
+        return self.resblocks[0].mlp.c_fc.weight.dtype
+
+class VisionTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads, need_weights=True)
+        self.attn = None
+        self.embed_dim = width
+        self.num_heads = heads
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+
+    @torch.no_grad()
+    def DAPM_replace(self, DPAM_layer):
+        if DPAM_layer is not None:
+            for i in range(1, DPAM_layer):
+                self.attn = Attention(self.embed_dim, self.embed_dim, self.num_heads, True)
+                self.attn.qkv.weight.data = self.transformer.resblocks[-i].attn.in_proj_weight.clone()
+                self.attn.qkv.bias.data = self.transformer.resblocks[-i].attn.in_proj_bias.clone()
+                self.attn.proj.weight.data = self.transformer.resblocks[-i].attn.out_proj.weight.clone()
+                self.attn.proj.bias.data = self.transformer.resblocks[-i].attn.out_proj.bias.clone()
+                self.transformer.resblocks[-i].attn = self.attn
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor, features_list, ori_patch = False, proj_use = True, DPAM_layer = None, ffn = False):
+
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        side = int((self.positional_embedding.shape[0] - 1) ** 0.5)
+        new_side = int((x.shape[1] - 1) ** 0.5)
+
+        # update the position embedding during inference for varied input size
+        if side != new_side:
+            new_pos = self.positional_embedding[1:, :].reshape(-1, side, side, x.shape[-1]).permute(0, 3, 1, 2)
+            new_pos = torch.nn.functional.interpolate(new_pos, (new_side, new_side), mode='bilinear')
+            new_pos = new_pos.reshape(-1, x.shape[-1], new_side * new_side).transpose(1, 2)
+            self.positional_embedding.data = torch.cat([self.positional_embedding[:1, :], new_pos[0]], 0)
+
+        pos = self.positional_embedding.to(x.dtype)
+        x = x + pos
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        [x, x_ori], patch_tokens = self.transformer(x, features_list, DPAM_layer = DPAM_layer, ffn = ffn)
+        
+
+        if True:
+            patch_token_list = []
+            for patch_token in patch_tokens:
+                patch_token = self.ln_post(patch_token.permute(1, 0, 2)) @ self.proj  # LND -> NLD
+                patch_token_list.append(patch_token)
+            patch_tokens = patch_token_list
+
+            return x_ori[0, :, :] @ self.proj, patch_tokens
+
+
+        return x
+
+
+from thop import profile
+class AnomalyCLIP(nn.Module):
+    def __init__(self,
+                 embed_dim: int,
+                 # vision
+                 image_resolution: int,
+                 vision_layers: Union[Tuple[int, int, int, int], int],
+                 vision_width: int,
+                 vision_patch_size: int,
+                 # text
+                 context_length: int,
+                 vocab_size: int,
+                 transformer_width: int,
+                 transformer_heads: int,
+                 transformer_layers: int,
+                 design_details = None
+                 ):
+        super().__init__()
+
+        self.context_length = context_length
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(
+                layers=vision_layers,
+                output_dim=embed_dim,
+                heads=vision_heads,
+                input_resolution=image_resolution,
+                width=vision_width
+            )
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisionTransformer(
+                input_resolution=image_resolution,
+                patch_size=vision_patch_size,
+                width=vision_width,
+                layers=vision_layers,
+                heads=vision_heads,
+                output_dim=embed_dim
+            )
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask(), text_layer=True, design_details=design_details
+        )
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image, feature_list = [], ori_patch = False, proj_use = True, DPAM_layer = None, ffn = False):
+        return self.visual(image.type(self.dtype), feature_list, ori_patch = ori_patch, proj_use = proj_use, DPAM_layer = DPAM_layer, ffn = ffn)
+
+
+    def encode_text(self, text):
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+        return x
+
+    def encode_text_learn(self, prompts, tokenized_prompts, deep_compound_prompts_text = None, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        # x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+        
+        # x = x + self.positional_embedding.to(cast_dtype)
+
+        x = prompts + self.positional_embedding.to(cast_dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        # print("test", x.shape, len(deep_compound_prompts_text))
+        if deep_compound_prompts_text is None:
+            x = self.transformer(x)
+        else:
+            x = self.transformer([x, deep_compound_prompts_text, 0])
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)  # [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), tokenized_prompts.argmax(dim=-1)] @ self.text_projection
+        return x
+    
+    def forward(self, image, text):
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logits_per_image.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        return logits_per_image, logits_per_text

AnomalyCLIP_lib/CLIP.py

@@ -0,0 +1,436 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu1(self.bn1(self.conv1(x)))
+        out = self.relu2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu3(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+
+        side = int((self.positional_embedding.shape[0] - 1) ** 0.5)
+        new_side = int((x.shape[0] - 1) ** 0.5)
+
+        # update the position embedding during inference for varied input size
+        if side != new_side:
+            new_pos = self.positional_embedding[1:, :].reshape(-1, side, side, x.shape[-1]).permute(0, 3, 1, 2)
+            new_pos = torch.nn.functional.interpolate(new_pos, (new_side, new_side), mode='bilinear')
+            new_pos = new_pos.reshape(-1, x.shape[-1], new_side * new_side).transpose(1, 2)
+            self.positional_embedding.data = torch.cat([self.positional_embedding[:1, :], new_pos[0]], 0)
+
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x, key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+
+        #return x[0]
+        return x.transpose(0, 1) # return both cls token and image tokens, B,N,C
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.relu1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.relu3 = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(2)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            x = self.relu1(self.bn1(self.conv1(x)))
+            x = self.relu2(self.bn2(self.conv2(x)))
+            x = self.relu3(self.bn3(self.conv3(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None, need_weights: bool = False):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+        self.need_weights = need_weights
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        if self.need_weights == False:
+            return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+        else:
+            return self.attn(x, x, x, need_weights=True, attn_mask=self.attn_mask)
+
+    def forward(self, x: torch.Tensor):
+        if self.need_weights == False:
+            x = x + self.attention(self.ln_1(x))
+            x = x + self.mlp(self.ln_2(x))
+            return x
+        else:
+            y, attn = self.attention(self.ln_1(x))
+            x = x + y
+            x = x + self.mlp(self.ln_2(x))
+            return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None, need_weights: bool = False):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask, need_weights if i == layers - 1 else False) for i in range(layers)])
+
+    def forward(self, x: torch.Tensor):
+        return self.resblocks(x)
+
+    def get_cast_dtype(self) -> torch.dtype:
+        return self.resblocks[0].mlp.c_fc.weight.dtype
+
+
+
+class VisionTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads, need_weights=True)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        
+        #####################################################################################
+        side = int((self.positional_embedding.shape[0] - 1) ** 0.5)
+        new_side = int((x.shape[1] - 1) ** 0.5)
+
+        # update the position embedding during inference for varied input size
+        if side != new_side:
+            new_pos = self.positional_embedding[1:, :].reshape(-1, side, side, x.shape[-1]).permute(0, 3, 1, 2)
+            new_pos = torch.nn.functional.interpolate(new_pos, (new_side, new_side), mode='bilinear')
+            new_pos = new_pos.reshape(-1, x.shape[-1], new_side * new_side).transpose(1, 2)
+            self.positional_embedding.data = torch.cat([self.positional_embedding[:1, :], new_pos[0]], 0)
+        #####################################################################################
+        
+        
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        #x = self.ln_post(x[:, 0, :])
+        x = self.ln_post(x) # return both cls token and image tokens
+
+        if self.proj is not None:
+            x = x @ self.proj
+
+        return x
+
+
+class CLIP(nn.Module):
+    def __init__(self,
+                 embed_dim: int,
+                 # vision
+                 image_resolution: int,
+                 vision_layers: Union[Tuple[int, int, int, int], int],
+                 vision_width: int,
+                 vision_patch_size: int,
+                 # text
+                 context_length: int,
+                 vocab_size: int,
+                 transformer_width: int,
+                 transformer_heads: int,
+                 transformer_layers: int
+                 ):
+        super().__init__()
+
+        self.context_length = context_length
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(
+                layers=vision_layers,
+                output_dim=embed_dim,
+                heads=vision_heads,
+                input_resolution=image_resolution,
+                width=vision_width
+            )
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisionTransformer(
+                input_resolution=image_resolution,
+                patch_size=vision_patch_size,
+                width=vision_width,
+                layers=vision_layers,
+                heads=vision_heads,
+                output_dim=embed_dim
+            )
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask()
+        )
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features ** -0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image):
+        return self.visual(image.type(self.dtype))
+
+    def encode_text(self, text):
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+        return x
+    
+    def encode_text_learn(self, prompts, tokenized_prompts, deep_compound_prompts_text = None, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        # x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+        
+        # x = x + self.positional_embedding.to(cast_dtype)
+
+        x = prompts + self.positional_embedding.to(cast_dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        # print("test", x.shape, len(deep_compound_prompts_text))
+        if deep_compound_prompts_text is None:
+            x = self.transformer(x)
+        else:
+            x = self.transformer([x, deep_compound_prompts_text, 0])
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)  # [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), tokenized_prompts.argmax(dim=-1)] @ self.text_projection
+        return x
+
+
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logits_per_image.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        return logits_per_image, logits_per_text

AnomalyCLIP_lib/__init__.py

@@ -0,0 +1 @@
+from .model_load import *

AnomalyCLIP_lib/bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917

AnomalyCLIP_lib/build_model.py

@@ -0,0 +1,50 @@
+from torch import nn
+from .CLIP import CLIP
+from .AnomalyCLIP import AnomalyCLIP
+
+def build_model(name: str, state_dict: dict, design_details = None):
+    vit = "visual.proj" in state_dict
+    
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_resolution = vision_patch_size * grid_size
+    else:
+        counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_resolution = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+    # print('name', name)
+    # if 'CS-' in name:
+    if design_details is not None:
+        model = AnomalyCLIP(
+            embed_dim,
+            image_resolution, vision_layers, vision_width, vision_patch_size,
+            context_length, vocab_size, transformer_width, transformer_heads, transformer_layers, design_details = design_details
+        )
+    else:
+        model = CLIP(
+            embed_dim,
+            image_resolution, vision_layers, vision_width, vision_patch_size,
+            context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
+        )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        if key in state_dict:
+            del state_dict[key]
+
+    #convert_weights(model)
+    model.load_state_dict(state_dict)
+    return model.eval()

AnomalyCLIP_lib/constants.py

@@ -0,0 +1,2 @@
+OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
+OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)

AnomalyCLIP_lib/model_load.py

@@ -0,0 +1,235 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Union, List
+from pkg_resources import packaging
+
+import torch
+from PIL import Image
+from torchvision.transforms import Compose, Resize, ToTensor, Normalize
+from tqdm import tqdm
+import numpy as np
+
+from .build_model import build_model
+from .simple_tokenizer import SimpleTokenizer as _Tokenizer
+from torchvision.transforms import InterpolationMode
+
+if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
+    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
+
+
+__all__ = ["available_models", "load", 
+           "get_similarity_map",  "compute_similarity"]
+_tokenizer = _Tokenizer()
+
+_MODELS = {
+    "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
+}
+
+
+def _download(
+        url: str,
+        cache_dir: Union[str, None] = None,
+):
+
+    if not cache_dir:
+        # cache_dir = os.path.expanduser("~/.cache/clip")
+        cache_dir = os.path.expanduser("/remote-home/iot_zhouqihang/root/.cache/clip")
+    os.makedirs(cache_dir, exist_ok=True)
+    filename = os.path.basename(url)
+
+    if 'openaipublic' in url:
+        expected_sha256 = url.split("/")[-2]
+    elif 'mlfoundations' in url:
+        expected_sha256 = os.path.splitext(filename)[0].split("-")[-1]
+    else:
+        expected_sha256 = ''
+
+    download_target = os.path.join(cache_dir, filename)
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        if expected_sha256:
+            if hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
+                return download_target
+            else:
+                warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+        else:
+            return download_target
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.headers.get("Content-Length")), ncols=80, unit='iB', unit_scale=True) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    if expected_sha256 and not hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
+        raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def _transform(n_px):
+    return Compose([
+        Resize((n_px, n_px), interpolation=InterpolationMode.BICUBIC),
+        #CenterCrop(n_px), # rm center crop to explain whole image
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+def available_models() -> List[str]:
+    """Returns the names of available CLIP models"""
+    return list(_MODELS.keys())
+
+
+def load_state_dict(checkpoint_path: str, map_location='cpu'):
+    checkpoint = torch.load(checkpoint_path, map_location=map_location)
+    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
+        state_dict = checkpoint['state_dict']
+    else:
+        state_dict = checkpoint
+    if next(iter(state_dict.items()))[0].startswith('module'):
+        state_dict = {k[7:]: v for k, v in state_dict.items()}
+    return state_dict
+
+def load_checkpoint(model, checkpoint_path, strict=True):
+    state_dict = load_state_dict(checkpoint_path)
+    # detect old format and make compatible with new format
+    if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
+        state_dict = convert_to_custom_text_state_dict(state_dict)
+    resize_pos_embed(state_dict, model)
+    incompatible_keys = model.load_state_dict(state_dict, strict=strict)
+    return incompatible_keys
+
+def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", design_details = None, jit: bool = False, download_root: str = None):
+    """Load a CLIP model
+
+    Parameters
+    ----------
+    name : str
+        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+
+    device : Union[str, torch.device]
+        The device to put the loaded model
+
+    jit : bool
+        Whether to load the optimized JIT model or more hackable non-JIT model (default).
+
+    download_root: str
+        path to download the model files; by default, it uses "~/.cache/clip"
+
+    Returns
+    -------
+    model : torch.nn.Module
+        The CLIP model
+
+    preprocess : Callable[[PIL.Image], torch.Tensor]
+        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+    """
+    print("name", name)
+    if name in _MODELS:
+        # model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
+        model_path = _download(_MODELS[name], download_root or os.path.expanduser("/remote-home/iot_zhouqihang/root/.cache/clip"))
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+    with open(model_path, 'rb') as opened_file:
+        try:
+            # loading JIT archive
+            model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
+            state_dict = None
+        except RuntimeError:
+            # loading saved state dict
+            if jit:
+                warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+                jit = False
+            state_dict = torch.load(opened_file, map_location="cpu")
+
+    if not jit:
+        model = build_model(name, state_dict or model.state_dict(), design_details).to(device)
+        if str(device) == "cpu":
+            model.float()
+        return model, _transform(model.visual.input_resolution)
+
+    # patch the device names
+    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+    def patch_device(module):
+        try:
+            graphs = [module.graph] if hasattr(module, "graph") else []
+        except RuntimeError:
+            graphs = []
+
+        if hasattr(module, "forward1"):
+            graphs.append(module.forward1.graph)
+
+        for graph in graphs:
+            for node in graph.findAllNodes("prim::Constant"):
+                if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
+                    node.copyAttributes(device_node)
+
+    model.apply(patch_device)
+    patch_device(model.encode_image)
+    patch_device(model.encode_text)
+
+    # patch dtype to float32 on CPU
+    if str(device) == "cpu":
+        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+        float_node = float_input.node()
+
+        def patch_float(module):
+            try:
+                graphs = [module.graph] if hasattr(module, "graph") else []
+            except RuntimeError:
+                graphs = []
+
+            if hasattr(module, "forward1"):
+                graphs.append(module.forward1.graph)
+
+            for graph in graphs:
+                for node in graph.findAllNodes("aten::to"):
+                    inputs = list(node.inputs())
+                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
+                        if inputs[i].node()["value"] == 5:
+                            inputs[i].node().copyAttributes(float_node)
+
+        model.apply(patch_float)
+        patch_float(model.encode_image)
+        patch_float(model.encode_text)
+
+        model.float()
+
+    return model, _transform(model.input_resolution.item())
+
+
+def get_similarity_map(sm, shape):
+    side = int(sm.shape[1] ** 0.5)
+    sm = sm.reshape(sm.shape[0], side, side, -1).permute(0, 3, 1, 2)
+    sm = torch.nn.functional.interpolate(sm, shape, mode='bilinear')
+    sm = sm.permute(0, 2, 3, 1)
+    return sm
+
+
+def compute_similarity(image_features, text_features, t=2):
+    prob_1 = image_features[:, :1, :] @ text_features.t()
+    b, n_t, n_i, c = image_features.shape[0], text_features.shape[0], image_features.shape[1], image_features.shape[2]
+    feats = image_features.reshape(b, n_i, 1, c) * text_features.reshape(1, 1, n_t, c)
+    similarity = feats.sum(-1)
+    return (similarity/0.07).softmax(-1), prob_1

AnomalyCLIP_lib/simple_tokenizer.py

@@ -0,0 +1,132 @@
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a signficant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path: str = default_bpe()):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+        merges = merges[1:49152-256-2+1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v+'</w>' for v in vocab]
+        for merge in merges:
+            vocab.append(''.join(merge))
+        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        return text

AnomalyCLIP_lib/transform.py

@@ -0,0 +1,133 @@
+import warnings
+from dataclasses import dataclass, asdict
+from typing import Any, Dict, Optional, Sequence, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as F
+
+from torchvision.transforms import Normalize, Compose, RandomResizedCrop, InterpolationMode, ToTensor, Resize, \
+    CenterCrop
+
+from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+
+
+@dataclass
+class AugmentationCfg:
+    scale: Tuple[float, float] = (0.9, 1.0)
+    ratio: Optional[Tuple[float, float]] = None
+    color_jitter: Optional[Union[float, Tuple[float, float, float]]] = None
+    interpolation: Optional[str] = None
+    re_prob: Optional[float] = None
+    re_count: Optional[int] = None
+    use_timm: bool = False
+
+
+class ResizeMaxSize(nn.Module):
+
+    def __init__(self, max_size, interpolation=InterpolationMode.BICUBIC, fn='max', fill=0):
+        super().__init__()
+        if not isinstance(max_size, int):
+            raise TypeError(f"Size should be int. Got {type(max_size)}")
+        self.max_size = max_size
+        self.interpolation = interpolation
+        self.fn = min if fn == 'min' else min
+        self.fill = fill
+
+    def forward(self, img):
+        if isinstance(img, torch.Tensor):
+            height, width = img.shape[:2]
+        else:
+            width, height = img.size
+        scale = self.max_size / float(max(height, width))
+        if scale != 1.0:
+            new_size = tuple(round(dim * scale) for dim in (height, width))
+            img = F.resize(img, new_size, self.interpolation)
+            pad_h = self.max_size - new_size[0]
+            pad_w = self.max_size - new_size[1]
+            img = F.pad(img, padding=[pad_w//2, pad_h//2, pad_w - pad_w//2, pad_h - pad_h//2], fill=self.fill)
+        return img
+
+
+def _convert_to_rgb(image):
+    return image.convert('RGB')
+
+
+def image_transform(
+        image_size: int,
+        is_train: bool,
+        mean: Optional[Tuple[float, ...]] = None,
+        std: Optional[Tuple[float, ...]] = None,
+        resize_longest_max: bool = False,
+        fill_color: int = 0,
+        aug_cfg: Optional[Union[Dict[str, Any], AugmentationCfg]] = None,
+):
+    mean = mean or OPENAI_DATASET_MEAN
+    if not isinstance(mean, (list, tuple)):
+        mean = (mean,) * 3
+
+    std = std or OPENAI_DATASET_STD
+    if not isinstance(std, (list, tuple)):
+        std = (std,) * 3
+
+    if isinstance(image_size, (list, tuple)) and image_size[0] == image_size[1]:
+        # for square size, pass size as int so that Resize() uses aspect preserving shortest edge
+        image_size = image_size[0]
+
+    if isinstance(aug_cfg, dict):
+        aug_cfg = AugmentationCfg(**aug_cfg)
+    else:
+        aug_cfg = aug_cfg or AugmentationCfg()
+    normalize = Normalize(mean=mean, std=std)
+    if is_train:
+        aug_cfg_dict = {k: v for k, v in asdict(aug_cfg).items() if v is not None}
+        use_timm = aug_cfg_dict.pop('use_timm', False)
+        if use_timm:
+            from timm.data import create_transform  # timm can still be optional
+            if isinstance(image_size, (tuple, list)):
+                assert len(image_size) >= 2
+                input_size = (3,) + image_size[-2:]
+            else:
+                input_size = (3, image_size, image_size)
+            # by default, timm aug randomly alternates bicubic & bilinear for better robustness at inference time
+            aug_cfg_dict.setdefault('interpolation', 'random')
+            aug_cfg_dict.setdefault('color_jitter', None)  # disable by default
+            train_transform = create_transform(
+                input_size=input_size,
+                is_training=True,
+                hflip=0.,
+                mean=mean,
+                std=std,
+                re_mode='pixel',
+                **aug_cfg_dict,
+            )
+        else:
+            train_transform = Compose([
+                RandomResizedCrop(
+                    image_size,
+                    scale=aug_cfg_dict.pop('scale'),
+                    interpolation=InterpolationMode.BICUBIC,
+                ),
+                _convert_to_rgb,
+                ToTensor(),
+                normalize,
+            ])
+            if aug_cfg_dict:
+                warnings.warn(f'Unused augmentation cfg items, specify `use_timm` to use ({list(aug_cfg_dict.keys())}).')
+        return train_transform
+    else:
+        if resize_longest_max:
+            transforms = [
+                ResizeMaxSize(image_size, fill=fill_color)
+            ]
+        else:
+            transforms = [
+                Resize(image_size, interpolation=InterpolationMode.BICUBIC),
+                CenterCrop(image_size),
+            ]
+        transforms.extend([
+            _convert_to_rgb,
+            ToTensor(),
+            normalize,
+        ])
+        return Compose(transforms)

README.md

@@ -1,117 +1,80 @@
 
-# CLIP 기반 제품 결함 탐지 모델 카드
-
-## 모델 세부사항
-
-### 모델 설명
-
-이 모델은 CLIP 기반의 이상 탐지 방법을 사용하여 제품 결함을 탐지합니다. 
-사전 훈련된 CLIP 모델을 fine-tuning하여 제품 이미지에서 결함을 식별하고, 생산 라인에서 품질 관리 및 결함 감지를 자동화합니다.
-
-- **Developed by:** 오석
-- **Funded by:** 4INLAB INC.
-- **Shared by:** zhou2023anomalyclip
-- **Model type:** CLIP based Anomaly Detection
-- **Language(s):** Python, PyTorch
-- **License:** Apache 2.0, MIT, GPL-3.0
-
-### 기술적 제한사항
-
-- 모델은 결함 탐지를 위한 충분하고 다양한 훈련 데이터를 필요로 합니다. 훈련 데이터셋이 부족하거나 불균형할 경우, 모델의 성능이 저하될 수 있습니다.
-- 실시간 결함 감지 성능은 하드웨어 사양에 따라 달라질 수 있으며, 높은 해상도에서 결함을 탐지하는 정확도가 떨어질 수 있습니다.
-- 결함이 미세하거나 제품 간 유사성이 매우 높은 경우, 모델이 결함을 정확하게 탐지하지 못할 수 있습니다.
-
-## 학습 세부사항
-
-### Hardware
-  - **CPU:** Intel Core i9-13900K (24 Cores, 32 Threads)
-  - **RAM:** 64GB DDR5
-  - **GPU:** NVIDIA RTX 4090Ti 24GB
-  - **Storage:** 1TB NVMe SSD + 2TB HDD
-  - **Operating System:** Windows 11 pro
-
-### 데이터셋 정보
-
-이 모델은 시계열 재고 데이터를 사용하여 훈련됩니다. 이 데이터는 재고 수준, 날짜 및 기타 관련 특성에 대한 정보를 포함하고 있습니다.
-데이터는 Conv1D와 BiLSTM 레이어에 적합하도록 MinMax 스케일링을 사용하여 전처리되고 정규화됩니다.
-
-
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/65e7d0935ea025ead9623dde/E8pMyLfUnlIQFCLbTiLba.png)
-
-- **Data sources:**    https://huggingface.co/datasets/quandao92/vision-inventory-prediction-data
-- **Training size:**
-   - 1차 : Few-shot learning with anomaly (10ea), good (4ea)
-   - 2차 : Few-shot learning with anomaly (10ea), good (10ea)
-   - 3차 : Few-shot learning with anomaly (10ea), good (110ea)
-     
-- **Time-step:** 5초 이내
-
-- **Data Processing Techniques:**
-  - normalization:
-      description: "이미지 픽셀 값을 평균 및 표준편차로 표준화"
-      method: "'Normalize' from 'torchvision.transforms'"
-  - max_resize:
-      description: "이미지의 최대 크기를 유지하며, 비율을 맞추고 패딩을 추가하여 크기 조정"
-      method: "Custom 'ResizeMaxSize' class"
-  - random_resized_crop:
-      description: "훈련 중에 이미지를 랜덤으로 자르고 크기를 조정하여 변형을 추가"
-      method: "'RandomResizedCrop' from 'torchvision.transforms'"
-  - resize:
-      description: "모델 입력에 맞게 이미지를 고정된 크기로 조정"
-      method: "'Resize' with BICUBIC interpolation"
-  - center_crop:
-      description: "이미지의 중앙 부분을 지정된 크기로 자르기"
-      method: "'CenterCrop'"
-  - to_tensor:
-      description: "이미지를 PyTorch 텐서로 변환"
-      method: "'ToTensor'"
-  - augmentation (optional):
-      description: "데이터 증강을 위해 다양한 랜덤 변환 적용, 'AugmentationCfg'로 설정 가능"
-      method: "Uses 'timm' library if specified"
-
-
-# AD-CLIP Model Architecture
-
-
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/65e7d0935ea025ead9623dde/1wFBzBCgF4sOefROGE7RO.png)
-
-- **model:**
-  - input_layer:
-      - image_size: [640, 640, 3]  # 표준 입력 이미지 크기
-  - backbone:
-      - name: CLIP (ViT-B-32)  # CLIP 모델의 비전 트랜스포머를 백본으로 사용
-      - filters: [32, 64, 128, 256, 512]  # 비전 트랜스포머의 각 레이어 필터 크기
-  - neck:
-      - name: Anomaly Detection Module  # 결함 탐지를 위한 추가 모듈
-      - method: Contrastive Learning  # CLIP 모델의 특징을 사용한 대조 학습 기법
-  - head:
-      - name: Anomaly Detection Head  # 결함 탐지를 위한 최종 출력 레이어
-      - outputs:
-          - anomaly_score: 1  # 이상 탐지 점수 (비정상/정상 구분)
-          - class_probabilities: N  # 각 클래스에 대한 확률 (결함 여부)
-
-# Optimizer and Loss Function
-- **training:**
-  - optimizer:
-      - name: AdamW  # AdamW 옵티마이저 (가중치 감쇠 포함)
-      - lr: 0.0001  # 학습률
-  - loss:
-      - classification_loss: 1.0  # 분류 손실 (교차 엔트로피)
-      - anomaly_loss: 1.0  # 결함 탐지 손실 (이상 탐지 모델에 대한 손실)
-      - contrastive_loss: 1.0  # 대조 학습 손실 (유사도 기반 손실)
-
-# Metrics
-- **metrics:**
-  - Precision  # 정밀도 (Precision)
-  - Recall  # 재현율 (Recall)
-  - mAP  # 평균 정밀도 (Mean Average Precision)
-  - F1-Score  # F1-점수 (균형 잡힌 평가 지표)
-
-# Training Parameters
- **하이퍼파라미터 설정**
-- Learning Rate: 0.001.
-- Batch Size: 8.
-- Epochs: 200.
+# CLIP based ANOMALY DETECTION
+
+<div align="center">
+
+[![Status](https://img.shields.io/badge/status-active-success.svg)]()
+[![GitHub Issues](https://img.shields.io/github/issues/kylelobo/The-Documentation-Compendium.svg)](https://github.com/kylelobo/The-Documentation-Compendium/issues)
+[![GitHub Pull Requests](https://img.shields.io/github/issues-pr/kylelobo/The-Documentation-Compendium.svg)](https://github.com/kylelobo/The-Documentation-Compendium/pulls)
+[![License](https://img.shields.io/badge/license-MIT-blue.svg)](/LICENSE)
+
+</div>
+
+---
+
+<p align="center"> Anomaly detection (AD) requires detection models trained using auxiliary data to detect anomalies without any training sample in a target dataset. AnomalyCLIP is to learn object-agnostic text prompts that capture generic normality and abnormality in an image regardless of its foreground objects. This allows our model to focus on the abnormal image regions rather than the object semantics, enabling generalized normality and abnormality recognition on diverse types of objects. All experiments are conducted in PyTorch-2.0.0 with a single NVIDIA RTX 4090 24GB.
+    <br> 
+</p>
+
+
+
+
+# 📝 Table of Contents
+
+- [Update](#update)
+- [Install & Dependence](#install--dependence)
+- [Dataset Preparation](#dataset-preparation)
+- [Pre-trained CLIP model](#pre-trained-clip-model)
+- [Usage](#usage)
+- [Code Details](#code-details)
+- [References](#references)
+
+# Update
+- 08.08.2024: Code has been released !!!
+
+
+# Install & Dependence
+
+### ⭕ Tested Platform
+- Software Information
+  ```
+  OS: Windows 11 64 bit
+  Python: 3.9.18 (anaconda)
+  PyTorch: 2.0.0
+  Cuda Toolkit: 11.8
+  CudDNN: 9.3.0.75 for cuda11
+  ```
+  ![analysis](./docs/CUDA_info.png) 
+
+- Hardware
+  ```
+  CPU: Intel(R) Core(TM) i7-14700K   3.40 GHz
+  RAM: 64GB
+  GPU: Nvidia RTX4090 (24GB)
+  ```
+
+
+- Install Python libraries
+  ```
+  pip install -r requirements.txt
+  ```
+
+# Dataset Preparation
+
+Download the dataset below:
+
+* Industrial Domain:
+
+| Dataset | Download |
+| ---     | ---   |
+| MVTec | [download](https://www.mvtec.com/company/research/datasets/mvtec-ad) |
+| VisA | [download](https://github.com/amazon-science/spot-diff) |
+| MPDD | [download](https://github.com/stepanje/MPDD) |
+| BTAD | [download](http://avires.dimi.uniud.it/papers/btad/btad.zip) |
+| SDD | [download](https://www.vicos.si/resources/kolektorsdd/) |
+| DAGM | [download](https://www.kaggle.com/datasets/mhskjelvareid/dagm-2007-competition-dataset-optical-inspection) |
+| DTD-Synthetic | [download](https://drive.google.com/drive/folders/10OyPzvI3H6llCZBxKxFlKWt1Pw1tkMK1) |
+
 
 # Pre-trained CLIP model
 | Model | Download |
@@ -121,75 +84,132 @@
 | ViT-L/14 | [download](https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt) |
 | ViT-L/14@336px | [download](https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt) |
 
-# Evaluation Parameters
-- F1-score: 95%이상.
+# Usage
+- for train (Fine-tuning)
+  ```ruby
+  python train.py
+  ```
+- for test with dataset (many data)
+  ```ruby
+  python test.py
+  ```
+- for simple test (개별 이미지 테스트)
+  ```ruby
+  python Simple_test_code.py
+  ```
+- for UI app test (simple app developed)
+  ```ruby
+  python monitor_check.py
+  ```
+- for real-time detection test (webcam and video tracking)
+  ```ruby
+  python real_time_CLIP.py
+  ```
+
+# Code Details 
+
+### ✅Dataset configuration
+
+- Dataset configuration as example below
+```
+├── data/
+│   ├── COMP_1/
+│   │   ├── product_1/
+│   │   │   ├──grouth_truth
+│   │   │   │   ├──anomaly_1
+│   │   │   │   ├──anomaly_2
+│   │   │   │
+│   │   │   ├──test/
+│   │   │   │   ├──good
+│   │   │   │   ├──anomaly_1
+│   │   │   │   ├──anomaly_2
+│   │   │   │
+│   │   │   ├──train/
+│   │   │   │   ├──good
+│   │   │   │   ├──anomaly_1
+│   │   │   │   ├──anomaly_2
+│   │   │   │ 
+│   │   ├── product_2/
+│   │   │   │
+│   │
+│   ├── COMP_2/
+│   │ 
+```
 
+- Generate JSON file storing all the above information of dataset ( -> meta_train.json, meta_test.json)
+```ruby
+cd dataset_config
+python dataset_get_json.py
+```
 
-# 학습 성능 및 테스트 결과
+- Making all grouth_truth (only anomaly mask) by hand
+```ruby
+cd dataset_config
+python image_ground_truth.py
+```
 
-- **학습성능 결과과 그래프**:
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/65e7d0935ea025ead9623dde/RduhNlkWiyPXj-vbAkJga.png)
+- Dataset configuration for train and test
+```ruby
+cd training_libs
+python dataset.py
+```
 
-<div style="display: flex; justify-content: space-between;">
-  <div style="text-align: center; margin-right: 20px;">
-    <img src="https://cdn-uploads.huggingface.co/production/uploads/65e7d0935ea025ead9623dde/_lUD77x-yueXycuIn7jya.png" height="80%" width="100%" style="margin-right:5px;">
-    <p>1차 학습 성능</p>
-  </div>
-  <div style="text-align: center; margin-right: 20px;">
-    <img src="https://cdn-uploads.huggingface.co/production/uploads/65e7d0935ea025ead9623dde/NHDH9N94cI-KqP8k-ASUN.png" height="80%" width="100%" style="margin-right:5px;">
-    <p>2차 학습 성능</p>
-  </div>
-  <div style="text-align: center; margin-right: 20px;">
-    <img src="https://cdn-uploads.huggingface.co/production/uploads/65e7d0935ea025ead9623dde/6n0DnnQjXD8Ql-p3Owxan.png" height="80%" width="100%" style="margin-right:5px;">
-    <p>3차 학습 성능</p>
-  </div>
-</div>
+ →  _ _init_ _ 메서드는 데이터셋의 루트 디렉토리, 변환 함수, 데이터셋 이름, 모드를 입력으로 받음  
+ →  메타 정보를 담은 JSON 파일 (meta_train.json)을 읽어와 클래스 이름 목록과 모든 데이터 항목을 리스트에 저장  
+ →  generate_class_info 함수를 호출하여 클래스 정보를 생성하고 클래스 이름을 클래스 ID에 매핑  
+ →  _ _len_ _ 메서드는 데이터셋의 샘플 수를 반환   
+ →  _ _getitem_ _ 메서드는 주어진 인덱스의 샘플 데이터를 반환  
+ →  이미지 경로를 통해 이미지를 읽고, 이상 여부에 따라 마스크 이미지를 생성  
+ →  필요시 이미지와 마스크에 변환 함수를 적용  
+ →  이미지, 마스크, 클래스 이름, 이상 여부, 이미지 경로, 클래스 ID를 포함한 딕셔너리를 반환 
 
-- **학습 결과표**:
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/65e7d0935ea025ead9623dde/kDxl9q6X2dxCRJm5nc7jR.png)
-
-- **테스트 결과**:
-<div style="display: flex; justify-content: space-between;">
-  <div style="text-align: center; margin-right: 20px;">
-    <img src="https://cdn-uploads.huggingface.co/production/uploads/65e7d0935ea025ead9623dde/A91V0GdrcUcX01cC-biG9.png" height="600" width="1000" style="margin-right:5px;">
-    <p>Anomaly Product</p>
-  </div>
-  <div style="text-align: center; margin-right: 20px;">
-    <img src="https://cdn-uploads.huggingface.co/production/uploads/65e7d0935ea025ead9623dde/PxleIhphzViTGCubVhWn7.png" height="600" width="1000" style="margin-right:5px;">
-    <p>Normal Product</p>
-  </div>
-</div>
 
+### ✅ Image pre-processing (transformation) for train and test
+```ruby
+  training_lib/utils.py
+```
+```ruby
+  AnomalyCLIP_lib/transform.py
+```
+⭐ **Data Processing Techniques**
+1. Normalization  
+ → Standardize image pixel values using mean and standard deviation  
+ → Utilized via *'Normalize'* from *'torchvision.transforms'*  
+  
+2. Normalization  
+ → Resize the image to a maximum dimension while maintaining aspect ratio, with padding  
+ → Custom *'ResizeMaxSize'* class  
+
+3. RandomResizedCrop  
+ → Randomly crop and resize images during training to create variability  
+ → Implemented via *'RandomResizedCrop'* from *'torchvision.transforms'*  
+  
+4. Resize  
+ → Resize images to a fixed size for model input  
+ → Done using *'Resize'* with BICUBIC interpolation  
+  
+5. Center Crop  
+ → Crop the central region of the image to the desired size  
+ → Applied using *'CenterCrop'*  
+
+6. ToTensor  
+ → Convert images to PyTorch tensors  
+ → Done with *'ToTensor'*
 
-# 설치 및 실행 가이라인
+7. Augmentation (Optional)  
+ → Apply various random transformations for data augmentation, configurable via *'AugmentationCfg' *   
+ → Uses *'timm'* library if specified  
 
-이 모델을 실행하려면 Python과 함께 다음 라이브러리가 필요합니다:
 
-- **ftfy==6.2.0**: 텍스트 정규화 및 인코딩 문제를 해결하는 라이브러리.
-- **matplotlib==3.9.0**: 데이터 시각화 및 그래프 생성을 위한 라이브러리.
-- **numpy==1.24.3**: 수치 연산을 위한 핵심 라이브러리.
-- **opencv_python==4.9.0.80**: 이미지 및 비디오 처리용 라이브러리.
-- **pandas==2.2.2**: 데이터 분석 및 조작을 위한 라이브러리.
-- **Pillow==10.3.0**: 이미지 파일 처리 및 변환을 위한 라이브러리.
-- **PyQt5==5.15.10**: GUI 애플리케이션 개발을 위한 프레임워크.
-- **PyQt5_sip==12.13.0**: PyQt5와 Python 간의 인터페이스를 제공하는 라이브러리.
-- **regex==2024.5.15**: 정규 표현식 처리를 위한 라이브러리.
-- **scikit_learn==1.2.2**: 기계 학습 및 데이터 분석을 위한 라이브러리.
-- **scipy==1.9.1**: 과학 및 기술 계산을 위한 라이브러리.
-- **setuptools==59.5.0**: Python 패키지 배포 및 설치를 위한 라이브러리.
-- **scikit-image**: 이미지 처리 및 분석을 위한 라이브러리.
-- **tabulate==0.9.0**: 표 형태로 데이터를 출력하는 라이브러리.
-- **thop==0.1.1.post2209072238**: PyTorch 모델의 FLOP 수를 계산하는 도구.
-- **timm==0.6.13**: 다양한 최신 이미지 분류 모델을 제공하는 라이브러리.
-- **torch==2.0.0**: PyTorch 딥러닝 프레임워크.
-- **torchvision==0.15.1**: 컴퓨터 비전 작업을 위한 PyTorch 확장 라이브러리.
-- **tqdm==4.65.0**: 진행 상황을 시각적으로 표시하는 라이브러리.
-- **pyautogui**: GUI 자동화를 위한 라이브러리.
 
+⭐ **Libraries Used**
+1. *'torch'*: Core deep learning library for tensor operations and model building
+2. *'torchvision'*: Provides image processing utilities like Resize, CenterCrop, Normalize, etc
+3. *'timm'*: Optional, for advanced augmentation and transformations
+4. *'AnomalyCLIP_lib'*: Custom library for dataset-specific constants and transformations
 
 
 
-### 모델 실행 단계:
 
 ### ✅ Prompt generating
 ```ruby
@@ -249,18 +269,8 @@ parser.add_argument("--dpam", type=int, default=20, help="dpam size")
 →  If you want to focus only on the final layers (where the model usually learns complex features), you can choose fewer DPAM layers.
 ```
 
-### ✅ Test process
 
-👍 **Load pre-trained and Fine tuned (Checkpoints) models**
-1. Pre-trained mode (./pre-trained model/): 
- ```ruby
-→ Contains the pre-trained model (ViT-B, ViT-L,....)
-→ Used as the starting point for training the CLIP model
-→ Pre-trained model helps speed up and improve training by leveraging previously learned features
-```
-2. Fine-tuned models (./checkpoint/): 
- ```ruby
-→ "epoch_N.pth" files in this folder store the model's states during the fine-tuning process.
-→ Each ".pth" file represents a version of the model fine-tuned from the pre-trained model
-→ These checkpoints can be used to resume fine-tuning, evaluate the model at different stages, or select the best-performing version
-```
+
+# References
+- AnomalyCLIP: Object-agnostic Prompt Learning for Zero-shot Anomaly Detection [[github](https://github.com/zqhang/AnomalyCLIP.git)]
+  

dataset_config/dataset_get_json.py

@@ -0,0 +1,51 @@
+import os
+import json
+
+
+class DATASolver(object):
+    
+    CLSNAMES = [
+        'shinpyung',
+        # 'gear'# Change to 'welding_test' for testing
+    ]
+
+    def __init__(self, root='data/4inlab'):
+        self.root = root
+        self.meta_path = f'{root}/meta_train.json'  # Change to meta_test.json for testing
+
+    def run(self):
+        info = dict(train={}, test={})
+        anomaly_samples = 0
+        normal_samples = 0
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
+                    img_names.sort()
+                    mask_names.sort() if mask_names is not None else None
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+                        if phase == 'test':
+                            if is_abnormal:
+                                anomaly_samples = anomaly_samples + 1
+                            else:
+                                normal_samples = normal_samples + 1
+                info[phase][cls_name] = cls_info
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+        print('normal_samples', normal_samples, 'anomaly_samples', anomaly_samples)
+if __name__ == '__main__':
+    runner = DATASolver(root='data/4inlab')
+    runner.run()

dataset_config/image_ground_truth.py

@@ -0,0 +1,68 @@
+#%%
+import cv2
+import numpy as np
+import os
+
+# Initialize global variables
+points = []
+drawing = False
+
+# Function to clear the drawn points
+def clear_points():
+    global points
+    points = []
+
+# Mouse callback function
+def draw_polygon(event, x, y, flags, param):
+    global points, drawing, img
+
+    if event == cv2.EVENT_LBUTTONDOWN:
+        drawing = True
+        points.append((x, y))
+
+    elif event == cv2.EVENT_MOUSEMOVE:
+        if drawing:
+            img_copy = img.copy()
+            for i in range(1, len(points)):
+                cv2.line(img_copy, points[i - 1], points[i], (255, 0, 0), 2)
+            if len(points) > 0:
+                cv2.circle(img_copy, points[-1], 3, (0, 0, 255), -1, lineType=cv2.LINE_AA)  # Hiển thị điểm chọn của chuột
+            cv2.imshow('image', img_copy)
+
+    elif event == cv2.EVENT_LBUTTONUP:
+        drawing = False
+        points.append((x, y))
+        pts = np.array(points, np.int32)
+        pts = pts.reshape((-1, 1, 2))
+        mask = np.zeros(img.shape[:2], dtype=np.uint8)
+        cv2.fillPoly(mask, [pts], 255)
+        cv2.imwrite(mask_path, mask)
+        cv2.imshow('image', img)
+
+# Function to process images in a folder
+def process_images_in_folder(folder_path):
+    global img, mask_path
+
+    for img_name in os.listdir(folder_path):
+        if img_name.endswith('.jpg'):
+            img_path = os.path.join(folder_path, img_name)
+            mask_path = os.path.join(folder_path, f'{os.path.splitext(img_name)[0]}_mask.jpg')
+            img = cv2.imread(img_path)
+
+            # Create a window and bind the mouse callback function
+            cv2.namedWindow('image')
+            cv2.setMouseCallback('image', draw_polygon)
+
+            while True:
+                cv2.imshow('image', img)
+                k = cv2.waitKey(1) & 0xFF
+                if k == 27:  # Press 'ESC' to exit
+                    break
+
+            clear_points()
+            cv2.destroyAllWindows()
+
+# Define folders to process
+folder_path=r'C:\Users\20240805\Documents\GitHub\AD-CLIP\data\4inlab\shinpyung\train\anomaly'
+process_images_in_folder(folder_path)
+# %%

dataset_config/image_resize.py

@@ -0,0 +1,13 @@
+from PIL import Image
+import os
+
+def resize_image(input_image_path, output_image_path, size=(518, 518)):
+    with Image.open(input_image_path) as image:
+        resized_image = image.resize(size)
+        resized_image.save(output_image_path)
+
+# Example usage:
+input_image = r'\4inlab\shinpyung\train\anomaly'
+output_image = r'\4inlab\shinpyung\train\anomaly\resize'
+resize_image(input_image, output_image)
+

requirements.txt

@@ -0,0 +1,20 @@
+ftfy==6.2.0
+matplotlib==3.9.0
+numpy==1.24.3
+opencv_python==4.9.0.80
+pandas==2.2.2
+Pillow==10.3.0
+PyQt5==5.15.10
+PyQt5_sip==12.13.0
+regex==2024.5.15
+scikit_learn==1.2.2
+scipy==1.9.1
+setuptools==59.5.0
+scikit-image
+tabulate==0.9.0
+thop==0.1.1.post2209072238
+timm==0.6.13
+torch==2.0.0
+torchvision==0.15.1
+tqdm==4.65.0
+pyautogui

test.py

@@ -0,0 +1,231 @@
+#%%
+import AnomalyCLIP_lib
+import torch
+import argparse
+import torch.nn.functional as F
+from training_libs.prompt_ensemble import AnomalyCLIP_PromptLearner
+from training_libs.loss import FocalLoss, BinaryDiceLoss
+from training_libs.utils import normalize
+from training_libs.dataset import Dataset_test
+from training_libs.logger import get_logger
+from tqdm import tqdm
+
+import os
+import random
+import numpy as np
+from tabulate import tabulate
+from training_libs.utils import get_transform
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+from training_libs.visualization import visualizer
+
+from training_libs.metrics import image_level_metrics, pixel_level_metrics
+from tqdm import tqdm
+from scipy.ndimage import gaussian_filter
+
+
+def test(args):
+    img_size = args.image_size
+    features_list = args.features_list
+    dataset_dir = args.data_path
+    save_path = args.save_path
+    dataset_name = args.dataset
+
+    logger = get_logger(args.save_path)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # device = "gpu"
+
+    AnomalyCLIP_parameters = {"Prompt_length": args.n_ctx, "learnabel_text_embedding_depth": args.depth, "learnabel_text_embedding_length": args.t_n_ctx}
+    model, _ = AnomalyCLIP_lib.load("pre-trained models/clip/ViT-B-32.pt", device=device, design_details = AnomalyCLIP_parameters)
+    model.eval()
+    # torch.save(model.state_dict(),"pre-trained models/clip")
+
+    preprocess, target_transform = get_transform(args)
+    test_data = Dataset_test(root=args.data_path, transform=preprocess, target_transform=target_transform, dataset_name = args.dataset)
+    test_dataloader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=False)
+    obj_list = test_data.obj_list
+
+
+    results = {}
+    metrics = {}
+    for obj in obj_list:
+        results[obj] = {}
+        results[obj]['gt_sp'] = []
+        results[obj]['pr_sp'] = []
+        results[obj]['imgs_masks'] = []
+        results[obj]['anomaly_maps'] = []
+        metrics[obj] = {}
+        metrics[obj]['pixel-auroc'] = 0
+        metrics[obj]['pixel-aupro'] = 0
+        metrics[obj]['image-auroc'] = 0
+        metrics[obj]['image-ap'] = 0
+
+    prompt_learner = AnomalyCLIP_PromptLearner(model.to(device=device), AnomalyCLIP_parameters)
+    
+    
+    #Add check-point from trained model with normal images
+    # checkpoint = torch.load("checkpoint/241120_SP_DPAM_13_518/epoch_500.pth",map_location=torch.device('cpu'))
+    # prompt_learner.load_state_dict(checkpoint["prompt_learner"])
+
+
+    #Add check-point from trained model with normal images
+    # checkpoint = torch.load(args.checkpoint_path,map_location=torch.device(device=device))
+    # prompt_learner.load_state_dict(checkpoint["prompt_learner"])
+
+
+    prompt_learner.to(device)
+    model.to(device)
+    model.visual.DAPM_replace(DPAM_layer = 13)
+
+    prompts, tokenized_prompts, compound_prompts_text = prompt_learner(cls_id = None)
+    print("print(prompts)")
+    print(prompts)
+
+
+
+    text_features = model.encode_text_learn(prompts, tokenized_prompts, compound_prompts_text).float()
+    text_features = torch.stack(torch.chunk(text_features, dim = 0, chunks = 2), dim = 1)
+    text_features = text_features/text_features.norm(dim=-1, keepdim=True)
+  
+    
+
+    model.to(device)
+    for idx, items in enumerate(tqdm(test_dataloader)):
+        image = items['img'].to(device)
+        cls_name = items['cls_name']
+        cls_id = items['cls_id']
+        
+        gt_mask_initial = items['img_mask']
+        #convert gt mask to good (0) and anomaly (1)
+        gt_mask = items['img_mask']
+        gt_mask[gt_mask > 0.5], gt_mask[gt_mask <= 0.5] = 1, 0
+
+
+        results[cls_name[0]]['imgs_masks'].append(gt_mask)  # px
+        results[cls_name[0]]['gt_sp'].extend(items['anomaly'].detach().cpu())
+
+        with torch.no_grad():
+            image_features, patch_features = model.encode_image(image, features_list, DPAM_layer = 20)
+            image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+
+            text_probs = image_features @ text_features.permute(0, 2, 1)
+            text_probs = (text_probs/0.07).softmax(-1)
+            text_probs = text_probs[:, 0, 1]
+            anomaly_map_list = []
+            for idx, patch_feature in enumerate(patch_features):
+                if idx >= args.feature_map_layer[0]:
+                    patch_feature = patch_feature/ patch_feature.norm(dim = -1, keepdim = True)
+                    similarity, _ = AnomalyCLIP_lib.compute_similarity(patch_feature, text_features[0])
+                    similarity_map = AnomalyCLIP_lib.get_similarity_map(similarity[:, 1:, :], args.image_size)
+                    anomaly_map = (similarity_map[...,1] + 1 - similarity_map[...,0])/2.0
+                    anomaly_map_list.append(anomaly_map)
+
+            anomaly_map = torch.stack(anomaly_map_list)
+            
+            anomaly_map = anomaly_map.sum(dim = 0)
+            results[cls_name[0]]['pr_sp'].extend(text_probs.detach().cpu())
+            anomaly_map = torch.stack([torch.from_numpy(gaussian_filter(i, sigma = args.sigma)) for i in anomaly_map.detach().cpu()], dim = 0 )
+            results[cls_name[0]]['anomaly_maps'].append(anomaly_map)
+
+            #Save the anomaly map images
+            visualizer(items['img_path'], anomaly_map.detach().cpu().numpy(), args.image_size, args.save_path, cls_name)
+    
+    print("print(results)")
+    torch.save(results,"results/results_shinpyung_0.pt")
+    # print(results)
+
+    table_ls = []
+    image_auroc_list = []
+    image_ap_list = []
+    pixel_auroc_list = []
+    pixel_aupro_list = []
+    for obj in obj_list:
+        table = []
+        table.append(obj)
+        results[obj]['imgs_masks'] = torch.cat(results[obj]['imgs_masks'])
+        results[obj]['anomaly_maps'] = torch.cat(results[obj]['anomaly_maps']).detach().cpu().numpy()
+        if args.metrics == 'image-level':
+            image_auroc = image_level_metrics(results, obj, "image-auroc")
+            image_ap = image_level_metrics(results, obj, "image-ap")
+            table.append(str(np.round(image_auroc * 100, decimals=1)))
+            table.append(str(np.round(image_ap * 100, decimals=1)))
+            image_auroc_list.append(image_auroc)
+            image_ap_list.append(image_ap) 
+        elif args.metrics == 'pixel-level':
+            pixel_auroc = pixel_level_metrics(results, obj, "pixel-auroc")
+            pixel_aupro = pixel_level_metrics(results, obj, "pixel-aupro")
+            table.append(str(np.round(pixel_auroc * 100, decimals=1)))
+            table.append(str(np.round(pixel_aupro * 100, decimals=1)))
+            pixel_auroc_list.append(pixel_auroc)
+            pixel_aupro_list.append(pixel_aupro)
+        elif args.metrics == 'image-pixel-level':
+            image_auroc = image_level_metrics(results, obj, "image-auroc")
+            image_ap = image_level_metrics(results, obj, "image-ap")
+            pixel_auroc = pixel_level_metrics(results, obj, "pixel-auroc")
+            pixel_aupro = pixel_level_metrics(results, obj, "pixel-aupro")
+            table.append(str(np.round(pixel_auroc * 100, decimals=1)))
+            table.append(str(np.round(pixel_aupro * 100, decimals=1)))
+            table.append(str(np.round(image_auroc * 100, decimals=1)))
+            table.append(str(np.round(image_ap * 100, decimals=1)))
+            image_auroc_list.append(image_auroc)
+            image_ap_list.append(image_ap) 
+            pixel_auroc_list.append(pixel_auroc)
+            pixel_aupro_list.append(pixel_aupro)
+        table_ls.append(table)
+
+    if args.metrics == 'image-level':
+        # logger
+        table_ls.append(['mean', 
+                        str(np.round(np.mean(image_auroc_list) * 100, decimals=1)),
+                        str(np.round(np.mean(image_ap_list) * 100, decimals=1))])
+        results = tabulate(table_ls, headers=['objects', 'image_auroc', 'image_ap'], tablefmt="pipe")
+    elif args.metrics == 'pixel-level':
+        # logger
+        table_ls.append(['mean', str(np.round(np.mean(pixel_auroc_list) * 100, decimals=1)),
+                        str(np.round(np.mean(pixel_aupro_list) * 100, decimals=1))
+                       ])
+        results = tabulate(table_ls, headers=['objects', 'pixel_auroc', 'pixel_aupro'], tablefmt="pipe")
+    elif args.metrics == 'image-pixel-level':
+        # logger
+        table_ls.append(['mean', str(np.round(np.mean(pixel_auroc_list) * 100, decimals=1)),
+                        str(np.round(np.mean(pixel_aupro_list) * 100, decimals=1)), 
+                        str(np.round(np.mean(image_auroc_list) * 100, decimals=1)),
+                        str(np.round(np.mean(image_ap_list) * 100, decimals=1))])
+        results = tabulate(table_ls, headers=['objects', 'pixel_auroc', 'pixel_aupro', 'image_auroc', 'image_ap'], tablefmt="pipe")
+    logger.info("\n%s", results)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser("AnomalyCLIP", add_help=True)
+    # paths
+    parser.add_argument("--data_path", type=str, default="./data/4inlab/", help="path to test dataset")
+    parser.add_argument("--save_path", type=str, default='./results/', help='path to save results')
+    parser.add_argument("--checkpoint_path", type=str, default='./checkpoint/241122_SP_DPAM_13_518', help='path to checkpoint')
+    # model
+    parser.add_argument("--dataset", type=str, default='4inlab')
+    parser.add_argument("--image_size", type=int, default=518, help="image size")
+    parser.add_argument("--depth", type=int, default=9, help="image size")
+    parser.add_argument("--n_ctx", type=int, default=12, help="zero shot")
+    parser.add_argument("--t_n_ctx", type=int, default=4, help="zero shot")
+    parser.add_argument("--metrics", type=str, default='image-pixel-level')
+    parser.add_argument("--seed", type=int, default=111, help="random seed")
+    parser.add_argument("--sigma", type=int, default=4, help="zero shot")
+        # Specify layers from which feature maps will be extracted (can pass multiple values)
+    parser.add_argument("--feature_map_layer", type=int, nargs="+", default=[0, 1, 2, 3], help="zero shot")
+    
+    # List of layers whose features will be used
+    parser.add_argument("--features_list", type=int, nargs="+", default=[6, 12, 18, 24], help="features used")
+
+    
+    args = parser.parse_args()
+    print(args)
+    setup_seed(args.seed)
+    test(args)
+#%%

train.py

@@ -0,0 +1,207 @@
+import os
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
+import AnomalyCLIP_lib
+import torch
+import argparse
+import torch.nn.functional as F
+from training_libs.prompt_ensemble import AnomalyCLIP_PromptLearner
+from training_libs.loss import FocalLoss, BinaryDiceLoss
+from training_libs.utils import normalize
+from training_libs.dataset import Dataset_train
+from training_libs.logger import get_logger
+from tqdm import tqdm
+import numpy as np
+import random
+from training_libs.utils import get_transform
+import matplotlib.pyplot as plt
+
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+class RealTimePlotter:  #
+    def __init__(self):
+        self.epochs = []
+        self.loss_list = []
+        self.image_loss_list = []
+        self.fig, (self.ax1, self.ax2) = plt.subplots(1, 2, figsize=(14, 6))
+        plt.ion()
+        self.fig.show()
+        self.fig.canvas.flush_events()
+
+    def update(self, epoch, loss, image_loss):
+        self.epochs.append(epoch)
+        self.loss_list.append(loss)
+        self.image_loss_list.append(image_loss)
+        
+        self.ax1.clear()
+        self.ax2.clear()
+
+        self.ax1.plot(self.epochs, self.loss_list, label='Training Loss')
+        self.ax1.set_title('Training Loss')
+        self.ax1.set_xlabel('Epochs')
+        self.ax1.set_ylabel('Loss')
+        self.ax1.legend()
+
+        self.ax2.plot(self.epochs, self.image_loss_list, label='Image Loss')
+        self.ax2.set_title('Image Loss')
+        self.ax2.set_xlabel('Epochs')
+        self.ax2.set_ylabel('Loss')
+        self.ax2.legend()
+
+        self.fig.canvas.flush_events()
+
+def train(args):
+
+    logger = get_logger(args.save_path)
+
+    preprocess, target_transform = get_transform(args)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # device = "cpu"
+
+    AnomalyCLIP_parameters = {"Prompt_length": args.n_ctx, "learnabel_text_embedding_depth": args.depth, "learnabel_text_embedding_length": args.t_n_ctx}
+
+    # model, _ = AnomalyCLIP_lib.load("ViT-L/14@336px", device=device, design_details = AnomalyCLIP_parameters)
+    model, _ = AnomalyCLIP_lib.load("pre-trained models/clip/ViT-B-32.pt", device=device, design_details = AnomalyCLIP_parameters)
+    model.eval()
+
+    train_data = Dataset_train(root=args.train_data_path, transform=preprocess, target_transform=target_transform, dataset_name = args.dataset)
+    train_dataloader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True)
+
+  ##########################################################################################
+    prompt_learner = AnomalyCLIP_PromptLearner(model.to(device), AnomalyCLIP_parameters)
+    prompt_learner.to(device)
+    model.to(device)
+    model.visual.DAPM_replace(DPAM_layer = args.dpam)
+    ##########################################################################################
+    optimizer = torch.optim.Adam(list(prompt_learner.parameters()), lr=args.learning_rate, betas=(0.5, 0.999))
+
+    # losses
+    loss_focal = FocalLoss()
+    loss_dice = BinaryDiceLoss()
+    
+    
+    model.eval()
+    prompt_learner.train()
+    # plotter = RealTimePlotter()
+
+    for epoch in tqdm(range(args.epoch)):
+        model.eval()
+        prompt_learner.train()
+        loss_list = []
+        image_loss_list = []
+
+        for items in tqdm(train_dataloader):
+            image = items['img'].to(device)
+            label =  items['anomaly']
+
+            gt = items['img_mask'].squeeze().to(device)
+            gt[gt > 0.5] = 1
+            gt[gt <= 0.5] = 0
+
+            with torch.no_grad():
+                # Apply DPAM to the layer from 6 to 24
+                # DPAM_layer represents the number of layer refined by DPAM from top to bottom
+                # DPAM_layer = 1, no DPAM is used
+                # DPAM_layer = 20 as default
+                image_features, patch_features = model.encode_image(image, args.features_list, DPAM_layer = args.dpam)
+                image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+                    
+           ####################################
+            prompts, tokenized_prompts, compound_prompts_text = prompt_learner(cls_id = None)
+            text_features = model.encode_text_learn(prompts, tokenized_prompts, compound_prompts_text).float()
+            text_features = torch.stack(torch.chunk(text_features, dim = 0, chunks = 2), dim = 1)
+            text_features = text_features/text_features.norm(dim=-1, keepdim=True)
+            # Apply DPAM surgery
+            text_probs = image_features.unsqueeze(1) @ text_features.permute(0, 2, 1)
+            text_probs = text_probs[:, 0, ...]/0.07
+
+            
+            
+            image_loss = F.cross_entropy(text_probs.squeeze(), label.long().cuda()) #Process with GPU
+            #image_loss = F.cross_entropy(text_probs.squeeze(), label.long())  #Without GPU and using CPU
+            image_loss_list.append(image_loss.item())
+            ######################################################################
+            similarity_map_list = []  
+            # similarity_map_list.append(similarity_map)
+            for idx, patch_feature in enumerate(patch_features):
+                if idx >= args.feature_map_layer[0]:
+                    patch_feature = patch_feature/ patch_feature.norm(dim = -1, keepdim = True)
+                    similarity, _ = AnomalyCLIP_lib.compute_similarity(patch_feature, text_features[0])
+                    similarity_map = AnomalyCLIP_lib.get_similarity_map(similarity[:, 1:, :], args.image_size).permute(0, 3, 1, 2)
+                    similarity_map_list.append(similarity_map)
+
+            loss = 0
+            for i in range(len(similarity_map_list)):
+                loss += loss_focal(similarity_map_list[i], gt)
+                loss += loss_dice(similarity_map_list[i][:, 1, :, :], gt)
+                loss += loss_dice(similarity_map_list[i][:, 0, :, :], 1-gt)
+
+            optimizer.zero_grad()
+            (loss+image_loss).backward()
+            optimizer.step()
+            loss_list.append(loss.item())
+        # logs
+        if (epoch + 1) % args.print_freq == 0:
+            avg_loss = np.mean(loss_list)
+            avg_image_loss = np.mean(image_loss_list)
+            logger.info('epoch [{}/{}], loss:{:.4f}, image_loss:{:.4f}'.format(epoch + 1, args.epoch, avg_loss, avg_image_loss))
+            # plotter.update(epoch + 1, avg_loss, avg_image_loss)  #Realtime training performance monitoring
+            
+        # save model
+        if (epoch + 1) % args.save_freq == 0:
+            ckp_path = os.path.join(args.save_path, 'epoch_' + str(epoch + 1) + '.pth')
+            torch.save({"prompt_learner": prompt_learner.state_dict(),"epoch":epoch+1}, ckp_path)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser("AnomalyCLIP", add_help=True)  # Initialize the argument parser
+    
+    # Define the path to the training dataset and model checkpoint saving
+    parser.add_argument("--train_data_path", type=str, default="./data/4inlab", help="train dataset path")
+    parser.add_argument("--save_path", type=str, default='./checkpoint/241122_SP_DPAM_13_518', help='path to save results')
+    
+    # Specify the name of the training dataset
+    parser.add_argument("--dataset", type=str, default='4inlab', help="train dataset name")
+    
+    # Set the depth parameter (Note: "image size" in help may be misleading)
+    parser.add_argument("--depth", type=int, default=9, help="image size")
+    
+    # Set the prompt length and learnable text embedding length for "zero-shot" learning
+    parser.add_argument("--n_ctx", type=int, default=12, help="zero shot")
+    parser.add_argument("--t_n_ctx", type=int, default=4, help="zero shot")
+    
+    # Specify layers from which feature maps will be extracted (can pass multiple values)
+    parser.add_argument("--feature_map_layer", type=int, nargs="+", default=[0, 1, 2, 3], help="zero shot")
+    
+    # List of layers whose features will be used
+    parser.add_argument("--features_list", type=int, nargs="+", default=[6, 12, 18, 24], help="features used")
+    
+    # Setting parameters for training
+    parser.add_argument("--epoch", type=int, default=400, help="epochs")
+    parser.add_argument("--learning_rate", type=float, default=0.0001, help="learning rate")
+    parser.add_argument("--batch_size", type=int, default=8, help="batch size")
+    
+    # Size/depth parameter for the DPAM (Deep Prompt Attention Mechanism)
+    parser.add_argument("--dpam", type=int, default=13, help="dpam size")
+    
+    # Define the size of input images used for training
+    parser.add_argument("--image_size", type=int, default=518, help="image size")
+    
+    # Frequency (in epochs) of logging training information and saving
+    parser.add_argument("--print_freq", type=int, default=1, help="print frequency")
+    parser.add_argument("--save_freq", type=int, default=1, help="save frequency")
+    parser.add_argument("--seed", type=int, default=111, help="random seed")
+    
+    args = parser.parse_args()  # Parse the command-line arguments and store them in the 'args' object
+    setup_seed(args.seed)  # Set the random seed for reproducibility using the provided seed value
+    train(args)  # Call the training function with the parsed arguments
+
+

training_libs/dataset.py

@@ -0,0 +1,116 @@
+import torch.utils.data as data
+import json
+import random
+from PIL import Image
+import numpy as np
+import torch
+import os
+
+def generate_class_info(dataset_name, mode='train'):
+    class_name_map_class_id = {}
+    if dataset_name == 'mvtec':
+        # obj_list = ['carpet', 'bottle', 'hazelnut', 'leather', 'cable', 'capsule', 'grid', 'pill',
+        #             'transistor', 'metal_nut', 'screw', 'toothbrush', 'zipper', 'tile', 'wood']
+        obj_list = ['bottle']
+    elif dataset_name == '4inlab':
+        if mode=='train':
+            obj_list = ['shinpyung']   # With training
+        elif mode=='test':
+            obj_list = ['shinpyung']    # With testing
+    elif dataset_name == 'task1':
+        if mode=='train':
+            obj_list = ['cup']
+    elif dataset_name == 'task2':
+        if mode=='train':
+            obj_list = ['fire']
+    elif dataset_name == 'smoke_cloud':
+        if mode=='train':
+            obj_list = ['fire']
+           
+    for k, index in zip(obj_list, range(len(obj_list))):
+        class_name_map_class_id[k] = index
+
+    return obj_list, class_name_map_class_id
+
+class Dataset_test(data.Dataset):
+    def __init__(self, root, transform, target_transform, dataset_name, mode="test"):
+        self.root = root
+        self.transform = transform
+        self.target_transform = target_transform
+        self.data_all = []
+        meta_info = json.load(open(f'{self.root}/meta_train.json', 'r'))
+        name = self.root.split('/')[-1]
+        meta_info = meta_info[mode]
+
+        self.cls_names = list(meta_info.keys())
+        for cls_name in self.cls_names:
+            self.data_all.extend(meta_info[cls_name])
+        self.length = len(self.data_all)
+
+        self.obj_list, self.class_name_map_class_id = generate_class_info(dataset_name,mode='test')
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, index):
+        data = self.data_all[index]
+        img_path, mask_path, cls_name, specie_name, anomaly = data['img_path'], data['mask_path'], data['cls_name'], \
+                                                              data['specie_name'], data['anomaly']
+        img = Image.open(os.path.join(self.root, img_path))
+        if anomaly == 0:
+            img_mask = Image.fromarray(np.zeros((img.size[0], img.size[1])), mode='L')
+        else:
+            if os.path.isdir(os.path.join(self.root, mask_path)):
+                # just for classification not report error
+                img_mask = Image.fromarray(np.zeros((img.size[0], img.size[1])), mode='L')
+            else:
+                img_mask = np.array(Image.open(os.path.join(self.root, mask_path)).convert('L')) > 0
+                img_mask = Image.fromarray(img_mask.astype(np.uint8) * 255, mode='L')
+        # transforms
+        img = self.transform(img) if self.transform is not None else img
+        img_mask = self.target_transform(   
+            img_mask) if self.target_transform is not None and img_mask is not None else img_mask
+        img_mask = [] if img_mask is None else img_mask
+        return {'img': img, 'img_mask': img_mask, 'cls_name': cls_name, 'anomaly': anomaly,
+                'img_path': os.path.join(self.root, img_path), "cls_id": self.class_name_map_class_id[cls_name]}    
+    
+
+class Dataset_train(data.Dataset):
+    def __init__(self, root, transform, target_transform, dataset_name, mode="train"):
+        self.root = root
+        self.transform = transform
+        self.target_transform = target_transform
+        self.data_all = []
+        meta_info = json.load(open(f'{self.root}/meta_train.json', 'r'))
+        name = self.root.split('/')[-1]
+        meta_info = meta_info[mode]
+
+        self.cls_names = list(meta_info.keys())
+        for cls_name in self.cls_names:
+            self.data_all.extend(meta_info[cls_name])
+        self.length = len(self.data_all)
+
+        self.obj_list, self.class_name_map_class_id = generate_class_info(dataset_name,mode='train')
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, index):
+        data = self.data_all[index]
+        img_path, mask_path, cls_name, specie_name, anomaly = data['img_path'], data['mask_path'], data['cls_name'], \
+                                                              data['specie_name'], data['anomaly']
+        img = Image.open(os.path.join(self.root, img_path))
+        if anomaly == 0:
+            img_mask = Image.fromarray(np.zeros((img.size[0], img.size[1])), mode='L')
+        else:
+            if os.path.isdir(os.path.join(self.root, mask_path)):
+                # just for classification not report error
+                img_mask = Image.fromarray(np.zeros((img.size[0], img.size[1])), mode='L')
+            else:
+                img_mask = np.array(Image.open(os.path.join(self.root, mask_path)).convert('L')) > 0
+                img_mask = Image.fromarray(img_mask.astype(np.uint8) * 255, mode='L')
+        # transforms
+        img = self.transform(img) if self.transform is not None else img
+        img_mask = self.target_transform(   
+            img_mask) if self.target_transform is not None and img_mask is not None else img_mask
+        img_mask = [] if img_mask is None else img_mask
+        return {'img': img, 'img_mask': img_mask, 'cls_name': cls_name, 'anomaly': anomaly,
+                'img_path': os.path.join(self.root, img_path), "cls_id": self.class_name_map_class_id[cls_name]}    

training_libs/logger.py

@@ -0,0 +1,25 @@
+
+import logging
+import os
+
+def get_logger(save_path):
+    if not os.path.exists(save_path):
+        os.makedirs(save_path)
+
+    txt_path = os.path.join(save_path, 'log.txt')
+    # logger
+    root_logger = logging.getLogger()
+    for handler in root_logger.handlers[:]:
+        root_logger.removeHandler(handler)
+    root_logger.setLevel(logging.WARNING)
+    logger = logging.getLogger('test')
+    formatter = logging.Formatter('%(asctime)s.%(msecs)03d - %(levelname)s: %(message)s',
+                                    datefmt='%y-%m-%d %H:%M:%S')
+    logger.setLevel(logging.INFO)
+    file_handler = logging.FileHandler(txt_path, mode='a')
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(formatter)
+    logger.addHandler(console_handler)
+    return logger

training_libs/loss.py

@@ -0,0 +1,125 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from math import exp
+
+class FocalLoss(nn.Module):
+    """
+    copy from: https://github.com/Hsuxu/Loss_ToolBox-PyTorch/blob/master/FocalLoss/FocalLoss.py
+    This is a implementation of Focal Loss with smooth label cross entropy supported which is proposed in
+    'Focal Loss for Dense Object Detection. (https://arxiv.org/abs/1708.02002)'
+        Focal_Loss= -1*alpha*(1-pt)*log(pt)
+    :param alpha: (tensor) 3D or 4D the scalar factor for this criterion
+    :param gamma: (float,double) gamma > 0 reduces the relative loss for well-classified examples (p>0.5) putting more
+                    focus on hard misclassified example
+    :param smooth: (float,double) smooth value when cross entropy
+    :param balance_index: (int) balance class index, should be specific when alpha is float
+    :param size_average: (bool, optional) By default, the losses are averaged over each loss element in the batch.
+    """
+
+    def __init__(self, apply_nonlin=None, alpha=None, gamma=2, balance_index=0, smooth=1e-5, size_average=True):
+        super(FocalLoss, self).__init__()
+        self.apply_nonlin = apply_nonlin
+        self.alpha = alpha
+        self.gamma = gamma
+        self.balance_index = balance_index
+        self.smooth = smooth
+        self.size_average = size_average
+
+        if self.smooth is not None:
+            if self.smooth < 0 or self.smooth > 1.0:
+                raise ValueError('smooth value should be in [0,1]')
+
+    def forward(self, logit, target):
+        if self.apply_nonlin is not None:
+            logit = self.apply_nonlin(logit)
+        num_class = logit.shape[1]
+
+        if logit.dim() > 2:
+            # N,C,d1,d2 -> N,C,m (m=d1*d2*...)
+            logit = logit.view(logit.size(0), logit.size(1), -1)
+            logit = logit.permute(0, 2, 1).contiguous()
+            logit = logit.view(-1, logit.size(-1))
+        target = torch.squeeze(target, 1)
+        target = target.view(-1, 1)
+        alpha = self.alpha
+
+        if alpha is None:
+            alpha = torch.ones(num_class, 1)
+        elif isinstance(alpha, (list, np.ndarray)):
+            assert len(alpha) == num_class
+            alpha = torch.FloatTensor(alpha).view(num_class, 1)
+            alpha = alpha / alpha.sum()
+        elif isinstance(alpha, float):
+            alpha = torch.ones(num_class, 1)
+            alpha = alpha * (1 - self.alpha)
+            alpha[self.balance_index] = self.alpha
+
+        else:
+            raise TypeError('Not support alpha type')
+
+        if alpha.device != logit.device:
+            alpha = alpha.to(logit.device)
+
+        idx = target.cpu().long()
+
+        one_hot_key = torch.FloatTensor(target.size(0), num_class).zero_()
+        one_hot_key = one_hot_key.scatter_(1, idx, 1)
+        if one_hot_key.device != logit.device:
+            one_hot_key = one_hot_key.to(logit.device)
+
+        if self.smooth:
+            one_hot_key = torch.clamp(
+                one_hot_key, self.smooth / (num_class - 1), 1.0 - self.smooth)
+        pt = (one_hot_key * logit).sum(1) + self.smooth
+        logpt = pt.log()
+
+        gamma = self.gamma
+
+        alpha = alpha[idx]
+        alpha = torch.squeeze(alpha)
+        loss = -1 * alpha * torch.pow((1 - pt), gamma) * logpt
+
+        if self.size_average:
+            loss = loss.mean()
+        return loss
+
+
+class BinaryDiceLoss(nn.Module):
+    def __init__(self):
+        super(BinaryDiceLoss, self).__init__()
+
+    def forward(self, input, targets):
+        # Get the size N of each batch
+        N = targets.size()[0]
+        # Smooth variable
+        smooth = 1
+        # Reshape the width and height to the same dimension
+        input_flat = input.view(N, -1)
+        targets_flat = targets.view(N, -1)
+
+        intersection = input_flat * targets_flat
+        N_dice_eff = (2 * intersection.sum(1) + smooth) / (input_flat.sum(1) + targets_flat.sum(1) + smooth)
+        # Calculate the loss average for each image in a batch
+        loss = 1 - N_dice_eff.sum() / N
+        return loss
+    
+def smooth(arr, lamda1):
+    new_array = arr
+    arr2 = torch.zeros_like(arr)
+    arr2[:, :-1, :] = arr[:, 1:, :]
+    arr2[:, -1, :] = arr[:, -1, :]
+
+    new_array2 = torch.zeros_like(new_array)
+    new_array2[:, :, :-1] = new_array[:, :, 1:]
+    new_array2[:, :, -1] = new_array[:, :, -1]
+    loss = (torch.sum((arr2 - arr) ** 2) + torch.sum((new_array2 - new_array) ** 2)) / 2
+    return lamda1 * loss
+
+def sparsity(arr, target, lamda2):
+    if target == 0:
+        loss = torch.mean(torch.norm(arr, dim=0))
+    else:
+        loss = torch.mean(torch.norm(1-arr, dim=0))
+    return lamda2 * loss

training_libs/metrics.py

@@ -0,0 +1,60 @@
+from sklearn.metrics import auc, roc_auc_score, average_precision_score, f1_score, precision_recall_curve, pairwise
+import numpy as np
+from skimage import measure
+
+def cal_pro_score(masks, amaps, max_step=200, expect_fpr=0.3):
+    # ref: https://github.com/gudovskiy/cflow-ad/blob/master/train.py
+    binary_amaps = np.zeros_like(amaps, dtype=bool)
+    min_th, max_th = amaps.min(), amaps.max()
+    delta = (max_th - min_th) / max_step
+    pros, fprs, ths = [], [], []
+    for th in np.arange(min_th, max_th, delta):
+        binary_amaps[amaps <= th], binary_amaps[amaps > th] = 0, 1
+        pro = []
+        for binary_amap, mask in zip(binary_amaps, masks):
+            for region in measure.regionprops(measure.label(mask)):
+                tp_pixels = binary_amap[region.coords[:, 0], region.coords[:, 1]].sum()
+                pro.append(tp_pixels / region.area)
+        inverse_masks = 1 - masks
+        fp_pixels = np.logical_and(inverse_masks, binary_amaps).sum()
+        fpr = fp_pixels / inverse_masks.sum()
+        pros.append(np.array(pro).mean())
+        fprs.append(fpr)
+        ths.append(th)
+    pros, fprs, ths = np.array(pros), np.array(fprs), np.array(ths)
+    idxes = fprs < expect_fpr
+    fprs = fprs[idxes]
+    fprs = (fprs - fprs.min()) / (fprs.max() - fprs.min())
+    pro_auc = auc(fprs, pros[idxes])
+    return pro_auc
+
+
+def image_level_metrics(results, obj, metric):
+    gt = results[obj]['gt_sp']
+    pr = results[obj]['pr_sp']
+    gt = np.array(gt)
+    pr = np.array(pr)
+    if metric == 'image-auroc':
+        performance = roc_auc_score(gt, pr)
+    elif metric == 'image-ap':
+        performance = average_precision_score(gt, pr)
+
+    return performance
+    # table.append(str(np.round(performance * 100, decimals=1)))
+
+
+def pixel_level_metrics(results, obj, metric):
+    gt = results[obj]['imgs_masks']
+    pr = results[obj]['anomaly_maps']
+    gt = np.array(gt)
+    pr = np.array(pr)
+    if metric == 'pixel-auroc':
+        performance = roc_auc_score(gt.ravel(), pr.ravel())
+    elif metric == 'pixel-aupro':
+        if len(gt.shape) == 4:
+            gt = gt.squeeze(1)
+        if len(pr.shape) == 4:
+            pr = pr.squeeze(1)
+        performance = cal_pro_score(gt, pr)
+    return performance
+    

training_libs/prompt_ensemble.py

@@ -0,0 +1,273 @@
+import os
+from typing import Union, List
+from pkg_resources import packaging
+import torch
+import numpy as np
+from AnomalyCLIP_lib.simple_tokenizer import SimpleTokenizer as _Tokenizer
+# from open_clip import tokenizer
+# simple_tokenizer = tokenizer.SimpleTokenizer()
+from copy import deepcopy
+import torch.nn as nn
+
+_tokenizer = _Tokenizer()
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]:
+    """
+    Returns the tokenized representation of given input string(s)
+
+    Parameters
+    ----------
+    texts : Union[str, List[str]]
+        An input string or a list of input strings to tokenize
+
+    context_length : int
+        The context length to use; all CLIP models use 77 as the context length
+
+    truncate: bool
+        Whether to truncate the text in case its encoding is longer than the context length
+
+    Returns
+    -------
+    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
+    We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    sot_token = _tokenizer.encoder["<|startoftext|>"]
+    eot_token = _tokenizer.encoder["<|endoftext|>"]
+    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+    if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+    else:
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
+
+    for i, tokens in enumerate(all_tokens):
+        if len(tokens) > context_length:
+            if truncate:
+                tokens = tokens[:context_length]
+                tokens[-1] = eot_token
+            else:
+                raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result
+
+# def encode_text_with_prompt_ensemble(model, texts, device):
+#     prompt_normal = ['{}', 'flawless {}', 'perfect {}', 'unblemished {}', '{} without flaw', '{} without defect', '{} without damage']
+#     prompt_abnormal = ['damaged {}', 'broken {}', '{} with flaw', '{} with defect', '{} with damage']
+#     prompt_state = [prompt_normal, prompt_abnormal]
+#     prompt_templates = ['a bad photo of a {}.', 'a low resolution photo of the {}.', 'a bad photo of the {}.', 'a cropped photo of the {}.', 'a bright photo of a {}.', 'a dark photo of the {}.', 'a photo of my {}.', 'a photo of the cool {}.', 'a close-up photo of a {}.', 'a black and white photo of the {}.', 'a bright photo of the {}.', 'a cropped photo of a {}.', 'a jpeg corrupted photo of a {}.', 'a blurry photo of the {}.', 'a photo of the {}.', 'a good photo of the {}.', 'a photo of one {}.', 'a close-up photo of the {}.', 'a photo of a {}.', 'a low resolution photo of a {}.', 'a photo of a large {}.', 'a blurry photo of a {}.', 'a jpeg corrupted photo of the {}.', 'a good photo of a {}.', 'a photo of the small {}.', 'a photo of the large {}.', 'a black and white photo of a {}.', 'a dark photo of a {}.', 'a photo of a cool {}.', 'a photo of a small {}.', 'there is a {} in the scene.', 'there is the {} in the scene.', 'this is a {} in the scene.', 'this is the {} in the scene.', 'this is one {} in the scene.']
+
+#     text_features = []
+#     for i in range(len(prompt_state)):
+#         prompted_state = [state.format(texts[0]) for state in prompt_state[i]]
+#         prompted_sentence = []
+#         for s in prompted_state:
+#             for template in prompt_templates:
+#                 prompted_sentence.append(template.format(s))
+#         prompted_sentence = tokenize(prompted_sentence)
+#         class_embeddings = model.encode_text(prompted_sentence.to(device))
+#         class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)
+#         class_embedding = class_embeddings.mean(dim=0)
+#         class_embedding /= class_embedding.norm()
+#         text_features.append(class_embedding)
+
+#     text_features = torch.stack(text_features, dim=1).to(device).t()
+
+#     return text_features
+
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([deepcopy(module) for i in range(N)])
+class AnomalyCLIP_PromptLearner(nn.Module):
+    def __init__(self, clip_model, design_details):
+        super().__init__()
+        classnames = ["object"]
+        self.n_cls = len(classnames)
+        self.n_ctx = design_details["Prompt_length"]
+        n_ctx_pos = self.n_ctx
+        n_ctx_neg = self.n_ctx
+        self.text_encoder_n_ctx = design_details["learnabel_text_embedding_length"] 
+        ctx_init_pos = ""
+        ctx_init_neg = ""
+        dtype = clip_model.transformer.get_cast_dtype()
+        device = clip_model.token_embedding.weight.device
+
+        ctx_dim = clip_model.ln_final.weight.shape[0]
+
+        
+        self.classnames = classnames
+
+        self.state_normal_list = [
+            "{}",
+        ]
+
+        self.state_anomaly_list = [
+            "damaged {}",
+        ]
+        
+        normal_num = len(self.state_normal_list)
+        anormaly_num = len(self.state_anomaly_list)
+        self.normal_num = normal_num
+        self.anormaly_num = anormaly_num
+
+        if ctx_init_pos and ctx_init_neg:
+            # use given words to initialize context vectors
+            ctx_init_pos = ctx_init_pos.replace("_", " ")
+            ctx_init_neg = ctx_init_neg.replace("_", " ")
+            n_ctx_pos = len(ctx_init_pos.split(" "))
+            n_ctx_neg = len(ctx_init_neg.split(" "))
+            # Initialize text into bpd encoding
+            prompt_pos = tokenize(ctx_init_pos)
+            prompt_neg = tokenize(ctx_init_neg)
+            with torch.no_grad():
+                # Generate corresponding text embedding
+                embedding_pos = clip_model.token_embedding(prompt_pos).type(dtype)
+                embedding_neg = clip_model.token_embedding(prompt_neg).type(dtype)
+            # Remove EOS and # CLS, EOS, and get the learnable textual prompt
+            ctx_vectors_pos = embedding_pos[0, 1: 1 + n_ctx_pos, :]
+            ctx_vectors_neg = embedding_neg[0, 1: 1 + n_ctx_neg, :]
+            prompt_prefix_pos = ctx_init_pos
+            prompt_prefix_neg = ctx_init_neg
+            if True:
+                ctx_vectors_pos_ = []
+                ctx_vectors_neg_ = []
+                for _ in range(self.n_cls):
+                    ctx_vectors_pos_.append(deepcopy(ctx_vectors_pos))
+                    ctx_vectors_neg_.append(deepcopy(ctx_vectors_neg))
+                ctx_vectors_pos = torch.stack(ctx_vectors_pos_, dim=0)
+                ctx_vectors_neg = torch.stack(ctx_vectors_neg_, dim=0)
+
+        else:
+            # Random Initialization
+            if True:
+                print("Initializing class-specific contexts")
+                # Here cls is the number of classes, n_ctx_pos represents the length of learnable tokens, ctx_dim indicates the dimension of the prompt
+                ctx_vectors_pos = torch.empty(self.n_cls, self.normal_num, n_ctx_pos, ctx_dim, dtype=dtype)
+                ctx_vectors_neg = torch.empty(self.n_cls, self.anormaly_num, n_ctx_neg, ctx_dim, dtype=dtype)
+            else:
+                print("Initializing a generic context")
+                ctx_vectors_pos = torch.empty(n_ctx_pos, ctx_dim, dtype=dtype)
+                ctx_vectors_neg = torch.empty(n_ctx_neg, ctx_dim, dtype=dtype)
+            nn.init.normal_(ctx_vectors_pos, std=0.02)
+            nn.init.normal_(ctx_vectors_neg, std=0.02)
+            prompt_prefix_pos = " ".join(["X"] * n_ctx_pos)
+            prompt_prefix_neg = " ".join(["X"] * n_ctx_neg)
+        self.compound_prompts_depth = design_details["learnabel_text_embedding_depth"]
+        self.compound_prompts_text = nn.ParameterList([nn.Parameter(torch.empty(self.text_encoder_n_ctx, ctx_dim))
+                                                      for _ in range(self.compound_prompts_depth - 1)])
+        for single_para in self.compound_prompts_text:
+            print("single_para", single_para.shape)
+            nn.init.normal_(single_para, std=0.02)
+
+        single_layer = nn.Linear(ctx_dim, 896)
+        self.compound_prompt_projections = _get_clones(single_layer, self.compound_prompts_depth - 1)
+
+
+        self.ctx_pos = nn.Parameter(ctx_vectors_pos)  # to be optimized
+        self.ctx_neg = nn.Parameter(ctx_vectors_neg)  # to be optimized
+
+        classnames = [name.replace("_", " ") for name in classnames]
+        name_lens = [len(_tokenizer.encode(name)) for name in classnames]
+
+
+        prompts_pos = [prompt_prefix_pos +  " " + template.format(name)+ "." for template in self.state_normal_list for name in classnames]
+        prompts_neg = [prompt_prefix_neg +  " " + template.format(name)+ "." for template in self.state_anomaly_list for name in classnames]
+        
+        # print("Normal Prompt:",prompts_pos )
+        # print("Anomaly Prompt:",prompts_neg )
+
+        tokenized_prompts_pos = []
+        tokenized_prompts_neg = []
+     
+        for p_pos in prompts_pos:
+            tokenized_prompts_pos.append(tokenize(p_pos))
+        for p_neg in prompts_neg:
+            tokenized_prompts_neg.append(tokenize(p_neg))
+            
+            
+        tokenized_prompts_pos = [tokenize(p_pos).to(device) for p_pos in prompts_pos]  # Move tokenized_prompts_pos to the same device
+        tokenized_prompts_neg = [tokenize(p_neg).to(device) for p_neg in prompts_neg]  # Move tokenized_prompts_neg to the same device
+
+        tokenized_prompts_pos = torch.cat(tokenized_prompts_pos)
+        tokenized_prompts_neg = torch.cat(tokenized_prompts_neg)
+        # Generate corresponding text embedding
+        with torch.no_grad():
+            embedding_pos = clip_model.token_embedding(tokenized_prompts_pos).type(dtype)
+            embedding_neg = clip_model.token_embedding(tokenized_prompts_neg).type(dtype)
+            n, l, d = embedding_pos.shape
+            print("embedding_pos", embedding_pos.shape)
+            embedding_pos = embedding_pos.reshape(normal_num, self.n_cls, l, d).permute(1, 0, 2, 3)
+            embedding_neg = embedding_neg.reshape(anormaly_num, self.n_cls, l, d).permute(1, 0, 2, 3)
+
+
+        self.register_buffer("token_prefix_pos", embedding_pos[:, :, :1, :] )
+        self.register_buffer("token_suffix_pos", embedding_pos[:, :,1 + n_ctx_pos:, :])
+        self.register_buffer("token_prefix_neg", embedding_neg[:,:, :1, :])
+        self.register_buffer("token_suffix_neg", embedding_neg[:, :, 1 + n_ctx_neg:, :])
+
+        n, d = tokenized_prompts_pos.shape
+        tokenized_prompts_pos = tokenized_prompts_pos.reshape(normal_num, self.n_cls, d).permute(1, 0, 2)
+
+        n, d = tokenized_prompts_neg.shape
+        tokenized_prompts_neg = tokenized_prompts_neg.reshape(anormaly_num, self.n_cls, d).permute(1, 0, 2)
+
+        self.n_ctx_pos = n_ctx_pos
+        self.n_ctx_neg = n_ctx_neg
+        # tokenized_prompts = torch.cat([tokenized_prompts_pos, tokenized_prompts_neg], dim=0)  # torch.Tensor
+        self.register_buffer("tokenized_prompts_pos", tokenized_prompts_pos)
+        self.register_buffer("tokenized_prompts_neg", tokenized_prompts_neg)
+        print("tokenized_prompts shape", self.tokenized_prompts_pos.shape, self.tokenized_prompts_neg.shape)
+
+
+
+    def forward(self, cls_id =None):
+        
+        ctx_pos = self.ctx_pos
+        ctx_neg = self.ctx_neg
+        ctx_pos = self.ctx_pos
+        ctx_neg = self.ctx_neg
+        # print("shape", self.ctx_pos[0:1].shape, ctx_pos.shape)
+        prefix_pos = self.token_prefix_pos
+        prefix_neg = self.token_prefix_neg
+        suffix_pos = self.token_suffix_pos
+        suffix_neg = self.token_suffix_neg
+
+        # print(prefix_pos.shape, prefix_neg.shape)
+
+        prompts_pos = torch.cat(
+            [
+                # N(the number of template), 1, dim
+                prefix_pos,  # (n_cls, 1, dim)
+                ctx_pos,  # (n_cls, n_ctx, dim)
+                suffix_pos,  # (n_cls, *, dim)
+            ],
+            dim=2,
+        )
+
+        prompts_neg = torch.cat(
+            [
+                prefix_neg,  # (n_cls, 1, dim)
+                ctx_neg,  # (n_cls, n_ctx, dim)
+                suffix_neg,  # (n_cls, *, dim)
+            ],
+            dim=2,
+        )
+        _, _, l, d = prompts_pos.shape
+        prompts_pos = prompts_pos.reshape(-1, l, d)
+        _, _, l, d = prompts_neg.shape
+        prompts_neg = prompts_neg.reshape(-1, l, d)
+        prompts = torch.cat([prompts_pos, prompts_neg], dim=0)
+
+
+        _, l, d = self.tokenized_prompts_pos.shape
+        tokenized_prompts_pos = self.tokenized_prompts_pos.reshape(-1,  d)
+        _, l, d = self.tokenized_prompts_neg.shape
+        tokenized_prompts_neg = self.tokenized_prompts_neg.reshape(-1,  d)
+        tokenized_prompts = torch.cat((tokenized_prompts_pos, tokenized_prompts_neg), dim = 0)
+
+
+        return prompts, tokenized_prompts, self.compound_prompts_text

training_libs/utils.py

@@ -0,0 +1,24 @@
+import torchvision.transforms as transforms
+# from torchvision.transforms import Compose, Resize, ToTensor, Normalize, InterpolationMode
+from AnomalyCLIP_lib.transform import image_transform
+from AnomalyCLIP_lib.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+
+
+
+def normalize(pred, max_value=None, min_value=None):
+    if max_value is None or min_value is None:
+        return (pred - pred.min()) / (pred.max() - pred.min())
+    else:
+        return (pred - min_value) / (max_value - min_value)
+
+def get_transform(args):
+    preprocess = image_transform(args.image_size, is_train=False, mean = OPENAI_DATASET_MEAN, std = OPENAI_DATASET_STD)
+    target_transform = transforms.Compose([
+        transforms.Resize((args.image_size, args.image_size)),
+        transforms.CenterCrop(args.image_size),
+        transforms.ToTensor()
+    ])
+    preprocess.transforms[0] = transforms.Resize(size=(args.image_size, args.image_size), interpolation=transforms.InterpolationMode.BICUBIC,
+                                                    max_size=None, antialias=None)
+    preprocess.transforms[1] = transforms.CenterCrop(size=(args.image_size, args.image_size))
+    return preprocess, target_transform

training_libs/visualization.py

@@ -0,0 +1,25 @@
+import cv2
+import os
+# from utils import normalize
+from sklearn.preprocessing import normalize
+import numpy as np
+
+def visualizer(pathes, anomaly_map, img_size, save_path, cls_name):
+    for idx, path in enumerate(pathes):
+        cls = path.split('/')[-2]
+        filename = path.split('/')[-1]
+        vis = cv2.cvtColor(cv2.resize(cv2.imread(path), (img_size, img_size)), cv2.COLOR_BGR2RGB)  # RGB
+        mask = normalize(anomaly_map[idx])
+        vis = apply_ad_scoremap(vis, mask)
+        vis = cv2.cvtColor(vis, cv2.COLOR_RGB2BGR)  # BGR
+        save_vis = os.path.join(save_path, 'imgs', cls_name[idx], cls)
+        if not os.path.exists(save_vis):
+            os.makedirs(save_vis)
+        cv2.imwrite(os.path.join(save_vis, filename), vis)
+
+def apply_ad_scoremap(image, scoremap, alpha=0.5):
+    np_image = np.asarray(image, dtype=float)
+    scoremap = (scoremap * 255).astype(np.uint8)
+    scoremap = cv2.applyColorMap(scoremap, cv2.COLORMAP_JET)
+    scoremap = cv2.cvtColor(scoremap, cv2.COLOR_BGR2RGB)
+    return (alpha * np_image + (1 - alpha) * scoremap).astype(np.uint8)