text: # -*- coding: utf-8 -*-
  """Siamese Network.ipynb
  
  Automatically generated by Colaboratory.
  
  Original file is located at
      https://colab.research.google.com/drive/1EVGtztJrs-M-YnBjL8ZxJhHJ6V7hIerz
  
  SCALE INVARIANT FEATURE TRANSFORM
  """
  """
  from google.colab import drive
  drive.mount('/content/drive')
  """
  
  """UPLOAD IMAGES
  
  """
  
  import torchvision.datasets as datasets
  import numpy as np
  from PIL import Image
  import os
  from collections import defaultdict
  import torch
  import torchvision
  import torch.nn as nn
  import torch.nn.functional as F
  from torch import optim
  import time
  import os
  import torch.utils.data as Data
  from tensorboardX import SummaryWriter
  from tqdm import tqdm
  
  ##pip install TensorboardX
  
  """
  IMG_EXTENSIONS.append('tiff')
  
  genuine_dir = ''
  imposter_dir = ''
  
  genuine = datasets.ImageFolder(genuine_dir, transform = transform)
  imposter = datasets.ImageFolder(imposter_dir, transform = transform)
  
  """
  
  def upload_pil(img_folder):
    
    ref = dict()
    dep = defaultdict(list)
    for file in os.listdir(img_folder):
      img_path = os.path.join(img_folder, file)
      #print('IMAGE PATH: ', img_path)
      #print('FILE: ', file)
      img_name = file.replace('.tif','')
      image = np.array(Image.open(img_path), dtype = int)/255
      image = image.astype('int64')
      #print(image.dtype)
      name_arr = img_name.split('_')
      key = str(name_arr[0])
      if name_arr[-1] == 'r':
        ref.update({key: image})
        #print('----------')
      else:
        dep[key].append(image)      
    print('UPLOAD DONE')
    return ref, dep
  
  """PREPROCESSING OF IMAGE"""
  
  def preprocess(img_dict, state):
    if state == 0:
      for k in img_dict.keys():
        if img_dict[k].shape[1] == 252:
          img_dict[k] = img_dict[k][18:-18, 14:-14]
        else:
          img_dict[k] = img_dict[k][22:-22, 18:-18]
        
        # write custom preprocess to eliminate black
    else:
      for k in img_dict.keys():
        for i, img in enumerate(img_dict[k]):
          if img_dict[k][i].shape[1] == 252:
            img_dict[k][i] = img_dict[k][i][18:-18, 14:-14]
          else:
            img_dict[k][i] = img_dict[k][i][22:-22, 18:-18]
        # write custom preprocess to eliminate black
          
    return img_dict
  
  """TESSALATION OF IMAGES"""
  
  def sliding_window(img, length = 32, width = 32):
    new_arr = []
  
    for hor in range(1,9):
      for vert in range(1,8):
        new_arr.append(img[(hor-1)*length:hor*length, (vert-1)*width:vert*width])
  
    #print(len(new_arr))
    return new_arr
  
  def tess(ref, dep, label):
    # Original patch image shape (224, 256)
  
    img_shape = (32,32)
    img_pair_ref = []
    img_pair_dep = []
    labels = []
  
    dict_len = len(ref)
  
    for key in range(1,dict_len+1):
      ref_img = ref[str(key)]
      #print(ref_img.shape)
      ref_arr = sliding_window(ref_img)
  
      for dep_img in dep[str(key)]:
        aln_arr = sliding_window(dep_img)
        
        for sub_ref, sub_aln in zip(ref_arr, aln_arr):
          #print('SUB_REF_SHAPE: ', sub_ref.shape)
          #print(type(sub_ref))
          img_pair_ref.append(sub_ref)
          img_pair_dep.append(sub_aln)
          labels.append(label)
  
    return img_pair_ref, img_pair_dep, labels
  
  """DATALOADER"""
  
  class Dataloader(Data.Dataset):
    def __init__(self, ref_pairs, dep_pairs, labels):
      super(Dataloader).__init__()
  
      self.ref_pairs = torch.from_numpy(ref_pairs)
      self.dep_pairs = torch.from_numpy(dep_pairs)
      self.labels = torch.from_numpy(labels)
      self.len = len(labels)
  
    def __len__(self):
      return self.len
    
    def __getitem__(self, idx):
      return (self.ref_pairs[idx], self.dep_pairs[idx], self.labels[idx])
  
  """MODEL ARCHITECTURE
  
  """
  
  class Net(nn.Module):
    def __init__(self, activ = True):
      super(Net, self).__init__()
  
      ## input image size : (32, 32, 1)
      self.activ = activ
  
      self.seq1 = nn.Sequential(
          nn.Conv2d(1, 2, kernel_size = (3, 3)),
          nn.SELU(inplace = activ),
          nn.BatchNorm2d(num_features = 2),
          nn.Dropout(0.2)
      )
      self.seq2_cnn1 = nn.Sequential(
        nn.Conv2d(2, 4, kernel_size = 3, stride = 2, dilation= 1),
        nn.SELU(inplace = activ)
      )
      self.seq2_cnn2 = nn.Sequential(
        nn.Conv2d(2, 4, kernel_size = 3, stride = 2, dilation = 2),
        nn.SELU(inplace = activ)
      )
      self.seq2_cnn3 = nn.Sequential(
        nn.Conv2d(2, 4, kernel_size = 3, stride = 2, dilation = 3),
        nn.SELU(inplace = activ)
      )
      self.seq3 = nn.Sequential(
          nn.Conv2d(4, 8, kernel_size = 3),
          nn.SELU(inplace= activ),
          nn.BatchNorm2d(num_features = 8)
      )
      self.seq4_cnn1 = nn.Sequential(
          nn.Conv2d(8, 16, kernel_size = 3, stride = 2, dilation = 1),
          nn.SELU(inplace= activ)
      )
      self.seq4_cnn2 = nn.Sequential(
          nn.Conv2d(8, 16, kernel_size = 3, stride = 2, dilation = 2),
          nn.SELU(inplace= activ)
      )
      self.seq4_cnn3 = nn.Sequential(
          nn.Conv2d(8, 16, kernel_size = 3, stride = 2, dilation = 3),
          nn.SELU(inplace= activ)
      )
      self.seq5 = nn.Sequential(
          nn.Conv2d(16, 32, kernel_size = 3),
          nn.SELU(inplace = activ),
          nn.Dropout(0.2),
  
          nn.Conv2d(32, 64, kernel_size = 3, stride  = 2),
          nn.SELU(inplace = activ),
          nn.BatchNorm2d(num_features = 64),
          nn.Dropout(0.2),
  
          nn.Conv2d(64, 64, kernel_size = 1),
          nn.SELU(inplace = activ),
          nn.AdaptiveAvgPool2d((1,1))
      )
  
      self.fc = nn.Sequential(
          nn.Linear(64, 512),
          nn.SELU(inplace = activ),
          nn.Linear(512, 64),
          nn.SELU(inplace = activ),
          nn.Linear(64, 1)
      )
  
    def forward_once(self, x):
      x1 = self.seq1(x)
      #print('BEFORE ENTERING: ', x1.shape)
      x2_1 = self.seq2_cnn1(x1)
      #print('X1: ', x2_1.shape)
      x2_2 = self.seq2_cnn2(x1)
      #print('X2: ', x2_2.shape)
      x2_2 = F.pad(input=x2_2, pad=(0, 1, 0, 1), mode='constant', value=0)
      #print('X2_PADDED: ', x2_2.shape)
      x2_3 = self.seq2_cnn3(x1)
      #print('X3: ', x2_3.shape)
      x2_3 = F.pad(input=x2_3, pad=(1, 1, 1, 1), mode='constant', value=0)
      #print('X3_PADDED: ', x2_3.shape)
  
      temp = x2_1.add(x2_2)
      x2 = temp.add(x2_3)
  
      x3 = self.seq3(x2)
      #print('BEFORE ENTERING: ', x3.shape)
      x4_1 = self.seq4_cnn1(x3)
      #print('X1: ', x4_1.shape)
      x4_2 = self.seq4_cnn2(x3)
      x4_2 = F.pad(input=x4_2, pad=(0, 1, 0, 1), mode='constant', value=0)
      #print('X2: ', x4_2.shape)
      x4_3 = self.seq4_cnn3(x3)
      x4_3 = F.pad(input=x4_3, pad=(1, 1, 1, 1), mode='constant', value=0)
      #print('X3: ', x4_3.shape)
  
      temp = x4_1.add(x4_2)
      x4 = temp.add(x4_3)
  
      out = self.seq5(x4)
      out = out.view(out.size()[0], -1)
  
      out = self.fc(out)
  
      return out
  
    def forward(self, input1, input2):
      output1 = self.forward_once(input1)
      output2 = self.forward_once(input2)
  
      return output1, output2
  
  """
  from torchsummary import summary
  model = Net()
  print(model)
  summary(model, (1,32,32))
  """
  
  # loss = nn.CosineEmbeddingLoss(output1, output2, gt_label)
  
  """CONTRASTIVE LOSS FUNCTION"""
  
  class ContrastiveLoss(nn.Module):
  
      def __init__(self, margin=1.0):
          super(ContrastiveLoss, self).__init__()
          self.margin = margin
  
      def forward(self, x0, x1, y):
          # euclidian distance
          diff = x0 - x1
          dist_sq = torch.sum(torch.pow(diff, 2), 1)
          dist = torch.sqrt(dist_sq)
  
          mdist = self.margin - dist
          dist = torch.clamp(mdist, min=0.0)
          loss = y * dist_sq + (1 - y) * torch.pow(dist, 2)
          loss = torch.sum(loss) / 2.0 / x0.size()[0]
          return loss
  
  """TRAINING THE MODEL"""
  
  
  """
  To view the data written by tensorboardX
  tensorboard --logdir <path of logs directory>
  In my case, pathdir = 'logs/'
  """
  
  #os.makedirs('/drive/MyDrive', SAVE_DIR, exist_ok=True)
  
  def init_weights(model):
    for name, param in model.named_parameters():
      nn.init.uniform_(param.data, -0.08, 0.08)
  
  def train(ref_imgs, dep_imgs, labels):
  
    data = Dataloader(ref_imgs, dep_imgs, labels)
    logger = SummaryWriter(os.path.join(HOME, LOG_DIR, TIME + ': Smit'))
    opt = {
        'batch_size': 16,
        'lr': 0.01,
        'epochs': 5,
        'train_len': int(0.75*len(labels)),
        'val_len': int(0.25*len(labels)),
    }
  
    model = Net()
    # model.apply(init_weights)
  
    train_loader = Data.DataLoader(Data.Subset(data, range(opt['train_len'])), batch_size = 16,shuffle = True)
    val_loader = Data.DataLoader(Data.Subset(data, range(opt['val_len'])), batch_size = 16, shuffle = True)
  
    optimizer = optim.Adam(model.parameters(), lr = opt['lr'])
    model.train()
    model.to(DEVICE)
  
    cosine_loss = nn.CosineEmbeddingLoss()
    #contrastive_loss = ContrastiveLoss()
  
    print('-----------------BEGIN TRAINING-------------------')
  
    for epoch in range(opt['epochs']):
      train_loss = 0.0
      train_corr = 0.0
      itr = 0
      model.train()
  
      for img1, img2, label in tqdm(train_loader, ascii = True, desc = 'Train' + str(epoch)):
        img1 = img1.unsqueeze(1).type(torch.cuda.FloatTensor).to(DEVICE)
        img2 = img2.unsqueeze(1).type(torch.cuda.FloatTensor).to(DEVICE)
        label = label.to(DEVICE)
  
        # print('PRE MODEL GTG')
        # print('IMG1: ',img1.shape)
        # print('IMG2: ',img2.shape)
        out1, out2 = model(img1, img2)
        # out1, out2 = out1.view(-1), out2.view(-1)
        loss = cosine_loss(out1, out2, label)
        optimizer.zero_grad()
        loss.backward()
  
        optimizer.step()
  
        train_loss += loss.item()
        train_corr += torch.sum(out1.round() == out2.round())
        itr+=1
  
      torch.cuda.empty_cache()
  
      logger.add_scalar('Training_loss', train_loss/itr, epoch)
      print('EPOCH: ', epoch, '---', train_loss/itr)
      logger.add_scalar('Training_accuracy', train_corr.double()/itr, epoch)
      print('EPOCH: ', epoch, '---', train_corr.double()/itr)
  
      print('----------------BEGIN VALIDATION---------------------')
      itr = 0
      val_loss = 0.0
      val_corr = 0.0
  
      model.eval()
      model.to(DEVICE)
      with torch.no_grad():
        for img1, img2, label in tqdm(val_loader, ascii = True, desc = 'Validation' + str(epoch)):
          img1 = img1.unsqueeze(1).type(torch.cuda.FloatTensor).to(DEVICE)
          img2 = img2.unsqueeze(1).type(torch.cuda.FloatTensor).to(DEVICE)
          label = label.to(DEVICE)
  
          out1, out2 = model(img1, img2)
          # out1, out2 = out1.view(-1), out2.view(-1)
          loss = cosine_loss(out1, out2, label)
          val_loss += loss.item()
          val_corr += torch.sum(out1.round() == out2.round())
          itr+=1
  
        torch.cuda.empty_cache()
  
        logger.add_scalar('Validation_loss', val_loss/itr, epoch)
        print('EPOCH: ', epoch, '---', val_loss/itr)
        logger.add_scalar('Validation_accuracy', val_corr.double()/itr, epoch)
        print('EPOCH: ', epoch, '---', val_corr.double()/itr)
    
    print('TRAINING DONE')
    torch.save({'state_dict': model.state_dict()}, os.path.join(HOME, 'multi_cnn.pt'))
    logger.close()
  
  if __name__ == '__main__':
     LOG_DIR = 'logs'
     HOME = '/drive/Mydrive'
     SAVE_DIR = 'save'
     TIME = time.strftime("%Y%m%d_%H%M%S")
     DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
  
  """MAIN
  
  """
  
  gen_dir = '/content/drive/MyDrive/aligned_genuine'
  #gen_dir = '/content/drive/MyDrive/1'
  gen_ref, gen_dep = upload_pil(gen_dir)
  
  imp_dir = '/content/drive/MyDrive/aligned_imposter/aligned_imposter'
  #imp_dir = '/content/drive/MyDrive/2'
  imp_ref, imp_dep = upload_pil(imp_dir)
  
  gen_ref_p, gen_dep_p = preprocess(gen_ref, 0), preprocess(gen_dep, 1)
  imp_ref_p, imp_dep_p = preprocess(imp_ref, 0), preprocess(imp_dep, 1)
  
  gen_pair_ref, gen_pair_dep, gen_labels = tess(gen_ref_p, gen_dep_p, 1)
  imp_pair_ref, imp_pair_dep, imp_labels = tess(imp_ref_p, imp_dep_p, -1)
  
  ref_pairs = [*gen_pair_ref, *imp_pair_ref]
  dep_pairs = [*gen_pair_dep, *imp_pair_dep]
  labels = [*gen_labels, *imp_labels]
  
  ref_pairs = np.asarray(ref_pairs)
  dep_pairs = np.asarray(dep_pairs)
  labels = np.asarray(labels)
  
  print(ref_pairs.dtype)
  print(dep_pairs.dtype)
  print(labels.dtype)
  
  train(ref_pairs, dep_pairs, labels)
  
  """TEST CODE
  
  """
  """
  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
  
  model = Net()
  checkpoint = torch.load(os.path.join(HOME, 'multi_cnn.pt'))
  #running_model = torch.load('multi_cnn.pt', map_location = device)
  model.load_state_dict(checkpoint['state_dict'])
  
  """
  # upload and convert to array
  # preprocess
  #