text: import numpy as np
  import json
  import os
  import sys
  import time
  import math
  import io
  import torch
  import torch.nn as nn
  import torch.optim as optim
  from torchvision import models
  import torchvision.datasets as dsets
  import torchvision.transforms as transforms
  from torchattacks.attack import Attack
  from utils import *
  from compression import *
  from decompression import *
  from PIL import ImageFile
  from info_attack import InfoDrop
  from Models.transformers import diet_tiny, diet_small, vit_tiny, vit_small
  
  
  from PIL import ImageFile
  import lpips
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  ImageFile.LOAD_TRUNCATED_IMAGES = True
  from timm.data import resolve_data_config
  from timm.data.transforms_factory import create_transform
  
  source_model = ["resnet"]
  target_model = ["resnet18"]
  
  
  model_t = ["resnet", vit_tiny, vit_small, diet_tiny, diet_small]
  q_sizes = [20, 60, 100]
  attacks = [False]
  model_names = ["ResNet50", "ViT_tiny", "ViT_small", "DieT_tiny", "DieT_small"]
  
  """
  Source s_model: Model on which the attack would be done, adversarial images
  will be generated using this s_model.
  target s_model: Model on which the adversarial images would be tested
  
  """
  
  class Normalize(nn.Module):
      def __init__(self, mean, std):
          super(Normalize, self).__init__()
          self.register_buffer('mean', torch.Tensor(mean))
          self.register_buffer('std', torch.Tensor(std))
  
      def forward(self, input):
          # Broadcasting
          input = input / 255.0
          mean = self.mean.reshape(1, 3, 1, 1)
          std = self.std.reshape(1, 3, 1, 1)
          return (input - mean.to(device=input.device)) / std.to(
              device=input.device)
  def pred_label_and_confidence(model, input_batch, labels_to_class):
      input_batch = input_batch.cuda()
      with torch.no_grad():
          out = model(input_batch)
      _, index = torch.max(out, 1)
  
      percentage = torch.nn.functional.softmax(out, dim=1) * 100
      # print(percentage.shape)
      pred_list = []
      for i in range(index.shape[0]):
          pred_class = labels_to_class[index[i]]
          pred_conf = str(round(percentage[i][index[i]].item(), 2))
          pred_list.append([pred_class, pred_conf])
      return pred_list
  
  def lpips_2imgs(img_batch0, img_batch1, version="0.1", use_gpu=True):
      loss_fn = lpips.LPIPS(net='alex', version=version)
  
      if (use_gpu):
          loss_fn.cuda()
  
      # img0 = lpips.im2tensor(lpips.load_image(path0))  # RGB image from [-1,1]
      # img1 = lpips.im2tensor(lpips.load_image(path1))
  
      if (use_gpu):
          img_batch0 = img_batch0.cuda()
          img_batch1 = img_batch1.cuda()
      dist01 = loss_fn.forward(img_batch0, img_batch1)
      # print('Distance: %.3f' % dist01)
      return dist01
  
  
  def build_model(model, adver=False):
      """
      Builds the model and the norm layer
      adver=True means the adversarial images will be passed through the model,
      no need of resizing for that case.
      """
      if model == "resnet":
          transform = transforms.Compose([
                                          transforms.Resize((224, 224)),
                                          transforms.ToTensor(), ]
                                          ) \
              if not adver else transforms.Compose([transforms.ToTensor(),])
          norm_layer = Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
          backbone = models.resnet50(pretrained=True)
  
      elif model == "resnet18":
          transform = transforms.Compose([
                                          transforms.Resize((224, 224)),
                                          transforms.ToTensor(), ]
                                          ) \
              if not adver else transforms.Compose([transforms.ToTensor(),])
          norm_layer = Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
          backbone = models.resnet18(pretrained=True)
  
      else:
  
          backbone = model()
          config = resolve_data_config({}, model=backbone)
          transform = create_transform(**config)
          transform.transforms.pop()
          transform = transform if not adver else transforms.Compose([transforms.ToTensor(),])
          norm_layer = Normalize(mean=config['mean'],
                                 std=config['std'])
      return backbone, norm_layer, transform
  
  if __name__ == "__main__":
      f = open("results/transferability.txt", "w")
  
      for att in attacks:
          idx_ = 0
          targetted_attack = att
          for s, t in zip(source_model, target_model):
  
  
              if targetted_attack:
                   name = model_names[idx_]+"_targetted"
              else:
                  name = model_names[idx_]+ "_untargetted"
              print(f"{idx_}::: model_name: {name}")
              for q_size in q_sizes:
  
                  device = torch.device(
                      "cuda" if torch.cuda.is_available() else "cpu")
                  class_idx = json.load(open("./imagenet_class_index.json"))
                  idx2label = [class_idx[str(k)][1] for k in range(len(class_idx))]
                  class2label = [class_idx[str(k)][0] for k in
                                 range(len(class_idx))]
  
                  backbone_s, norm_layer_s, transform_s = build_model(s, adver=False)
                  backbone_t, norm_layer_t, transform_t = build_model(t, adver=True)
  
                  s_model = nn.Sequential(norm_layer_s, backbone_s.to(device))
                  s_model = s_model.eval()
                  s_model_name = name
  
                  t_model = nn.Sequential(norm_layer_t, backbone_t.to(device))
                  t_model = t_model.eval()
  
                  batch_size = 20
                  q_size = q_size
                  cur_cnt = 0
                  suc_cnt = 0
                  data_dir = "./test-data"
                  save_dir = "./results"
                  data_clean(data_dir)
                  normal_data = image_folder_custom_label(root=data_dir,
                                                          transform=transform_s,
                                                          idx2label=class2label)
                  normal_loader = torch.utils.data.DataLoader(normal_data,
                                                              batch_size=batch_size,
                                                              shuffle=False)
  
                  i = 0
                  fool_rate = 0
                  file_number = 0
                  lpips_score = 0
  
                  for i, (images, labels) in enumerate(normal_loader):  # in range(tar_cnt//batch_size):
  
                      print("Iter: ", i)
                      gt_labels = labels
                      if targetted_attack:
                          labels = torch.from_numpy(np.random.randint(0, 1000, size=images.shape[0]))
  
                      images = images * 255.0
                      steps = 500 if targetted_attack else 50   # change again from 20 to 500
                      attack = InfoDrop(s_model, batch_size=images.shape[0],
                                        q_size=q_size, steps=steps,
                                        targeted=targetted_attack)
                      # Add unclamped adv images
                      at_images, at_labels, suc_step, at_images_unclamped = attack(images, labels)
  
                      ### Calculate fool rate on the target s_model
  
                      labels_before_attack = t_model(images.to(device="cuda"))
                      _, labels_before_attack = torch.max(labels_before_attack.data, 1)
  
                      labels_after_attack = t_model(at_images_unclamped.to(device="cuda"))
                      _, labels_after_attack = torch.max(
                          labels_after_attack.data, 1)
                      if (labels_after_attack == at_labels).sum() != 20:
                          print("different results")
                      fool_rate += torch.sum(labels_before_attack != labels_after_attack)
  
  
                      # outputs_pre_attack = s_model(images.to(device="cuda"))
                      # _, pred_pre_attack_label = torch.max(outputs_pre_attack.data,
                      #                                      1)
                      # fool_rate += torch.sum(pred_pre_attack_label != at_labels)
  
                      dist = lpips_2imgs(at_images.to(device="cuda"), images.to(device="cuda"))
                      dist = dist.sum()/at_images.shape[0]
                     # print(f"Avg Sim Batch {dist}")
                      lpips_score += dist
  
                      labels = labels.to(device)
                      # Success rate of adversarial examples on the target model
                      # check target labels with the predictions of the target model
                      if targetted_attack:
                          suc_cnt += (labels_after_attack == labels).sum().item()
                      else:
                          suc_cnt += (labels_after_attack != labels).sum().item()
                      print("Current suc. rate: ", suc_cnt / ((i + 1) * batch_size))
  
                  score_list = np.zeros(len(normal_data))
                  score_list[:suc_cnt] = 1.0
                  stderr_dist = np.std(np.array(score_list)) / np.sqrt(
                      len(score_list))
                  print('Avg suc rate: %.5f +/- %.5f' % (
                  suc_cnt / len(normal_data), stderr_dist))
                  print(f"Fool Rate {q_size} is : {fool_rate / len(normal_data)}")
                  print(f"Average Similarity score: {lpips_score / len(normal_loader)}")
                  f.write(
                      f"{name}_{q_size},{(suc_cnt / len(normal_data))}, {stderr_dist}, {fool_rate / len(normal_data)}, {lpips_score / len(normal_loader)} \n")
              idx_ += 1
  
      f.close()
  
  