Upload 11 files

combined_sampler.py

@@ -0,0 +1,30 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data.sampler import Sampler
+from tqdm import *
+
+class BalancedSampler(Sampler):
+    def __init__(self, data_source, batch_size, images_per_class=3):
+        self.data_source = data_source
+        self.ys = np.array(data_source.all_labels)
+        self.num_groups = batch_size // images_per_class
+        self.batch_size = batch_size
+        self.num_instances = images_per_class
+        self.num_samples = len(self.ys)
+        self.num_classes = len(set(self.ys))
+
+    def __len__(self):
+        return self.num_samples
+
+    def __iter__(self):
+        num_batches = len(self.data_source) // self.batch_size
+        ret = []
+        while num_batches > 0:
+            sampled_classes = np.random.choice(self.num_classes, self.num_groups, replace=False)
+            for i in range(len(sampled_classes)):
+                ith_class_idxs = np.nonzero(self.ys == sampled_classes[i])[0]
+                class_sel = np.random.choice(ith_class_idxs, size=self.num_instances, replace=True)
+                ret.extend(np.random.permutation(class_sel))
+            num_batches -= 1
+        return iter(ret) 

hkpoly_evaluation_phase1.py

@@ -0,0 +1,106 @@
+# script to evaluated HKPolyU testing dataset on finetuned model after phase 1
+import torch
+from datasets.hkpoly_test import hktest
+from utils import Prev_RetMetric, l2_norm, compute_recall_at_k
+import numpy as np
+from tqdm import tqdm
+from model import SwinModel_domain_agnostic as Model
+from sklearn.metrics import roc_curve, auc
+import json
+
+def calculate_tar_at_far(fpr, tpr, target_fars):
+    tar_at_far = {}
+    for far in target_fars:
+        if far in fpr:
+            tar = tpr[np.where(fpr == far)][0]
+        else:
+            tar = np.interp(far, fpr, tpr)
+        tar_at_far[far] = tar
+    return tar_at_far
+
+if __name__ == '__main__':
+    device = torch.device('cuda')
+    data = hktest(split = 'test')
+    dataloader = torch.utils.data.DataLoader(data,batch_size = 16, num_workers = 1, pin_memory = True)
+    model = Model().to(device)
+    checkpoint = torch.load("ridgeformer_checkpoints/phase1_ft_hkpoly.pt",map_location = torch.device('cpu'))
+    model.load_state_dict(checkpoint,strict=False)
+    model.eval()
+
+    cl_feats, cb_feats, cl_labels, cb_labels, cl_feats_unnormed, cb_feats_unnormed = list(),list(),list(),list(),list(),list()
+    with torch.no_grad():
+        for (x_cl, x_cb, label) in tqdm(dataloader):
+            x_cl, x_cb, label = x_cl.to(device), x_cb.to(device), label.to(device)
+            x_cl_feat, x_cl_token = model.get_embeddings(x_cl,'contactless')
+            x_cb_feat,x_cb_token = model.get_embeddings(x_cb,'contactbased')
+            cl_feats_unnormed.append(x_cl_feat.cpu().detach().numpy())
+            cb_feats_unnormed.append(x_cb_feat.cpu().detach().numpy())
+            x_cl_feat = l2_norm(x_cl_feat).cpu().detach().numpy()
+            x_cb_feat = l2_norm(x_cb_feat).cpu().detach().numpy()
+            label = label.cpu().detach().numpy()
+            cl_feats.append(x_cl_feat)
+            cb_feats.append(x_cb_feat)
+            cl_labels.append(label)
+            cb_labels.append(label)
+
+    cl_feats = np.concatenate(cl_feats)
+    cb_feats = np.concatenate(cb_feats)
+    cl_feats_unnormed = np.concatenate(cl_feats_unnormed)
+    cb_feats_unnormed = np.concatenate(cb_feats_unnormed)
+    cl_label = torch.from_numpy(np.concatenate(cl_labels))
+    cb_label = torch.from_numpy(np.concatenate(cb_labels))
+
+    # CB2CL
+    squared_diff = np.sum(np.square(cl_feats_unnormed[:, np.newaxis] - cb_feats_unnormed), axis=2)
+    distance     = -1 * np.sqrt(squared_diff)
+    similarities = np.dot(cl_feats,np.transpose(cb_feats))
+    scores_mat = similarities + 0.1 * distance
+
+    scores = scores_mat.flatten().tolist()
+    labels = torch.eq(cl_label.view(-1,1) - cb_label.view(1,-1),0.0).flatten().tolist()
+    ids_mod = list()
+    for i in labels:
+        if i==True:
+            ids_mod.append(1)
+        else:
+            ids_mod.append(0)
+
+    fpr,tpr,thresh = roc_curve(labels,scores,drop_intermediate=True)
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.01)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.01)
+    tar_far_102 = tpr[upper_fpr_idx]
+    print(tpr[lower_fpr_idx], lower_fpr_idx, fpr[lower_fpr_idx], thresh[lower_fpr_idx])
+    print(tpr[upper_fpr_idx], upper_fpr_idx, fpr[upper_fpr_idx], thresh[upper_fpr_idx])
+
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.001)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.001)
+    tar_far_103 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    print(tpr[lower_fpr_idx], lower_fpr_idx, fpr[lower_fpr_idx])
+    print(tpr[upper_fpr_idx], upper_fpr_idx, fpr[upper_fpr_idx])
+
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.0001)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.0001)
+    tar_far_104 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    print(tpr[lower_fpr_idx], lower_fpr_idx, fpr[lower_fpr_idx])
+    print(tpr[upper_fpr_idx], upper_fpr_idx, fpr[upper_fpr_idx])
+
+    fnr = 1 - tpr
+    EER = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+    roc_auc = auc(fpr, tpr)
+    print(f"ROCAUC for CB2CL: {roc_auc * 100} %")
+    print(f"EER for CB2CL: {EER * 100} %")
+    eer_cb2cl = EER * 100
+    
+    cbcltf102 = tar_far_102 * 100
+    cbcltf103 = tar_far_103 * 100
+    cbcltf104 = tar_far_104 * 100
+    cl_label = cl_label.cpu().detach()
+    cb_label = cb_label.cpu().detach()
+    print(f"TAR@FAR=10^-2 for CB2CL: {tar_far_102 * 100} %")
+    print(f"TAR@FAR=10^-3 for CB2CL: {tar_far_103 * 100} %")
+    print(f"TAR@FAR=10^-4 for CB2CL: {tar_far_104 * 100} %")
+    
+    print(f"R@1 for CB2CL: {compute_recall_at_k(torch.from_numpy(scores_mat), cl_label, cb_label, 1) * 100} %")
+    print(f"R@10 for CB2CL: {compute_recall_at_k(torch.from_numpy(scores_mat), cl_label, cb_label, 10) * 100} %")
+    print(f"R@50 for CB2CL: {compute_recall_at_k(torch.from_numpy(scores_mat), cl_label, cb_label, 50) * 100} %")
+    print(f"R@100 for CB2CL: {compute_recall_at_k(torch.from_numpy(scores_mat), cl_label, cb_label, 100) * 100} %")

hkpoly_evaluation_phase2.py

@@ -0,0 +1,114 @@
+# script to evaluate HKPolyU testing dataset on finetuned model after phase 2
+import torch
+from datasets.hkpoly_test import hktest
+from utils import Prev_RetMetric, l2_norm, compute_recall_at_k
+import numpy as np
+from tqdm import tqdm
+from model import SwinModel_Fusion as Model
+from sklearn.metrics import roc_curve, auc
+import json
+
+def calculate_tar_at_far(fpr, tpr, target_fars):
+    tar_at_far = {}
+    for far in target_fars:
+        if far in fpr:
+            tar = tpr[np.where(fpr == far)][0]
+        else:
+            tar = np.interp(far, fpr, tpr)
+        tar_at_far[far] = tar
+    return tar_at_far
+
+def get_fused_cross_score_matrix(model, cl_tokens, cb_tokens):
+    cl_tokens   = torch.cat(cl_tokens)
+    cb_tokens   = torch.cat(cb_tokens)
+    batch_size  = cl_tokens.shape[0]
+    shard_size  = 20
+    similarity_matrix = torch.zeros((batch_size, batch_size))
+    for i_start in tqdm(range(0, batch_size, shard_size)):
+        i_end   = min(i_start + shard_size, batch_size)
+        shard_i = cl_tokens[i_start:i_end]
+        for j_start in range(0, batch_size, shard_size):
+            j_end               = min(j_start + shard_size, batch_size)
+            shard_j             = cb_tokens[j_start:j_end]
+            batch_i             = shard_i.unsqueeze(1)
+            batch_j             = shard_j.unsqueeze(0)
+            pairwise_i          = batch_i.expand(-1, shard_size, -1, -1)
+            pairwise_j          = batch_j.expand(shard_size, -1, -1, -1)
+            similarity_scores, distances   = model.combine_features(pairwise_i.reshape(-1, 197, 1024), pairwise_j.reshape(-1, 197, 1024))
+            scores = similarity_scores - 0.1 * distances
+            scores   = scores.reshape(shard_size, shard_size)
+            similarity_matrix[i_start:i_end, j_start:j_end] = scores.cpu().detach()
+    return similarity_matrix
+
+if __name__ == '__main__':
+    device = torch.device('cuda')
+    data = hktest(split = 'test')
+    dataloader = torch.utils.data.DataLoader(data,batch_size = 16, num_workers = 1, pin_memory = True)
+    model = Model().to(device)
+    checkpoint = torch.load("ridgeformer_checkpoints/phase2_ft_hkpoly.pt",map_location = torch.device('cpu'))
+    model.load_state_dict(checkpoint,strict=False)
+    model.eval()
+
+    cl_feats, cb_feats, cl_labels, cb_labels, cl_feats_unnormed, cb_feats_unnormed = list(),list(),list(),list(),list(),list()
+    with torch.no_grad():
+        for (x_cl, x_cb, label) in tqdm(dataloader):
+            x_cl, x_cb, label = x_cl.to(device), x_cb.to(device), label.to(device)
+            x_cl_token  = model.get_tokens(x_cl,'contactless')
+            x_cb_token  = model.get_tokens(x_cb,'contactbased')
+            label = label.cpu().detach().numpy()
+            cl_feats.append(x_cl_token)
+            cb_feats.append(x_cb_token)
+            cl_labels.append(label)
+            cb_labels.append(label)
+
+    cl_label = torch.from_numpy(np.concatenate(cl_labels))
+    cb_label = torch.from_numpy(np.concatenate(cb_labels))
+
+    # CB2CL
+    scores_mat = get_fused_cross_score_matrix(model, cl_feats, cb_feats)
+    scores = scores_mat.cpu().detach().numpy().flatten().tolist()
+    labels = torch.eq(cl_label.view(-1,1) - cb_label.view(1,-1),0.0).flatten().tolist()
+    ids_mod = list()
+    for i in labels:
+        if i==True:
+            ids_mod.append(1)
+        else:
+            ids_mod.append(0)
+
+    fpr,tpr,thresh = roc_curve(labels,scores,drop_intermediate=True)
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.01)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.01)
+    tar_far_102 = tpr[upper_fpr_idx]
+
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.001)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.001)
+    tar_far_103 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.0001)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.0001)
+    tar_far_104 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+
+    fnr = 1 - tpr
+    EER = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+    roc_auc = auc(fpr, tpr)
+    print(f"ROCAUC for CB2CL: {roc_auc * 100} %")
+    print(f"EER for CB2CL: {EER * 100} %")
+    eer_cb2cl = EER * 100
+    cbcltf102 = tar_far_102 * 100
+    cbcltf103 = tar_far_103 * 100
+    cbcltf104 = tar_far_104 * 100
+    cl_label = cl_label.cpu().detach()
+    cb_label = cb_label.cpu().detach()
+
+    print(f"TAR@FAR=10^-2 for CB2CL: {tar_far_102 * 100} %")
+    print(f"TAR@FAR=10^-3 for CB2CL: {tar_far_103 * 100} %")
+    print(f"TAR@FAR=10^-4 for CB2CL: {tar_far_104 * 100} %")
+
+    recall_dict = dict()
+    for i in range(1,101):
+        recall_dict[i] = compute_recall_at_k(scores_mat, cl_label, cb_label, i)
+
+    print(f"R@1 for CB2CL: {compute_recall_at_k(scores_mat, cl_label, cb_label, 1) * 100} %")
+    print(f"R@10 for CB2CL: {compute_recall_at_k(scores_mat, cl_label, cb_label, 10) * 100} %")
+    print(f"R@50 for CB2CL: {compute_recall_at_k(scores_mat, cl_label, cb_label, 50) * 100} %")
+    print(f"R@100 for CB2CL: {compute_recall_at_k(scores_mat, cl_label, cb_label, 100) * 100} %")

loss.py

@@ -0,0 +1,377 @@
+from pytorch_metric_learning import losses 
+import torch
+import torch.nn as nn
+import torch.nn.init
+import torchvision.models as models
+from torch.autograd import Variable
+from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
+from torch.nn.utils.weight_norm import weight_norm
+import torch.backends.cudnn as cudnn
+from torch.nn.utils.clip_grad import clip_grad_norm
+import numpy as np
+import os
+import torch.nn.functional as F
+import itertools
+
+torch.autograd.set_detect_anomaly(True)
+       
+class DualMSLoss_FineGrained(nn.Module):
+    """
+    Compute contrastive loss
+    """
+    def __init__(self, margin=0, max_violation=False):
+        super(DualMSLoss_FineGrained, self).__init__()
+        self.margin = margin
+        self.max_violation = max_violation
+        self.thresh = 0.5
+        self.margin = 0.7 # 0.1
+        self.scale_pos = 2
+        self.scale_neg = 40.0
+
+    def ms_sample(self,sim_mat,label):
+        pos_exp     = torch.exp(-self.scale_pos*(sim_mat-self.thresh))
+        neg_exp     = torch.exp( self.scale_neg*(sim_mat-self.thresh))
+        pos_mask    = torch.eq(label.view(-1,1)-label.view(1,-1),0.0).float().cuda()
+        neg_mask    = 1 - pos_mask
+        P_sim       = torch.where(pos_mask == 1,sim_mat,torch.ones_like(pos_exp)*1e16)
+        N_sim       = torch.where(neg_mask == 1,sim_mat,torch.ones_like(neg_exp)*-1e16)
+        min_P_sim,_ = torch.min(P_sim,dim=1,keepdim=True)
+        max_N_sim,_ = torch.max(N_sim,dim=1,keepdim=True)
+        hard_P_sim  = torch.where(P_sim - self.margin < max_N_sim,pos_exp,torch.zeros_like(pos_exp)).sum(dim=-1)
+        hard_N_sim  = torch.where(N_sim + self.margin > min_P_sim,neg_exp,torch.zeros_like(neg_exp)).sum(dim=-1)
+        pos_loss    = torch.log(1+hard_P_sim).sum()/self.scale_pos
+        neg_loss    = torch.log(1+hard_N_sim).sum()/self.scale_neg
+        
+        return pos_loss + neg_loss
+
+    def ms_sample_cbcb_clcl(self,sim_mat,label):
+        pos_exp = torch.exp(-self.scale_pos*(sim_mat-self.thresh))
+        neg_exp = torch.exp( self.scale_neg*(sim_mat-self.thresh))
+        pos_mask = torch.eq(label.view(-1,1)-label.view(1,-1),0.0).float().cuda()
+        
+        pos_mask = pos_mask + torch.eye(pos_mask.shape[0]).cuda()
+        P_sim = torch.where(pos_mask == 1,sim_mat,torch.ones_like(pos_exp)*1e16)
+        N_sim = torch.where(pos_mask == 0,sim_mat,torch.ones_like(neg_exp)*-1e16)
+        min_P_sim,_ = torch.min(P_sim,dim=1,keepdim=True)
+        max_N_sim,_ = torch.max(N_sim,dim=1,keepdim=True)
+        hard_P_sim = torch.where(P_sim - self.margin < max_N_sim,pos_exp,torch.zeros_like(pos_exp)).sum(dim=-1)
+        hard_N_sim = torch.where(N_sim + self.margin > min_P_sim,neg_exp,torch.zeros_like(neg_exp)).sum(dim=-1)
+        pos_loss = torch.log(1+hard_P_sim).sum()/self.scale_pos
+        neg_loss = torch.log(1+hard_N_sim).sum()/self.scale_neg
+        
+        return pos_loss + neg_loss
+    
+    def ms_sample_cbcb_clcl_trans(self,sim_mat,label):
+        pos_exp = torch.exp(-self.scale_pos*(sim_mat-self.thresh))
+        neg_exp = torch.exp( self.scale_neg*(sim_mat-self.thresh))
+        pos_mask = torch.eq(label.view(-1,1)-label.view(1,-1),0.0).float().cuda()
+        
+        n_sha = pos_mask.shape[0]
+        mask_pos = torch.ones(n_sha, n_sha, dtype=torch.bool)
+        mask_pos = mask_pos.triu(1) | mask_pos.tril(-1)
+        pos_mask = torch.transpose(torch.transpose(pos_mask[mask_pos].reshape(n_sha, n_sha-1),0,1),0,1)
+
+        neg_mask = 1-pos_mask
+        P_sim = torch.where(pos_mask == 1,sim_mat,torch.ones_like(pos_exp)*1e16)
+        N_sim = torch.where(neg_mask == 1,sim_mat,torch.ones_like(neg_exp)*-1e16)
+        min_P_sim,_ = torch.min(P_sim,dim=1,keepdim=True)
+        max_N_sim,_ = torch.max(N_sim,dim=1,keepdim=True)
+        hard_P_sim = torch.where(P_sim - self.margin < max_N_sim,pos_exp,torch.zeros_like(pos_exp)).sum(dim=-1)
+        hard_N_sim = torch.where(N_sim + self.margin > min_P_sim,neg_exp,torch.zeros_like(neg_exp)).sum(dim=-1)
+        pos_loss = torch.log(1+hard_P_sim).sum()/self.scale_pos
+        neg_loss = torch.log(1+hard_N_sim).sum()/self.scale_neg
+        
+        return pos_loss + neg_loss
+
+    def compute_sharded_cosine_similarity(self, tensor1, tensor2, shard_size):
+        B, T, D = tensor1.shape
+        average_sim_matrix = torch.zeros((B, B), device=tensor1.device)
+
+        for start_idx in range(0, T, shard_size):
+            end_idx = min(start_idx + shard_size, T)
+
+            # Get the shard
+            shard_tensor1 = tensor1[:, start_idx:end_idx, :]
+            shard_tensor2 = tensor2[:, start_idx:end_idx, :]
+
+            # Reshape and expand
+            shard_tensor1_expanded = shard_tensor1.unsqueeze(1).unsqueeze(3)
+            shard_tensor2_expanded = shard_tensor2.unsqueeze(0).unsqueeze(2)
+
+            # Compute cosine similarity for the shard
+            shard_cos_sim = F.cosine_similarity(shard_tensor1_expanded, shard_tensor2_expanded, dim=-1)
+
+            # Accumulate the sum of cosine similarities
+            average_sim_matrix += torch.sum(shard_cos_sim, dim=[2, 3])
+
+        # Normalize by the total number of elements (T*T)
+        average_sim_matrix /= (T * T)
+
+        return average_sim_matrix
+
+    def forward(self, x_contactless, x_contactbased, x_cl_tokens, x_cb_tokens, labels, device):
+                
+        sim_mat_clcl = F.linear(self.l2_norm(x_contactless), self.l2_norm(x_contactless))
+        n = sim_mat_clcl.shape[0]
+        sim_mat_cbcb = F.linear(self.l2_norm(x_contactbased), self.l2_norm(x_contactbased))
+        sim_mat_cbcl = F.linear(self.l2_norm(x_contactbased), self.l2_norm(x_contactless))
+
+        loss2              = self.ms_sample_cbcb_clcl(sim_mat_clcl, labels).cuda() + self.ms_sample_cbcb_clcl(sim_mat_clcl.t(), labels).cuda()
+        loss3              = self.ms_sample_cbcb_clcl(sim_mat_cbcb, labels).cuda() + self.ms_sample_cbcb_clcl(sim_mat_cbcb.t(), labels).cuda()
+
+        loss4              = self.ms_sample(sim_mat_cbcl, labels).cuda() + self.ms_sample(sim_mat_cbcl.t(), labels).cuda()
+        return loss4 + loss2 + loss3#+ (1.5*loss2) + (1.5*loss3)  # + loss2 + loss3#+ loss5 # 0.1*loss5  + loss3
+
+    def l2_norm(self, input):
+        input_size = input.size()
+        buffer = torch.pow(input, 2)
+        normp = torch.sum(buffer, 1).add_(1e-12)
+        norm = torch.sqrt(normp)
+        _output = torch.div(input, norm.view(-1, 1).expand_as(input))
+        output = _output.view(input_size)
+        
+        return output
+        
+class DualMSLoss_FineGrained_domain_agnostic(nn.Module):
+    """
+    Compute contrastive loss
+    """
+    def __init__(self, margin=0, max_violation=False):
+        super(DualMSLoss_FineGrained_domain_agnostic, self).__init__()
+        self.margin = margin
+        self.max_violation = max_violation
+        self.thresh = 0.5
+        self.margin = 0.5 # 0.1
+        self.scale_pos = 2
+        self.scale_neg = 40.0
+        self.criterion = nn.CrossEntropyLoss()
+
+    def ms_sample(self,sim_mat,label):
+        pos_exp     = torch.exp(-self.scale_pos*(sim_mat-self.thresh))
+        neg_exp     = torch.exp( self.scale_neg*(sim_mat-self.thresh))
+        pos_mask    = torch.eq(label.view(-1,1)-label.view(1,-1),0.0).float().cuda()
+        neg_mask    = 1 - pos_mask
+        P_sim       = torch.where(pos_mask == 1,sim_mat,torch.ones_like(pos_exp)*1e16)
+        N_sim       = torch.where(neg_mask == 1,sim_mat,torch.ones_like(neg_exp)*-1e16)
+        min_P_sim,_ = torch.min(P_sim,dim=1,keepdim=True)
+        max_N_sim,_ = torch.max(N_sim,dim=1,keepdim=True)
+        hard_P_sim  = torch.where(P_sim - self.margin < max_N_sim,pos_exp,torch.zeros_like(pos_exp)).sum(dim=-1)
+        hard_N_sim  = torch.where(N_sim + self.margin > min_P_sim,neg_exp,torch.zeros_like(neg_exp)).sum(dim=-1)
+        pos_loss    = torch.log(1+hard_P_sim).sum()/self.scale_pos
+        neg_loss    = torch.log(1+hard_N_sim).sum()/self.scale_neg
+        
+        return pos_loss + neg_loss
+
+    def ms_sample_cbcb_clcl(self,sim_mat,label):
+        pos_exp = torch.exp(-self.scale_pos*(sim_mat-self.thresh))
+        neg_exp = torch.exp( self.scale_neg*(sim_mat-self.thresh))
+        pos_mask = torch.eq(label.view(-1,1)-label.view(1,-1),0.0).float().cuda()
+        
+        pos_mask = pos_mask + torch.eye(pos_mask.shape[0]).cuda()
+        P_sim = torch.where(pos_mask == 1,sim_mat,torch.ones_like(pos_exp)*1e16)
+        N_sim = torch.where(pos_mask == 0,sim_mat,torch.ones_like(neg_exp)*-1e16)
+        min_P_sim,_ = torch.min(P_sim,dim=1,keepdim=True)
+        max_N_sim,_ = torch.max(N_sim,dim=1,keepdim=True)
+        hard_P_sim = torch.where(P_sim - self.margin < max_N_sim,pos_exp,torch.zeros_like(pos_exp)).sum(dim=-1)
+        hard_N_sim = torch.where(N_sim + self.margin > min_P_sim,neg_exp,torch.zeros_like(neg_exp)).sum(dim=-1)
+        pos_loss = torch.log(1+hard_P_sim).sum()/self.scale_pos
+        neg_loss = torch.log(1+hard_N_sim).sum()/self.scale_neg
+        
+        return pos_loss + neg_loss
+    
+    def ms_sample_cbcb_clcl_trans(self,sim_mat,label):
+        pos_exp = torch.exp(-self.scale_pos*(sim_mat-self.thresh))
+        neg_exp = torch.exp( self.scale_neg*(sim_mat-self.thresh))
+        pos_mask = torch.eq(label.view(-1,1)-label.view(1,-1),0.0).float().cuda()
+        
+        n_sha = pos_mask.shape[0]
+        mask_pos = torch.ones(n_sha, n_sha, dtype=torch.bool)
+        mask_pos = mask_pos.triu(1) | mask_pos.tril(-1)
+        pos_mask = torch.transpose(torch.transpose(pos_mask[mask_pos].reshape(n_sha, n_sha-1),0,1),0,1)
+
+        neg_mask = 1-pos_mask
+        P_sim = torch.where(pos_mask == 1,sim_mat,torch.ones_like(pos_exp)*1e16)
+        N_sim = torch.where(neg_mask == 1,sim_mat,torch.ones_like(neg_exp)*-1e16)
+        min_P_sim,_ = torch.min(P_sim,dim=1,keepdim=True)
+        max_N_sim,_ = torch.max(N_sim,dim=1,keepdim=True)
+        hard_P_sim = torch.where(P_sim - self.margin < max_N_sim,pos_exp,torch.zeros_like(pos_exp)).sum(dim=-1)
+        hard_N_sim = torch.where(N_sim + self.margin > min_P_sim,neg_exp,torch.zeros_like(neg_exp)).sum(dim=-1)
+        pos_loss = torch.log(1+hard_P_sim).sum()/self.scale_pos
+        neg_loss = torch.log(1+hard_N_sim).sum()/self.scale_neg
+        
+        return pos_loss + neg_loss
+
+    def compute_sharded_cosine_similarity(self, tensor1, tensor2, shard_size):
+        B, T, D = tensor1.shape
+        average_sim_matrix = torch.zeros((B, B), device=tensor1.device)
+
+        for start_idx in range(0, T, shard_size):
+            end_idx = min(start_idx + shard_size, T)
+
+            # Get the shard
+            shard_tensor1 = tensor1[:, start_idx:end_idx, :]
+            shard_tensor2 = tensor2[:, start_idx:end_idx, :]
+
+            # Reshape and expand
+            shard_tensor1_expanded = shard_tensor1.unsqueeze(1).unsqueeze(3)
+            shard_tensor2_expanded = shard_tensor2.unsqueeze(0).unsqueeze(2)
+
+            # Compute cosine similarity for the shard
+            shard_cos_sim = F.cosine_similarity(shard_tensor1_expanded, shard_tensor2_expanded, dim=-1)
+
+            # Accumulate the sum of cosine similarities
+            average_sim_matrix += torch.sum(shard_cos_sim, dim=[2, 3])
+
+        # Normalize by the total number of elements (T*T)
+        average_sim_matrix /= (T * T)
+
+        return average_sim_matrix
+
+    def forward(self, x_contactless, x_contactbased, x_cl_tokens, x_cb_tokens, labels, device, domain_class_cl, domain_class_cb, domain_class_cl_gt, domain_class_cb_gt):
+                
+        sim_mat_clcl = F.linear(self.l2_norm(x_contactless), self.l2_norm(x_contactless))
+        n = sim_mat_clcl.shape[0]
+
+        sim_mat_cbcb = F.linear(self.l2_norm(x_contactbased), self.l2_norm(x_contactbased))
+        sim_mat_cbcl = F.linear(self.l2_norm(x_contactbased), self.l2_norm(x_contactless))
+
+        loss2                = self.ms_sample_cbcb_clcl(sim_mat_clcl, labels).cuda() + self.ms_sample_cbcb_clcl(sim_mat_clcl.t(), labels).cuda()
+        loss3                = self.ms_sample_cbcb_clcl(sim_mat_cbcb, labels).cuda() + self.ms_sample_cbcb_clcl(sim_mat_cbcb.t(), labels).cuda()
+
+        loss4                = self.ms_sample(sim_mat_cbcl, labels).cuda() + self.ms_sample(sim_mat_cbcl.t(), labels).cuda()
+
+        pred = torch.cat([domain_class_cl,    domain_class_cb])
+        gt = torch.cat([domain_class_cl_gt,   domain_class_cb_gt])
+        
+        domain_class_loss = self.criterion(pred,gt)
+        return loss4  + loss2 + loss3 + (3*domain_class_loss)
+
+    def l2_norm(self, input):
+        input_size = input.size()
+        buffer = torch.pow(input, 2)
+        normp = torch.sum(buffer, 1).add_(1e-12)
+        norm = torch.sqrt(normp)
+        _output = torch.div(input, norm.view(-1, 1).expand_as(input))
+        output = _output.view(input_size)
+        
+        return output
+    
+class DualMSLoss_FineGrained_domain_agnostic_ft(nn.Module):
+    """
+    Compute contrastive loss
+    """
+    def __init__(self, margin=0, max_violation=False):
+        super(DualMSLoss_FineGrained_domain_agnostic_ft, self).__init__()
+        self.margin = margin
+        self.max_violation = max_violation
+        self.thresh = 0.5
+        self.margin = 0.7 # 0.1
+        self.scale_pos = 2
+        self.scale_neg = 40.0
+        self.criterion = nn.CrossEntropyLoss()
+
+    def ms_sample(self,sim_mat,label):
+        pos_exp     = torch.exp(-self.scale_pos*(sim_mat-self.thresh))
+        neg_exp     = torch.exp( self.scale_neg*(sim_mat-self.thresh))
+        pos_mask    = torch.eq(label.view(-1,1)-label.view(1,-1),0.0).float().cuda()
+        neg_mask    = 1 - pos_mask
+        P_sim       = torch.where(pos_mask == 1,sim_mat,torch.ones_like(pos_exp)*1e16)
+        N_sim       = torch.where(neg_mask == 1,sim_mat,torch.ones_like(neg_exp)*-1e16)
+        min_P_sim,_ = torch.min(P_sim,dim=1,keepdim=True)
+        max_N_sim,_ = torch.max(N_sim,dim=1,keepdim=True)
+        hard_P_sim  = torch.where(P_sim - self.margin < max_N_sim,pos_exp,torch.zeros_like(pos_exp)).sum(dim=-1)
+        hard_N_sim  = torch.where(N_sim + self.margin > min_P_sim,neg_exp,torch.zeros_like(neg_exp)).sum(dim=-1)
+        pos_loss    = torch.log(1+hard_P_sim).sum()/self.scale_pos
+        neg_loss    = torch.log(1+hard_N_sim).sum()/self.scale_neg
+        
+        return pos_loss + neg_loss
+
+    def ms_sample_cbcb_clcl(self,sim_mat,label):
+        pos_exp = torch.exp(-self.scale_pos*(sim_mat-self.thresh))
+        neg_exp = torch.exp( self.scale_neg*(sim_mat-self.thresh))
+        pos_mask = torch.eq(label.view(-1,1)-label.view(1,-1),0.0).float().cuda()
+        
+        pos_mask = pos_mask + torch.eye(pos_mask.shape[0]).cuda()
+        P_sim = torch.where(pos_mask == 1,sim_mat,torch.ones_like(pos_exp)*1e16)
+        N_sim = torch.where(pos_mask == 0,sim_mat,torch.ones_like(neg_exp)*-1e16)
+        min_P_sim,_ = torch.min(P_sim,dim=1,keepdim=True)
+        max_N_sim,_ = torch.max(N_sim,dim=1,keepdim=True)
+        hard_P_sim = torch.where(P_sim - self.margin < max_N_sim,pos_exp,torch.zeros_like(pos_exp)).sum(dim=-1)
+        hard_N_sim = torch.where(N_sim + self.margin > min_P_sim,neg_exp,torch.zeros_like(neg_exp)).sum(dim=-1)
+        pos_loss = torch.log(1+hard_P_sim).sum()/self.scale_pos
+        neg_loss = torch.log(1+hard_N_sim).sum()/self.scale_neg
+        
+        return pos_loss + neg_loss
+    
+    def ms_sample_cbcb_clcl_trans(self,sim_mat,label):
+        pos_exp = torch.exp(-self.scale_pos*(sim_mat-self.thresh))
+        neg_exp = torch.exp( self.scale_neg*(sim_mat-self.thresh))
+        pos_mask = torch.eq(label.view(-1,1)-label.view(1,-1),0.0).float().cuda()
+        
+        n_sha = pos_mask.shape[0]
+        mask_pos = torch.ones(n_sha, n_sha, dtype=torch.bool)
+        mask_pos = mask_pos.triu(1) | mask_pos.tril(-1)
+        pos_mask = torch.transpose(torch.transpose(pos_mask[mask_pos].reshape(n_sha, n_sha-1),0,1),0,1)
+
+        neg_mask = 1-pos_mask
+        P_sim = torch.where(pos_mask == 1,sim_mat,torch.ones_like(pos_exp)*1e16)
+        N_sim = torch.where(neg_mask == 1,sim_mat,torch.ones_like(neg_exp)*-1e16)
+        min_P_sim,_ = torch.min(P_sim,dim=1,keepdim=True)
+        max_N_sim,_ = torch.max(N_sim,dim=1,keepdim=True)
+        hard_P_sim = torch.where(P_sim - self.margin < max_N_sim,pos_exp,torch.zeros_like(pos_exp)).sum(dim=-1)
+        hard_N_sim = torch.where(N_sim + self.margin > min_P_sim,neg_exp,torch.zeros_like(neg_exp)).sum(dim=-1)
+        pos_loss = torch.log(1+hard_P_sim).sum()/self.scale_pos
+        neg_loss = torch.log(1+hard_N_sim).sum()/self.scale_neg
+        
+        return pos_loss + neg_loss
+
+    def compute_sharded_cosine_similarity(self, tensor1, tensor2, shard_size):
+        B, T, D = tensor1.shape
+        average_sim_matrix = torch.zeros((B, B), device=tensor1.device)
+
+        for start_idx in range(0, T, shard_size):
+            end_idx = min(start_idx + shard_size, T)
+
+            # Get the shard
+            shard_tensor1 = tensor1[:, start_idx:end_idx, :]
+            shard_tensor2 = tensor2[:, start_idx:end_idx, :]
+
+            # Reshape and expand
+            shard_tensor1_expanded = shard_tensor1.unsqueeze(1).unsqueeze(3)
+            shard_tensor2_expanded = shard_tensor2.unsqueeze(0).unsqueeze(2)
+
+            # Compute cosine similarity for the shard
+            shard_cos_sim = F.cosine_similarity(shard_tensor1_expanded, shard_tensor2_expanded, dim=-1)
+
+            # Accumulate the sum of cosine similarities
+            average_sim_matrix += torch.sum(shard_cos_sim, dim=[2, 3])
+
+        # Normalize by the total number of elements (T*T)
+        average_sim_matrix /= (T * T)
+
+        return average_sim_matrix
+
+    def forward(self, x_contactless, x_contactbased, x_cl_tokens, x_cb_tokens, labels, device, domain_class_cl, domain_class_cb, domain_class_cl_gt, domain_class_cb_gt):
+                
+        sim_mat_clcl = F.linear(self.l2_norm(x_contactless), self.l2_norm(x_contactless))
+        n = sim_mat_clcl.shape[0]
+
+        sim_mat_cbcb = F.linear(self.l2_norm(x_contactbased), self.l2_norm(x_contactbased))
+        sim_mat_cbcl = F.linear(self.l2_norm(x_contactbased), self.l2_norm(x_contactless))
+
+        loss2                = self.ms_sample_cbcb_clcl(sim_mat_clcl, labels).cuda() + self.ms_sample_cbcb_clcl(sim_mat_clcl.t(), labels).cuda()
+        loss3                = self.ms_sample_cbcb_clcl(sim_mat_cbcb, labels).cuda() + self.ms_sample_cbcb_clcl(sim_mat_cbcb.t(), labels).cuda()
+
+        loss4                = self.ms_sample(sim_mat_cbcl, labels).cuda() + self.ms_sample(sim_mat_cbcl.t(), labels).cuda()
+
+        return loss4  + loss2 + loss3 
+
+    def l2_norm(self, input):
+        input_size = input.size()
+        buffer = torch.pow(input, 2)
+        normp = torch.sum(buffer, 1).add_(1e-12)
+        norm = torch.sqrt(normp)
+        _output = torch.div(input, norm.view(-1, 1).expand_as(input))
+        output = _output.view(input_size)
+        
+        return output

model.py

@@ -0,0 +1,207 @@
+from __future__ import print_function
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.optim.lr_scheduler import StepLR
+import torchvision.models as models
+import timm
+from pprint import pprint
+import numpy as np
+from tqdm import tqdm
+from torch.utils.data.sampler import BatchSampler
+from gradient_reversal.module import GradientReversal
+
+class SwinModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.swin_cl = timm.create_model('vit_large_patch16_224_in21k', pretrained=True, num_classes=0)
+        self.swin_cb = self.swin_cl
+        
+        self.linear_cl = nn.Sequential(nn.Linear(1024, 1024),
+                                       nn.ReLU(),
+                                       nn.Linear(1024, 1024))
+        self.linear_cb = nn.Linear(1024, 1024)
+
+    def freeze_encoder(self):
+        for param in self.swin_cl.parameters():
+            param.requires_grad = False
+        for param in self.swin_cb.parameters():
+            param.requires_grad = False
+
+    def unfreeze_encoder(self):
+        for param in self.swin_cl.parameters():
+            param.requires_grad = True
+        for param in self.swin_cb.parameters():
+            param.requires_grad = True
+
+    def get_embeddings(self, image, ftype):
+        linear = self.linear_cl if ftype == "contactless" else self.linear_cl
+        swin   = self.swin_cl   if ftype == "contactless" else self.swin_cb
+        
+        tokens = swin(image)
+        emb_mean = tokens.mean(dim=1)
+        feat = linear(emb_mean)
+        tokens_transformed = linear(tokens)
+        return feat, tokens
+
+    def forward(self, x_cl, x_cb):
+        x_cl_tokens = self.swin_cl(x_cl)
+        x_cb_tokens = self.swin_cb(x_cb)
+
+        x_cl_mean = x_cl_tokens.mean(dim=1)
+        x_cb_mean = x_cb_tokens.mean(dim=1)
+
+        x_cl = self.linear_cl(x_cl_mean)
+        x_cl_tokens_transformed = self.linear_cl(x_cl_tokens)
+
+        x_cb = self.linear_cl(x_cb_mean)
+        x_cb_tokens_transformed = self.linear_cl(x_cb_tokens)
+
+        return x_cl, x_cb, x_cl_tokens, x_cb_tokens
+
+class SwinModel_domain_agnostic(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.swin_cl = timm.create_model('vit_large_patch16_224_in21k', pretrained=True, num_classes=0)
+        self.swin_cb = self.swin_cl #timm.create_model('vit_large_patch16_224_in21k', pretrained=True, num_classes=0)
+        
+        self.linear_cl = nn.Sequential(nn.Linear(1024, 1024),
+                                       nn.ReLU(),
+                                       nn.Linear(1024, 1024))
+        self.linear_cb = nn.Linear(1024, 1024)
+        self.classify = nn.Sequential(GradientReversal(alpha=0.6),  # original 0.8
+                                      nn.Linear(1024,512),
+                                      nn.ReLU(),
+                                      nn.Linear(512,8))
+
+    def freeze_encoder(self):
+        for param in self.swin_cl.parameters():
+            param.requires_grad = False
+        for param in self.swin_cb.parameters():
+            param.requires_grad = False
+
+    def unfreeze_encoder(self):
+        for param in self.swin_cl.parameters():
+            param.requires_grad = True
+        for param in self.swin_cb.parameters():
+            param.requires_grad = True
+
+    def get_embeddings(self, image, ftype):
+        linear = self.linear_cl if ftype == "contactless" else self.linear_cl
+        swin   = self.swin_cl   if ftype == "contactless" else self.swin_cb
+        
+        tokens = swin(image)
+        emb_mean = tokens.mean(dim=1)
+        feat = linear(emb_mean)
+        tokens_transformed = linear(tokens)
+        return feat, tokens
+
+    def forward(self, x_cl, x_cb):
+        x_cl_tokens = self.swin_cl(x_cl)
+        x_cb_tokens = self.swin_cb(x_cb)
+
+        x_cl_mean = x_cl_tokens.mean(dim=1)
+        x_cb_mean = x_cb_tokens.mean(dim=1)
+
+        x_cl = self.linear_cl(x_cl_mean)
+        x_cl_tokens_transformed = self.linear_cl(x_cl_tokens)
+
+        x_cb = self.linear_cl(x_cb_mean)
+        x_cb_tokens_transformed = self.linear_cl(x_cb_tokens)
+
+        domain_class_cl = self.classify(x_cl_mean)
+        domain_class_cb = self.classify(x_cb_mean)
+
+        return x_cl, x_cb, x_cl_tokens, x_cb_tokens, domain_class_cl, domain_class_cb
+
+class SwinModel_Fusion(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.feature_dim        = 1024
+        self.swin_cl            = timm.create_model('vit_large_patch16_224_in21k', pretrained=True, num_classes=0)  
+        self.encoder_layer      = nn.TransformerEncoderLayer(d_model=self.feature_dim, nhead=4, dropout=0.5, batch_first=True, norm_first=True, activation="gelu")
+        self.fusion             = nn.TransformerEncoder(self.encoder_layer, num_layers=2)
+        self.sep_token          = nn.Parameter(torch.randn(1, 1, self.feature_dim))
+        self.output_logit_mlp   = nn.Sequential(nn.Linear(1024, 512),
+                                       nn.ReLU(),
+                                       nn.Dropout(),
+                                       nn.Linear(512, 1))
+        self.linear_cl          = nn.Sequential(nn.Linear(1024, 1024),
+                                       nn.ReLU(),
+                                       nn.Linear(1024, 1024))
+        
+    def load_pretrained_models(self, swin_cl_path, fusion_ckpt_path):
+        swin_cl_state_dict = torch.load(swin_cl_path)
+        new_dict = {}
+        for key in swin_cl_state_dict.keys():
+            if "swin_cl" in key:
+                new_dict[key.replace("swin_cl.","")] = swin_cl_state_dict[key] 
+        self.swin_cl.load_state_dict(new_dict)
+
+        fusion_params = torch.load(fusion_ckpt_path)
+        new_dict = {}
+        for key in fusion_params.keys():
+            if "encoder_layer" in key:
+                new_dict[key.replace("encoder_layer.","")] = fusion_params[key] 
+        self.encoder_layer.load_state_dict(new_dict)
+        
+        new_dict = {}
+        for key in fusion_params.keys():
+            if "fusion" in key:
+                new_dict[key.replace("fusion.","")] = fusion_params[key] 
+        self.fusion.load_state_dict(new_dict)
+        
+        self.sep_token = nn.Parameter(fusion_params["sep_token"])   
+
+        new_dict = {}
+        for key in fusion_params.keys():
+            if "output_logit_mlp" in key:
+                new_dict[key.replace("output_logit_mlp.","")] = fusion_params[key] 
+        self.output_logit_mlp.load_state_dict(new_dict)
+
+    def l2_norm(self,input):
+        input_size = input.shape[0]
+        buffer     = torch.pow(input, 2)
+        normp      = torch.sum(buffer, 1).add_(1e-12)
+        norm       = torch.sqrt(normp)
+        _output    = torch.div(input, norm.view(-1, 1).expand_as(input))
+        return _output
+
+    def combine_features(self, fingerprint_1_tokens, fingerprint_2_tokens):
+        # This function takes a pair of embeddings [B, 49, 1024], [B, 49, 1024] and returns a B logit scores [B]
+        # fingerprint_1_tokens        = self.linear_cl(fingerprint_1_tokens)
+        # fingerprint_2_tokens        = self.linear_cl(fingerprint_2_tokens)
+        batch_size                  = fingerprint_1_tokens.shape[0]
+        sep_token                   = self.sep_token.repeat(batch_size, 1, 1)
+        combine_features            = torch.cat((fingerprint_1_tokens, sep_token, fingerprint_2_tokens), dim=1)            
+        fused_match_representation  = self.fusion(combine_features) 
+        fingerprint_1 = fused_match_representation[:,:197,:].mean(dim=1)
+        fingerprint_2 = fused_match_representation[:,198:,:].mean(dim=1)
+
+        fingerprint_1_norm = self.l2_norm(fingerprint_1)
+        fingerprint_2_norm = self.l2_norm(fingerprint_2)
+        
+        similarities = torch.sum(fingerprint_1_norm * fingerprint_2_norm, axis=1)
+
+        differences  = fingerprint_1 - fingerprint_2
+        squared_differences = differences ** 2
+        sum_squared_differences = torch.sum(squared_differences, axis=1)
+        distances = torch.sqrt(sum_squared_differences)
+        return similarities, distances
+    
+    def get_tokens(self, image, ftype):
+        swin   = self.swin_cl
+        tokens = swin(image)
+        return tokens
+
+    def freeze_backbone(self):
+        for param in self.swin_cl.parameters():
+            param.requires_grad = False
+
+    def forward(self, x_cl, x_cb):
+        x_cl_tokens = self.swin_cl(x_cl)
+        x_cb_tokens = self.swin_cl(x_cb)
+        return x_cl_tokens, x_cb_tokens

rb_evaluation_phase1.py

@@ -0,0 +1,148 @@
+import torch
+from datasets.rb_loader import RB_loader
+from utils import Prev_RetMetric, l2_norm, compute_recall_at_k
+import numpy as np
+from tqdm import tqdm
+from model import SwinModel_domain_agnostic as Model
+from sklearn.metrics import roc_curve, auc
+import json
+import torch.nn.functional as F
+
+if __name__ == '__main__':
+    device = torch.device('cuda')
+    data = RB_loader(split = 'test')
+    dataloader = torch.utils.data.DataLoader(data,batch_size = 16, num_workers = 1, pin_memory = True)
+    model = Model().to(device)
+    checkpoint = torch.load("ridgeformer_checkpoints/phase1_scratch.pt",map_location = torch.device('cpu'))
+    model.load_state_dict(checkpoint,strict=False)
+
+    model.eval()
+    cl_feats, cb_feats, cl_labels, cb_labels, cl_fnames, cb_fnames, cl_feats_unnormed, cb_feats_unnormed = list(),list(),list(),list(),list(),list(),list(),list()
+    print("Computing Test Recall")
+    with torch.no_grad():
+        for (x_cl, x_cb, target, cl_fname, cb_fname) in tqdm(dataloader):
+            x_cl, x_cb, target = x_cl.to(device), x_cb.to(device), target.to(device)
+            x_cl, _ = model.get_embeddings(x_cl, ftype="contactless")
+            x_cb, _ = model.get_embeddings(x_cb, ftype="contactbased")
+            cl_feats_unnormed.append(x_cl.cpu().detach().numpy())
+            cb_feats_unnormed.append(x_cb.cpu().detach().numpy())
+            x_cl = l2_norm(x_cl).cpu().detach().numpy()
+            x_cb = l2_norm(x_cb).cpu().detach().numpy()
+            target = target.cpu().detach().numpy()
+            cl_feats.append(x_cl)
+            cb_feats.append(x_cb)
+            cl_labels.append(target)
+            cb_labels.append(target)
+            cl_fnames.extend(cl_fname)
+            cb_fnames.extend(cb_fname)
+
+    cl_feats = torch.from_numpy(np.concatenate(cl_feats))
+    cb_feats = torch.from_numpy(np.concatenate(cb_feats))
+    cl_labels = torch.from_numpy(np.concatenate(cl_labels))
+    cb_labels = torch.from_numpy(np.concatenate(cb_labels))
+    cl_feats_unnormed = torch.from_numpy(np.concatenate(cl_feats_unnormed))
+    cb_feats_unnormed = torch.from_numpy(np.concatenate(cb_feats_unnormed))
+
+    unique_labels, indices  = torch.unique(cb_labels, return_inverse=True)
+    unique_feats            = torch.stack([cb_feats[indices == i].mean(dim=0) for i in range(len(unique_labels))])
+    cb_feats                = unique_feats
+    unique_labels, indices  = torch.unique(cb_labels, return_inverse=True)
+    unique_feats            = torch.stack([cb_feats_unnormed[indices == i].mean(dim=0) for i in range(len(unique_labels))])
+    cb_labels               = unique_labels
+    cb_feats_unnormed       = unique_feats
+
+    # CL2CB <---------------------------------------->
+    cl_feats  = cl_feats.numpy()
+    cb_feats  = cb_feats.numpy()
+    cb_feats_unnormed = cb_feats_unnormed.numpy()
+    cl_feats_unnormed = cl_feats_unnormed.numpy()
+
+    squared_diff = np.sum(np.square(cl_feats_unnormed[:, np.newaxis] - cb_feats_unnormed), axis=2)
+    distance     = -1 * np.sqrt(squared_diff)
+    similarities = np.dot(cl_feats,np.transpose(cb_feats))
+    scores_mat = similarities + 0.1 * distance
+    scores = scores_mat.flatten().tolist()
+
+    ids = torch.eq(cl_labels.view(-1,1)-cb_labels.view(1,-1),0.0).flatten().tolist()
+    ids_mod = list()
+    for x in ids:
+        if x==True:
+            ids_mod.append(1)
+        else:
+            ids_mod.append(0)
+    fpr,tpr,thresh = roc_curve(ids_mod,scores,drop_intermediate=True)
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.01)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.01)
+    tar_far_102 = tpr[upper_fpr_idx]
+    print(tpr[lower_fpr_idx], lower_fpr_idx, fpr[lower_fpr_idx], thresh[lower_fpr_idx])
+    print(tpr[upper_fpr_idx], upper_fpr_idx, fpr[upper_fpr_idx], thresh[upper_fpr_idx])
+
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.001)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.001)
+    tar_far_103 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    print(tpr[lower_fpr_idx], lower_fpr_idx, fpr[lower_fpr_idx])
+    print(tpr[upper_fpr_idx], upper_fpr_idx, fpr[upper_fpr_idx])
+
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.0001)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.0001)
+    tar_far_104 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    print(tpr[lower_fpr_idx], lower_fpr_idx, fpr[lower_fpr_idx])
+    print(tpr[upper_fpr_idx], upper_fpr_idx, fpr[upper_fpr_idx])
+
+    fnr = 1 - tpr
+    EER = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+    roc_auc = auc(fpr, tpr)
+    print(f"ROCAUC for CB2CL: {roc_auc * 100} %")
+    print(f"EER for CB2CL: {EER * 100} %")
+    eer_cb2cl = EER * 100
+    cbcltf102 = tar_far_102 * 100
+    cbcltf103 = tar_far_103 * 100
+    cbcltf104 = tar_far_104 * 100
+    cl_labels = cl_labels.cpu().detach()
+    cb_labels = cb_labels.cpu().detach()
+    print(f"TAR@FAR=10^-2 for CB2CL: {tar_far_102 * 100} %")
+    print(f"TAR@FAR=10^-3 for CB2CL: {tar_far_103 * 100} %")
+    print(f"TAR@FAR=10^-4 for CB2CL: {tar_far_104 * 100} %")
+    print(f"R@1 for CB2CL: {compute_recall_at_k(torch.from_numpy(scores_mat), cl_labels, cb_labels, 1) * 100} %")
+    print(f"R@10 for CB2CL: {compute_recall_at_k(torch.from_numpy(scores_mat), cl_labels, cb_labels, 10) * 100} %")
+    print(f"R@50 for CB2CL: {compute_recall_at_k(torch.from_numpy(scores_mat), cl_labels, cb_labels, 50) * 100} %")
+    print(f"R@100 for CB2CL: {compute_recall_at_k(torch.from_numpy(scores_mat), cl_labels, cb_labels, 100) * 100} %")
+
+    ################################################################################
+
+    # CL2CL
+    scores = torch.from_numpy(np.dot(cl_feats,np.transpose(cl_feats)))
+    row, col = torch.triu_indices(row=scores.size(0), col=scores.size(1), offset=1)
+    scores = scores[row, col]
+    scores = scores.numpy().flatten().tolist()
+    labels = torch.eq(cl_labels.view(-1,1) - cl_labels.view(1,-1),0.0).float().cuda()
+    labels = labels[torch.triu(torch.ones(labels.shape),diagonal = 1) == 1].tolist()
+    fpr,tpr,_ = roc_curve(labels,scores)
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.01)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.01)
+    tar_far_102 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.001)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.001)
+    tar_far_103 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.0001)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.0001)
+    tar_far_104 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    clcltf102 = tar_far_102 * 100
+    clcltf103 = tar_far_103 * 100
+    clcltf104 = tar_far_104 * 100
+    fnr = 1 - tpr
+    EER = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+    roc_auc = auc(fpr, tpr)
+    print(f"ROCAUC for CL2CL: {roc_auc * 100} %")
+    print(f"EER for CL2CL: {EER * 100} %")
+    eer_cl2cl = EER * 100
+    print(f"TAR@FAR=10^-2 for CL2CL: {tar_far_102 * 100} %")
+    print(f"TAR@FAR=10^-3 for CL2CL: {tar_far_103 * 100} %")
+    print(f"TAR@FAR=10^-4 for CL2CL: {tar_far_104 * 100} %")
+    cl_labels = cl_labels.cpu().detach().numpy()
+    recall_score = Prev_RetMetric([cl_feats,cl_feats],[cl_labels,cl_labels],cl2cl = True)
+    cl2clk1 = recall_score.recall_k(k=1) * 100
+    print(f"R@1 for CL2CL: {recall_score.recall_k(k=1) * 100} %")
+    print(f"R@10 for CL2CL: {recall_score.recall_k(k=10) * 100} %")
+    print(f"R@50 for CL2CL: {recall_score.recall_k(k=50) * 100} %")
+    print(f"R@100 for CL2CL: {recall_score.recall_k(k=100) * 100} %")

rb_evaluation_phase2.py

@@ -0,0 +1,164 @@
+import torch
+from datasets.rb_loader_cl import RB_loader_cl
+from datasets.rb_loader_cb import RB_loader_cb
+from utils import Prev_RetMetric, l2_norm, compute_recall_at_k
+import numpy as np
+from tqdm import tqdm
+from model import SwinModel_Fusion as Model
+from sklearn.metrics import roc_curve, auc
+import json
+import torch.nn.functional as F
+
+def get_fused_cross_score_matrix(model, cl_tokens, cb_tokens):
+    cl_tokens   = torch.cat(cl_tokens)
+    cb_tokens   = torch.cat(cb_tokens)
+    
+    batch_size_cl = cl_tokens.shape[0]
+    batch_size_cb = cb_tokens.shape[0]
+    shard_size  = 20
+    similarity_matrix = torch.zeros((batch_size_cl, batch_size_cb))
+    for i_start in tqdm(range(0, batch_size_cl, shard_size)):
+        i_end   = min(i_start + shard_size, batch_size_cl)
+        shard_i = cl_tokens[i_start:i_end]
+        for j_start in range(0, batch_size_cb, shard_size):
+            j_end               = min(j_start + shard_size, batch_size_cb)
+            shard_j             = cb_tokens[j_start:j_end]
+            batch_i             = shard_i.unsqueeze(1)
+            batch_j             = shard_j.unsqueeze(0)
+
+            pairwise_i          = batch_i.expand(-1, shard_j.shape[0], -1, -1)
+            pairwise_j          = batch_j.expand(shard_i.shape[0], -1, -1, -1)
+            
+            similarity_scores, distances = model.combine_features(
+                pairwise_i.reshape(-1, 197, shard_i.shape[-1]), 
+                pairwise_j.reshape(-1, 197, shard_j.shape[-1])
+            )
+            scores = similarity_scores - 0.1 * distances  #-0.1
+            scores   = scores.reshape(shard_i.shape[0], shard_j.shape[0])
+            similarity_matrix[i_start:i_end, j_start:j_end] = scores.cpu().detach()
+    return similarity_matrix
+
+device = torch.device('cuda')
+data_cl = RB_loader_cl(split="test")
+data_cb = RB_loader_cb(split="test")
+dataloader_cb = torch.utils.data.DataLoader(data_cb,batch_size = 16, num_workers = 1, pin_memory = True)
+dataloader_cl = torch.utils.data.DataLoader(data_cl,batch_size = 16, num_workers = 1, pin_memory = True)
+model = Model().to(device)
+checkpoint = torch.load("ridgeformer_checkpoints/phase2_scratch.pt",map_location = torch.device('cpu'))
+model.load_state_dict(checkpoint,strict=False)
+
+model.eval()
+cl_feats, cb_feats, cl_labels, cb_labels, cl_fnames, cb_fnames, cl_feats_unnormed, cb_feats_unnormed = list(),list(),list(),list(),list(),list(),list(),list()
+print("Computing Test Recall")
+
+with torch.no_grad():
+    for (x_cb, target) in tqdm(dataloader_cb):
+        x_cb, label = x_cb.to(device), target.to(device)
+        x_cb_token  = model.get_tokens(x_cb,'contactbased')
+        label = label.cpu().detach().numpy()
+        cb_feats.append(x_cb_token)
+        cb_labels.append(label)
+
+with torch.no_grad():
+    for (x_cl, target) in tqdm(dataloader_cl):
+        x_cl, label = x_cl.to(device), target.to(device)
+        x_cl_token  = model.get_tokens(x_cl,'contactless')
+        label = label.cpu().detach().numpy()
+        cl_feats.append(x_cl_token)
+        cl_labels.append(label)
+
+cl_label = torch.from_numpy(np.concatenate(cl_labels))
+cb_label = torch.from_numpy(np.concatenate(cb_labels))
+
+# CB2CL <---------------------------------------->
+scores_mat = get_fused_cross_score_matrix(model, cl_feats, cb_feats)
+scores = scores_mat.cpu().detach().numpy().flatten().tolist()
+labels = torch.eq(cl_label.view(-1,1) - cb_label.view(1,-1),0.0).flatten().tolist()
+ids_mod = list()
+for i in labels:
+    if i==True:
+        ids_mod.append(1)
+    else:
+        ids_mod.append(0)
+fpr,tpr,thresh = roc_curve(labels,scores,drop_intermediate=True)
+lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.01)
+upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.01)
+tar_far_102 = tpr[upper_fpr_idx]#(tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+
+lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.001)
+upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.001)
+tar_far_103 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+
+lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.0001)
+upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.0001)
+tar_far_104 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+
+fnr = 1 - tpr
+
+EER = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+roc_auc = auc(fpr, tpr)
+print(f"ROCAUC for CB2CL: {roc_auc * 100} %")
+print(f"EER for CB2CL: {EER * 100} %")
+eer_cb2cl = EER * 100
+cbcltf102 = tar_far_102 * 100
+cbcltf103 = tar_far_103 * 100
+cbcltf104 = tar_far_104 * 100
+cl_label = cl_label.cpu().detach()
+cb_label = cb_label.cpu().detach()
+print(f"TAR@FAR=10^-2 for CB2CL: {tar_far_102 * 100} %")
+print(f"TAR@FAR=10^-3 for CB2CL: {tar_far_103 * 100} %")
+print(f"TAR@FAR=10^-4 for CB2CL: {tar_far_104 * 100} %")
+
+print(f"R@1 for CB2CL: {compute_recall_at_k(scores_mat, cl_label, cb_label, 1) * 100} %")
+print(f"R@10 for CB2CL: {compute_recall_at_k(scores_mat, cl_label, cb_label, 10) * 100} %")
+print(f"R@50 for CB2CL: {compute_recall_at_k(scores_mat, cl_label, cb_label, 50) * 100} %")
+print(f"R@100 for CB2CL: {compute_recall_at_k(scores_mat, cl_label, cb_label, 100) * 100} %")
+
+# CL2CL -------------------------
+scores = get_fused_cross_score_matrix(model, cl_feats, cl_feats)
+scores_mat = scores
+row, col = torch.triu_indices(row=scores.size(0), col=scores.size(1), offset=1)
+scores = scores[row, col]
+labels = torch.eq(cl_label.view(-1,1) - cl_label.view(1,-1),0.0).float().cuda()
+labels = labels[torch.triu(torch.ones(labels.shape),diagonal = 1) == 1]
+scores = scores.cpu().detach().numpy().flatten().tolist()
+labels = labels.flatten().tolist()
+ids_mod = list()
+for i in labels:
+    if i==True:
+        ids_mod.append(1)
+    else:
+        ids_mod.append(0)
+fpr,tpr,thresh = roc_curve(labels,scores,drop_intermediate=True)
+
+lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.01)
+upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.01)
+tar_far_102 = tpr[upper_fpr_idx]#(tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+
+lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.001)
+upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.001)
+tar_far_103 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+
+lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.0001)
+upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.0001)
+tar_far_104 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+
+fnr = 1 - tpr
+
+EER = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+roc_auc = auc(fpr, tpr)
+print(f"ROCAUC for CL2CL: {roc_auc * 100} %")
+print(f"EER for CL2CL: {EER * 100} %")
+eer_cb2cl = EER * 100
+cbcltf102 = tar_far_102 * 100
+cbcltf103 = tar_far_103 * 100
+cbcltf104 = tar_far_104 * 100
+cl_label = cl_label.cpu().detach()
+print(f"TAR@FAR=10^-2 for CL2CL: {tar_far_102 * 100} %")
+print(f"TAR@FAR=10^-3 for CL2CL: {tar_far_103 * 100} %")
+print(f"TAR@FAR=10^-4 for CL2CL: {tar_far_104 * 100} %")
+
+print(f"R@1 for CL2CL: {compute_recall_at_k(scores_mat, cl_label, cl_label, 1) * 100} %")
+print(f"R@10 for CL2CL: {compute_recall_at_k(scores_mat, cl_label, cl_label, 10) * 100} %")
+print(f"R@50 for CL2CL: {compute_recall_at_k(scores_mat, cl_label, cl_label, 50) * 100} %")
+print(f"R@100 for CL2CL: {compute_recall_at_k(scores_mat, cl_label, cl_label, 100) * 100} %")

requirements.txt

@@ -0,0 +1,202 @@
+# This file may be used to create an environment using:
+# $ conda create --name <env> --file <this file>
+# platform: linux-64
+_libgcc_mutex=0.1=main
+_openmp_mutex=5.1=1_gnu
+absl-py=2.1.0=pypi_0
+addict=2.4.0=pypi_0
+aliyun-python-sdk-core=2.15.0=pypi_0
+aliyun-python-sdk-kms=2.16.2=pypi_0
+attrs=23.2.0=pypi_0
+blas=1.0=mkl
+bzip2=1.0.8=h5eee18b_6
+ca-certificates=2024.7.2=h06a4308_0
+cachetools=5.4.0=pypi_0
+certifi=2024.2.2=pypi_0
+cffi=1.16.0=pypi_0
+charset-normalizer=2.1.1=pypi_0
+click=8.1.7=pypi_0
+colorama=0.4.6=pypi_0
+coloredlogs=15.0.1=pypi_0
+contourpy=1.1.1=pypi_0
+crcmod=1.7=pypi_0
+cryptography=42.0.5=pypi_0
+cuda-cudart=11.8.89=0
+cuda-cudart_linux-64=12.4.127=hd681fbe_0
+cuda-cupti=11.8.87=0
+cuda-libraries=11.8.0=0
+cuda-nvrtc=11.8.89=0
+cuda-nvtx=11.8.86=0
+cuda-opencl=12.4.127=h6a678d5_0
+cuda-runtime=11.8.0=0
+cuda-version=12.4=hbda6634_3
+cycler=0.12.1=pypi_0
+entrypoints=0.4=pypi_0
+ffmpeg=4.3=hf484d3e_0
+flatbuffers=24.3.25=pypi_0
+fonttools=4.53.1=pypi_0
+freetype=2.12.1=h4a9f257_0
+fsspec=2024.3.1=pypi_0
+gmp=6.2.1=h295c915_3
+gnutls=3.6.15=he1e5248_0
+google-auth=2.33.0=pypi_0
+google-auth-oauthlib=1.0.0=pypi_0
+grpcio=1.65.4=pypi_0
+httpcore=1.0.5=pypi_0
+httpx=0.27.0=pypi_0
+huggingface-hub=0.22.1=pypi_0
+humanfriendly=10.0=pypi_0
+idna=3.6=pypi_0
+imageio=2.34.2=pypi_0
+importlib-metadata=7.1.0=pypi_0
+importlib-resources=6.4.0=pypi_0
+intel-openmp=2023.1.0=hdb19cb5_46306
+jinja2=3.1.3=pypi_0
+jmespath=0.10.0=pypi_0
+joblib=1.4.2=pypi_0
+jpeg=9e=h5eee18b_2
+jsonschema=4.21.1=pypi_0
+jsonschema-specifications=2023.12.1=pypi_0
+kaleido=0.2.1=pypi_0
+kiwisolver=1.4.5=pypi_0
+lame=3.100=h7b6447c_0
+lcms2=2.12=h3be6417_0
+ld_impl_linux-64=2.38=h1181459_1
+lerc=3.0=h295c915_0
+libcublas=11.11.3.6=0
+libcufft=10.9.0.58=0
+libcufile=1.9.1.3=h99ab3db_1
+libcurand=10.3.5.147=h99ab3db_1
+libcusolver=11.4.1.48=0
+libcusparse=11.7.5.86=0
+libdeflate=1.17=h5eee18b_1
+libffi=3.4.4=h6a678d5_1
+libgcc-ng=11.2.0=h1234567_1
+libgomp=11.2.0=h1234567_1
+libiconv=1.16=h5eee18b_3
+libidn2=2.3.4=h5eee18b_0
+libjpeg-turbo=2.0.0=h9bf148f_0
+libnpp=11.8.0.86=0
+libnvfatbin=12.4.127=h7934f7d_2
+libnvjitlink=12.4.99=0
+libnvjpeg=11.9.0.86=0
+libpng=1.6.39=h5eee18b_0
+libstdcxx-ng=11.2.0=h1234567_1
+libtasn1=4.19.0=h5eee18b_0
+libtiff=4.5.1=h6a678d5_0
+libunistring=0.9.10=h27cfd23_0
+libwebp-base=1.3.2=h5eee18b_0
+llvm-openmp=14.0.6=h9e868ea_0
+llvmlite=0.41.1=pypi_0
+lz4-c=1.9.4=h6a678d5_1
+markdown=3.6=pypi_0
+markdown-it-py=3.0.0=pypi_0
+markupsafe=2.1.5=pypi_0
+matplotlib=3.7.5=pypi_0
+mdit-py-plugins=0.4.0=pypi_0
+mkl=2023.1.0=h213fc3f_46344
+mmcv=2.1.0=dev_0
+mmdet=3.3.0=dev_0
+mmengine=0.10.3=pypi_0
+model-index=0.1.11=pypi_0
+mpc=1.1.0=h10f8cd9_1
+mpfr=4.0.2=hb69a4c5_1
+mpmath=1.3.0=py38h06a4308_0
+ncurses=6.4=h6a678d5_0
+nettle=3.7.3=hbbd107a_1
+networkx=3.1=py38h06a4308_0
+numba=0.58.1=pypi_0
+numpy=1.24.4=pypi_0
+nvidia-cublas-cu11=11.11.3.6=pypi_0
+nvidia-cuda-cupti-cu11=11.8.87=pypi_0
+nvidia-cuda-nvrtc-cu11=11.8.89=pypi_0
+nvidia-cuda-runtime-cu11=11.8.89=pypi_0
+nvidia-cudnn-cu11=8.7.0.84=pypi_0
+nvidia-cufft-cu11=10.9.0.58=pypi_0
+nvidia-curand-cu11=10.3.0.86=pypi_0
+nvidia-cusolver-cu11=11.4.1.48=pypi_0
+nvidia-cusparse-cu11=11.7.5.86=pypi_0
+nvidia-nccl-cu11=2.19.3=pypi_0
+nvidia-nvtx-cu11=11.8.86=pypi_0
+oauthlib=3.2.2=pypi_0
+ocl-icd=2.3.2=h5eee18b_1
+onnxruntime=1.18.1=pypi_0
+opencv-python=4.10.0.84=pypi_0
+opencv-python-headless=4.10.0.84=pypi_0
+opendatalab=0.0.10=pypi_0
+openh264=2.1.1=h4ff587b_0
+openjpeg=2.4.0=h9ca470c_2
+openmim=0.3.9=pypi_0
+openssl=3.0.14=h5eee18b_0
+openxlab=0.0.37=pypi_0
+ordered-set=4.1.0=pypi_0
+oss2=2.17.0=pypi_0
+packaging=24.0=pypi_0
+pandas=2.0.3=pypi_0
+pillow=9.0.1=pypi_0
+pip=23.3.1=pypi_0
+pkgutil-resolve-name=1.3.10=pypi_0
+platformdirs=4.2.0=pypi_0
+plotly=5.23.0=pypi_0
+pooch=1.8.2=pypi_0
+protobuf=5.27.3=pypi_0
+pyasn1=0.6.0=pypi_0
+pyasn1-modules=0.4.0=pypi_0
+pycocotools=2.0.7=pypi_0
+pycparser=2.21=pypi_0
+pygments=2.17.2=pypi_0
+pymatting=1.1.12=pypi_0
+pyparsing=3.1.2=pypi_0
+python=3.8.19=h955ad1f_0
+python-dateutil=2.9.0.post0=pypi_0
+pytorch-cuda=11.8=h7e8668a_5
+pytorch-metric-learning=2.5.0=pypi_0
+pytorch-mutex=1.0=cuda
+pytz=2023.4=pypi_0
+pywavelets=1.4.1=pypi_0
+pyyaml=6.0.1=py38h5eee18b_0
+readline=8.2=h5eee18b_0
+referencing=0.34.0=pypi_0
+rembg=2.0.58=pypi_0
+requests=2.28.2=pypi_0
+requests-oauthlib=2.0.0=pypi_0
+rich=13.4.2=pypi_0
+rpds-py=0.18.0=pypi_0
+rsa=4.9=pypi_0
+safetensors=0.4.2=pypi_0
+scikit-image=0.19.3=pypi_0
+scikit-learn=1.3.2=pypi_0
+scipy=1.10.1=pypi_0
+setuptools=60.2.0=pypi_0
+shapely=2.0.3=pypi_0
+six=1.16.0=pypi_0
+sqlite=3.45.3=h5eee18b_0
+sympy=1.12=py38h06a4308_0
+tabulate=0.9.0=pypi_0
+tbb=2021.8.0=hdb19cb5_0
+tenacity=9.0.0=pypi_0
+tensorboard=2.14.0=pypi_0
+tensorboard-data-server=0.7.2=pypi_0
+termcolor=2.4.0=pypi_0
+terminaltables=3.1.10=pypi_0
+threadpoolctl=3.5.0=pypi_0
+tifffile=2023.7.10=pypi_0
+timm=0.5.0=dev_0
+tk=8.6.14=h39e8969_0
+tomli=2.0.1=pypi_0
+torch=2.2.2+cu118=pypi_0
+torchaudio=2.2.2+cu118=pypi_0
+torchvision=0.17.2+cu118=pypi_0
+tqdm=4.66.5=pypi_0
+triton=2.2.0=pypi_0
+typing-extensions=4.10.0=pypi_0
+tzdata=2024.1=pypi_0
+urllib3=1.26.18=pypi_0
+werkzeug=3.0.3=pypi_0
+wheel=0.41.2=pypi_0
+xz=5.4.6=h5eee18b_1
+yaml=0.2.5=h7b6447c_0
+yapf=0.40.2=pypi_0
+zipp=3.18.1=pypi_0
+zlib=1.2.13=h5eee18b_1
+zstd=1.5.5=hc292b87_2

train_combined.py

@@ -0,0 +1,273 @@
+from __future__ import print_function
+import argparse
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.optim.lr_scheduler import StepLR, MultiStepLR
+from datasets.hkpoly_test import hktest
+from datasets.original_combined_train import Combined_original
+from datasets.rb_loader import RB_loader
+from loss import DualMSLoss_FineGrained_domain_agnostic_ft, DualMSLoss_FineGrained, DualMSLoss_FineGrained_domain_agnostic
+import timm
+from utils import Prev_RetMetric, RetMetric, compute_recall_at_k, l2_norm, compute_sharded_cosine_similarity, count_parameters
+from pprint import pprint
+import numpy as np
+from tqdm import tqdm
+from combined_sampler import BalancedSampler
+from torch.utils.data.sampler import BatchSampler
+from torch.nn.parallel import DataParallel
+from model import SwinModel_domain_agnostic as Model
+import matplotlib.pyplot as plt
+from sklearn.metrics import roc_curve, auc
+import json
+from torch.utils.tensorboard import SummaryWriter
+
+def train(args, model, device, train_loader, test_loader, optimizers, epoch, loss_func, pl_arg, stepping, log_writer):
+    model.train()
+    steploss = list()
+    for batch_idx, (x_cl, x_cb, target, category_cl, category_cb) in enumerate(pbar := tqdm(train_loader)):
+        x_cl, x_cb, target, category_cl, category_cb = x_cl.to(device), x_cb.to(device), target.to(device), category_cl.to(device), category_cb.to(device)
+        for optimizer in optimizers:
+            optimizer.zero_grad()
+        x_cl, x_cb, x_cl_tokens, x_cb_tokens, domain_class_cl, domain_class_cb = model(x_cl, x_cb)
+        loss = loss_func(x_cl, x_cb, x_cl_tokens, x_cb_tokens, target, device, domain_class_cl, domain_class_cb, category_cl, category_cb)
+        loss.backward()
+        for optimizer in optimizers:
+            optimizer.step()
+        if batch_idx % args.log_interval == 0:
+            if args.dry_run:
+                break
+        pbar.set_description(f"Loss {loss}")
+        steploss.append(loss)
+    return sum(steploss)/len(steploss), stepping
+
+def l2_norm(input):
+    input_size = input.size()
+    buffer = torch.pow(input, 2)
+    normp = torch.sum(buffer, 1).add_(1e-12)
+    norm = torch.sqrt(normp)
+    _output = torch.div(input, norm.view(-1, 1).expand_as(input))
+    output = _output.view(input_size)
+    return output
+
+def hkpoly_test_fn(model,device,test_loader,epoch,plot_argument):
+    model.eval()
+    cl_feats, cb_feats, cl_labels, cb_labels = list(),list(),list(),list()
+    with torch.no_grad():
+        for (x_cl, x_cb, label) in tqdm(test_loader):
+            x_cl, x_cb, label = x_cl.to(device), x_cb.to(device), label.to(device)
+            x_cl_feat, x_cl_token = model.get_embeddings(x_cl,'contactless')
+            x_cb_feat,x_cb_token = model.get_embeddings(x_cb,'contactbased')
+            x_cl_feat = l2_norm(x_cl_feat).cpu().detach().numpy()
+            x_cb_feat = l2_norm(x_cb_feat).cpu().detach().numpy()
+            label = label.cpu().detach().numpy()
+            cl_feats.append(x_cl_feat)
+            cb_feats.append(x_cb_feat)
+            cl_labels.append(label)
+            cb_labels.append(label)
+
+    cl_feats = np.concatenate(cl_feats)
+    cb_feats = np.concatenate(cb_feats)
+    cl_label = torch.from_numpy(np.concatenate(cl_labels))
+    cb_label = torch.from_numpy(np.concatenate(cb_labels))
+
+    # CB2CL
+    scores = np.dot(cl_feats,np.transpose(cb_feats))
+    np.save("combined_models_scores/task1_cb2cl_score_matrix_"+str(epoch)+"_"+plot_argument[0]+"_"+plot_argument[1]+"_"+plot_argument[2]+"_"+plot_argument[3]+".npy", scores)
+    scores = scores.flatten().tolist()
+    labels = torch.eq(cl_label.view(-1,1) - cb_label.view(1,-1),0.0).flatten().tolist()
+    ids_mod = list()
+    for i in labels:
+        if i==True:
+            ids_mod.append(1)
+        else:
+            ids_mod.append(0)
+    fpr,tpr,_ = roc_curve(labels,scores)
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.01)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.01)
+    tar_far_102 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.001)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.001)
+    tar_far_103 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.0001)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.0001)
+    tar_far_104 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    fnr = 1 - tpr
+    EER = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+    roc_auc = auc(fpr, tpr)
+    plt.figure()
+    plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % roc_auc)
+    plt.plot([0, 1], [0, 1], 'k--', label='No Skill')
+    plt.xlim([0, 1])
+    plt.ylim([0, 1])
+    plt.xlabel('False Positive Rate')
+    plt.ylabel('True Positive Rate')
+    plt.title('ROC Curve CB2CL task1')
+    plt.legend(loc="lower right")
+    plt.savefig("combined_models_scores/roc_curve_cb2cl_task1_"+"_"+plot_argument[0]+"_"+plot_argument[1]+"_"+plot_argument[2]+"_"+plot_argument[3]+str(epoch)+".png", dpi=300, bbox_inches='tight')
+    print(f"ROCAUC for CB2CL: {roc_auc * 100} %")
+    print(f"EER for CB2CL: {EER * 100} %")
+    eer_cb2cl = EER * 100
+    print(f"TAR@FAR=10^-2 for CB2CL: {tar_far_102 * 100} %")
+    print(f"TAR@FAR=10^-3 for CB2CL: {tar_far_103 * 100} %")
+    print(f"TAR@FAR=10^-4 for CB2CL: {tar_far_104 * 100} %")
+    cbcltf102 = tar_far_102 * 100
+    cbcltf103 = tar_far_103 * 100
+    cbcltf104 = tar_far_104 * 100
+    cl_label = cl_label.cpu().detach().numpy()
+    cb_label = cb_label.cpu().detach().numpy()
+    recall_score = Prev_RetMetric([cb_feats,cl_feats],[cb_label,cl_label],cl2cl = False)
+    cl2cbk1 = recall_score.recall_k(k=1) * 100
+    print(f"R@1 for CB2CL: {recall_score.recall_k(k=1) * 100} %")
+    print(f"R@10 for CB2CL: {recall_score.recall_k(k=10) * 100} %")
+    print(f"R@50 for CB2CL: {recall_score.recall_k(k=50) * 100} %")
+    print(f"R@100 for CB2CL: {recall_score.recall_k(k=100) * 100} %")
+
+    return cl2cbk1,eer_cb2cl,cbcltf102,cbcltf103,cbcltf104
+
+def main():
+    # Training settings
+    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
+    parser.add_argument('--manifest-list', type=list, default=mani_lst,
+                        help='list of manifest files from different datasets to train on')
+    parser.add_argument('--batch-size', type=int, default=32, metavar='N',
+                        help='input batch size for training (default: 64)')
+    parser.add_argument('--test-batch-size', type=int, default=16, metavar='N',
+                        help='input batch size for testing (default: 1000)')
+    parser.add_argument('--epochs', type=int, default=50, metavar='N',
+                        help='number of epochs to train (default: 14)')
+    parser.add_argument('--lr_linear', type=float, default=1.0, metavar='LR',
+                        help='learning rate (default: 1.0)')
+    parser.add_argument('--lr_swin', type=float, default=1.0, metavar='LR',
+                        help='learning rate (default: 1.0)')
+    parser.add_argument('--gamma', type=float, default=0.9, metavar='M',
+                        help='Learning rate step gamma (default: 0.7)')
+    parser.add_argument('--no-cuda', action='store_true', default=False,
+                        help='disables CUDA training')
+    parser.add_argument('--dry-run', action='store_true', default=False,
+                        help='quickly check a single pass')
+    parser.add_argument('--seed', type=int, default=1, metavar='S',
+                        help='random seed (default: 1)')
+    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
+                        help='how many batches to wait before logging training status')
+    parser.add_argument('--warmup', type=int, default=2, metavar='N',
+                        help='warm up rate for feature extractor')
+    parser.add_argument('--model-name', type=str, default="ridgeformer",
+                        help='Name of the model for checkpointing')
+    args = parser.parse_args()
+
+    checkpoint_save_path = "ridgeformer_checkpoints/"
+    use_cuda = not args.no_cuda and torch.cuda.is_available()
+
+    if not os.path.exists("experiment_logs/"+args.model_name):
+        os.mkdir("experiment_logs/"+args.model_name)
+
+    log_writer = SummaryWriter("experiment_logs/"+args.model_name+"/",comment = str(args.batch_size)+str(args.lr_linear)+str(args.lr_swin))
+
+    torch.manual_seed(args.seed)
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    print("loading Normal RGB images -----------------------------")
+    train_dataset    = Combined_original(args.manifest_list,split="train")
+    val_dataset      = hktest(split="test")
+    
+    balanced_sampler = BalancedSampler(train_dataset, batch_size = args.batch_size, images_per_class = 2)
+    batch_sampler    = BatchSampler(balanced_sampler, batch_size = args.batch_size, drop_last = True)
+    
+    train_kwargs     = {'batch_sampler': batch_sampler}
+    test_kwargs      = {'batch_size':    args.test_batch_size}
+    
+    if use_cuda:
+        cuda_kwargs = {
+                       'num_workers': 1,
+                       'pin_memory': True
+                       }
+        train_kwargs.update(cuda_kwargs)
+        test_kwargs.update(cuda_kwargs)
+
+    
+    train_loader = torch.utils.data.DataLoader(train_dataset, **train_kwargs)
+    test_loader = torch.utils.data.DataLoader(val_dataset, **test_kwargs)
+
+    model = Model().to(device)
+    ckpt = torch.load("ridgeformer_checkpoints/phase1_scratch.pt", map_location=torch.device('cpu'))
+    model.load_state_dict(ckpt,strict=False)
+    print("Number of Trainable Parameters: - ", count_parameters(model))
+
+    loss_func = DualMSLoss_FineGrained_domain_agnostic()
+    # loss_func = DualMSLoss_FineGrained_domain_agnostic_ft()
+    
+    optimizer_swin = optim.AdamW(
+        [
+            {"params": model.swin_cl.parameters(), "lr":args.lr_swin},
+            {"params": model.classify.parameters(), "lr":args.lr_linear},
+            {"params": model.linear_cl.parameters(), "lr":args.lr_linear},
+            {"params": model.linear_cb.parameters(), "lr":args.lr_linear},
+        ],
+        weight_decay=0.000001,
+        lr=args.lr_swin)
+
+    scheduler_swin = MultiStepLR(optimizer_swin, milestones = [100], gamma=0.7)
+
+    cl2cl_lst = list()
+    cb2cl_lst = list()
+    eer_cl2cl_lst = list()
+    eer_cb2cl_lst = list()
+    cbcltf102_lst,cbcltf103_lst,cbcltf104_lst,clcltf102_lst,clcltf103_lst,clcltf104_lst = list(),list(),list(),list(),list(),list()
+    stepping = 1
+    for epoch in range(1, args.epochs + 1):
+        print(f"running epoch------ {epoch}")            
+        if (epoch > args.warmup):
+            print("Training with Swin")
+            model.unfreeze_encoder()
+        else:
+            print("Training only linear")
+            model.freeze_encoder()
+
+        avg_step_loss,stepping = train(args, model, device, train_loader, test_loader, [optimizer_swin], epoch, loss_func, [args.model_name,str(args.batch_size),str(args.lr_linear),str(args.lr_swin)],stepping,log_writer)
+        
+        print(f"Learning Rate for {epoch} for swin = {scheduler_swin.get_last_lr()}")
+        
+        log_writer.add_scalar('Swin_LR/epoch',scheduler_swin.get_last_lr()[0],epoch)
+
+        if (epoch > args.warmup):
+            scheduler_swin.step()
+
+        cl2clk1,cl2cbk1,eer_cb2cl,eer_cl2cl,cbcltf102,cbcltf103,cbcltf104,clcltf102,clcltf103,clcltf104 = hkpoly_test_fn(model, device, test_loader, epoch,[args.model_name,str(args.batch_size),str(args.lr_linear),str(args.lr_swin)])
+        cl2cl_lst.append(cl2clk1)
+        cb2cl_lst.append(cl2cbk1)
+        eer_cl2cl_lst.append(eer_cl2cl)
+        eer_cb2cl_lst.append(eer_cb2cl)
+        cbcltf102_lst.append(cbcltf102)
+        cbcltf103_lst.append(cbcltf103)
+        cbcltf104_lst.append(cbcltf104)
+        clcltf102_lst.append(clcltf102)
+        clcltf103_lst.append(clcltf103)
+        clcltf104_lst.append(clcltf104)
+
+        log_writer.add_scalars('recall@1/epoch',{'CL2CL':cl2clk1,'CB2CL':cl2cbk1},epoch)
+        log_writer.add_scalars('EER/epoch',{'CL2CL':eer_cl2cl,'CB2CL':eer_cb2cl},epoch)
+        log_writer.add_scalars('TARFAR10^-2/epoch',{'CL2CL':clcltf102,'CB2CL':cbcltf102},epoch)
+        log_writer.add_scalars('TARFAR10^-3/epoch',{'CL2CL':clcltf103,'CB2CL':cbcltf103},epoch)
+        log_writer.add_scalars('TARFAR10^-4/epoch',{'CL2CL':clcltf104,'CB2CL':cbcltf104},epoch)
+        log_writer.add_scalar('AvgLoss/epoch',avg_step_loss,epoch)
+
+        torch.save(model.state_dict(), checkpoint_save_path + "combinedtrained_hkpolytest_" + args.model_name + "_" + str(args.lr_linear) + "_" + str(args.lr_swin) + "_" + str(args.batch_size) + str(epoch) + "_" + str(cl2clk1)+ "_" + str(cl2cbk1) + ".pt")
+    log_writer.close()
+
+    print(f"Maximum recall@1 for CL2CL: {max(cl2cl_lst)} at epoch {cl2cl_lst.index(max(cl2cl_lst))+1}")
+    print(f"Maximum recall@1 for CB2CL: {max(cb2cl_lst)} at epoch {cb2cl_lst.index(max(cb2cl_lst))+1}")
+    print(f"Minimum EER for CL2CL: {min(eer_cl2cl_lst)} at epoch {eer_cl2cl_lst.index(min(eer_cl2cl_lst))+1}")
+    print(f"Minimum EER for CB2CL: {min(eer_cb2cl_lst)} at epoch {eer_cb2cl_lst.index(min(eer_cb2cl_lst))+1}")
+    print(f"Maximum TAR@FAR=10^-2 for CB2CL: {max(cbcltf102_lst)} at epoch {cbcltf102_lst.index(max(cbcltf102_lst))+1}")
+    print(f"Maximum TAR@FAR=10^-3 for CB2CL: {max(cbcltf103_lst)} at epoch {cbcltf103_lst.index(max(cbcltf103_lst))+1}")
+    print(f"Maximum TAR@FAR=10^-4 for CB2CL: {max(cbcltf104_lst)} at epoch {cbcltf104_lst.index(max(cbcltf104_lst))+1}")
+    print(f"Maximum TAR@FAR=10^-2 for CL2CL: {max(clcltf102_lst)} at epoch {clcltf102_lst.index(max(clcltf102_lst))+1}")
+    print(f"Maximum TAR@FAR=10^-3 for CL2CL: {max(clcltf103_lst)} at epoch {clcltf103_lst.index(max(clcltf103_lst))+1}")
+    print(f"Maximum TAR@FAR=10^-4 for CL2CL: {max(clcltf104_lst)} at epoch {clcltf104_lst.index(max(clcltf104_lst))+1}")
+
+if __name__ == '__main__':
+    main()

train_combined_fusion.py

@@ -0,0 +1,301 @@
+from __future__ import print_function
+import argparse
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.optim.lr_scheduler import StepLR, MultiStepLR
+from datasets.hkpoly_test import hktest
+from datasets.original_combined_train import Combined_original
+from datasets.rb_loader import RB_loader
+from loss import DualMSLoss_FineGrained_domain_agnostic_ft, DualMSLoss_FineGrained, DualMSLoss_FineGrained_domain_agnostic
+import timm
+from utils import Prev_RetMetric, RetMetric, compute_recall_at_k, l2_norm, compute_sharded_cosine_similarity, count_parameters
+from pprint import pprint
+import numpy as np
+from tqdm import tqdm
+from combined_sampler import BalancedSampler
+from torch.utils.data.sampler import BatchSampler
+from torch.nn.parallel import DataParallel
+from model import SwinModel_Fusion as Model
+import matplotlib.pyplot as plt
+from sklearn.metrics import roc_curve, auc
+import json
+from torch.utils.tensorboard import SummaryWriter
+
+def train(args, model, device, train_loader, test_loader, optimizers, epoch, loss_func, pl_arg, stepping, log_writer, checkpoint_save_path):
+    model.train()
+    steploss = list()
+    for batch_idx, (x_cl, x_cb, target,_,_) in enumerate(pbar := tqdm(train_loader)):
+        x_cl, x_cb, target = x_cl.to(device), x_cb.to(device), target.to(device)
+        for optimizer in optimizers:
+            optimizer.zero_grad()
+        x_cl_tokens, x_cb_tokens = model(x_cl, x_cb)
+
+        N, M, D = x_cl_tokens.shape
+
+        index_i = torch.arange(N).unsqueeze(1)  # Shape: (100, 1)
+        index_j = torch.arange(N).unsqueeze(0)  # Shape: (1, 100)
+
+        x = x_cl_tokens[index_i]  # Shape: (100, 100, 197, 1024)
+        y = x_cb_tokens[index_j]  # Shape: (100, 100, 197, 1024)
+
+        x = x.expand(N, N, M, D).reshape(N * N, M, D)  # Shape: (10000, 197, 1024)
+        y = y.expand(N, N, M, D).reshape(N * N, M, D)  # Shape: (10000, 197, 1024)
+        sim_matrix,_ = model.combine_features(x, y)
+        sim_matrix = sim_matrix.view(N, N).to(device)
+
+        loss = loss_func.ms_sample(sim_matrix, target).cuda() + loss_func.ms_sample(sim_matrix.t(), target.t()).cuda()
+        loss.backward()
+        for optimizer in optimizers:
+            optimizer.step()
+        if batch_idx % args.log_interval == 0:
+            if args.dry_run:
+                break
+        pbar.set_description(f"Loss {loss}")
+        steploss.append(loss)
+        if (batch_idx + 1)%50 == 0:
+            cl2clk1,cl2cbk1,eer_cb2cl,eer_cl2cl,cbcltf102,cbcltf103,cbcltf104,clcltf102,clcltf103,clcltf104 = hkpoly_test_fn(model, device, test_loader, epoch, pl_arg)
+            log_writer.add_scalars('recall@1/step',{'CL2CL':cl2clk1,'CB2CL':cl2cbk1},stepping)
+            log_writer.add_scalars('EER/step',{'CL2CL':eer_cl2cl,'CB2CL':eer_cb2cl},stepping)
+            log_writer.add_scalars('TARFAR10^-2/step',{'CL2CL':clcltf102,'CB2CL':cbcltf102},stepping)
+            log_writer.add_scalars('TARFAR10^-4/step',{'CL2CL':clcltf104,'CB2CL':cbcltf104},stepping)
+            stepping+=1
+            
+    return sum(steploss)/len(steploss), stepping
+
+def l2_norm(input):
+    input_size = input.size()
+    buffer = torch.pow(input, 2)
+    normp = torch.sum(buffer, 1).add_(1e-12)
+    norm = torch.sqrt(normp)
+    _output = torch.div(input, norm.view(-1, 1).expand_as(input))
+    output = _output.view(input_size)
+    return output
+
+def get_fused_cross_score_matrix(model, cl_tokens, cb_tokens):
+    cl_tokens   = torch.cat(cl_tokens)
+    cb_tokens   = torch.cat(cb_tokens)
+    batch_size  = cl_tokens.shape[0]
+    shard_size  = 20
+    similarity_matrix = torch.zeros((batch_size, batch_size))
+    for i_start in tqdm(range(0, batch_size, shard_size)):
+        i_end   = min(i_start + shard_size, batch_size)
+        shard_i = cl_tokens[i_start:i_end]  
+        for j_start in range(0, batch_size, shard_size):
+            j_end               = min(j_start + shard_size, batch_size)
+            shard_j             = cb_tokens[j_start:j_end]
+            batch_i             = shard_i.unsqueeze(1)
+            batch_j             = shard_j.unsqueeze(0)
+            pairwise_i          = batch_i.expand(-1, shard_size, -1, -1)
+            pairwise_j          = batch_j.expand(shard_size, -1, -1, -1)
+
+            similarity_scores, distances   = model.combine_features(pairwise_i.reshape(-1, 197, 1024), pairwise_j.reshape(-1, 197, 1024))
+            scores = similarity_scores - 0.1 * distances
+            scores = scores.reshape(shard_size, shard_size)
+            similarity_matrix[i_start:i_end, j_start:j_end] = scores.cpu().detach()
+    return similarity_matrix
+
+def hkpoly_test_fn(model,device,test_loader,epoch,plot_argument):
+    model.eval()
+    cl_feats, cb_feats, cl_labels, cb_labels = list(),list(),list(),list()
+    with torch.no_grad():
+        for (x_cl, x_cb, label) in tqdm(test_loader):
+            x_cl, x_cb, label = x_cl.to(device), x_cb.to(device), label.to(device)
+            x_cl_token  = model.get_tokens(x_cl,'contactless')
+            x_cb_token  = model.get_tokens(x_cb,'contactbased')
+            label = label.cpu().detach().numpy()
+            cl_feats.append(x_cl_token)
+            cb_feats.append(x_cb_token)
+            cl_labels.append(label)
+            cb_labels.append(label)
+
+    cl_label = torch.from_numpy(np.concatenate(cl_labels))
+    cb_label = torch.from_numpy(np.concatenate(cb_labels))
+
+    # CB2CL
+    scores_mat = get_fused_cross_score_matrix(model, cl_feats, cb_feats)
+    np.save("combined_models_scores/task1_cb2cl_score_matrix_"+str(epoch)+"_"+plot_argument[0]+"_"+plot_argument[1]+"_"+plot_argument[2]+".npy", scores_mat)
+    scores = scores_mat.cpu().detach().numpy().flatten().tolist()
+    labels = torch.eq(cb_label.view(-1,1) - cl_label.view(1,-1),0.0).flatten().tolist()
+    ids_mod = list()
+    for i in labels:
+        if i==True:
+            ids_mod.append(1)
+        else:
+            ids_mod.append(0)
+    fpr,tpr,_ = roc_curve(labels,scores)
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.01)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.01)
+    tar_far_102 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.001)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.001)
+    tar_far_103 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    lower_fpr_idx = max(i for i, val in enumerate(fpr) if val < 0.0001)
+    upper_fpr_idx = min(i for i, val in enumerate(fpr) if val >= 0.0001)
+    tar_far_104 = (tpr[lower_fpr_idx]+tpr[upper_fpr_idx])/2
+    fnr = 1 - tpr
+    EER = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+    roc_auc = auc(fpr, tpr)
+    plt.figure()
+    plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % roc_auc)
+    plt.plot([0, 1], [0, 1], 'k--', label='No Skill')
+    plt.xlim([0, 1])
+    plt.ylim([0, 1])
+    plt.xlabel('False Positive Rate')
+    plt.ylabel('True Positive Rate')
+    plt.title('ROC Curve CB2CL task1')
+    plt.legend(loc="lower right")
+    plt.savefig("combined_models_scores/roc_curve_cb2cl_task1_"+"_"+plot_argument[0]+"_"+plot_argument[1]+"_"+plot_argument[2]+str(epoch)+".png", dpi=300, bbox_inches='tight')
+    print(f"ROCAUC for CB2CL: {roc_auc * 100} %")   
+    print(f"EER for CB2CL: {EER * 100} %")
+    eer_cb2cl = EER * 100
+    print(f"TAR@FAR=10^-2 for CB2CL: {tar_far_102 * 100} %")
+    print(f"TAR@FAR=10^-3 for CB2CL: {tar_far_103 * 100} %")
+    print(f"TAR@FAR=10^-4 for CB2CL: {tar_far_104 * 100} %")
+    cbcltf102 = tar_far_102 * 100
+    cbcltf103 = tar_far_103 * 100
+    cbcltf104 = tar_far_104 * 100
+    cl2cbk1   = compute_recall_at_k(scores_mat, cl_label, cb_label, 1) * 100
+    print(f"R@1 for CB2CL: {compute_recall_at_k(scores_mat, cl_label, cb_label, 1) * 100} %")
+    print(f"R@10 for CB2CL: {compute_recall_at_k(scores_mat, cl_label, cb_label, 10) * 100} %")
+    print(f"R@50 for CB2CL: {compute_recall_at_k(scores_mat, cl_label, cb_label, 50) * 100} %")
+    print(f"R@100 for CB2CL: {compute_recall_at_k(scores_mat, cl_label, cb_label, 100) * 100} %")
+    torch.cuda.empty_cache()
+
+    return cl2cbk1,eer_cb2cl,cbcltf102,cbcltf103,cbcltf104
+
+def main():
+    # Training settings
+    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
+    parser.add_argument('--manifest-list', type=list, default=mani_lst,
+                        help='list of manifest files')
+    parser.add_argument('--batch-size', type=int, default=32, metavar='N',
+                        help='input batch size for training (default: 64)')
+    parser.add_argument('--test-batch-size', type=int, default=16, metavar='N',
+                        help='input batch size for testing (default: 1000)')
+    parser.add_argument('--epochs', type=int, default=50, metavar='N',
+                        help='number of epochs to train (default: 14)')
+    parser.add_argument('--lr_fusion', type=float, default=1.0, metavar='LR',
+                        help='learning rate (default: 1.0)')
+    parser.add_argument('--gamma', type=float, default=0.9, metavar='M',
+                        help='Learning rate step gamma (default: 0.7)')
+    parser.add_argument('--no-cuda', action='store_true', default=False,
+                        help='disables CUDA training')
+    parser.add_argument('--dry-run', action='store_true', default=False,
+                        help='quickly check a single pass')
+    parser.add_argument('--seed', type=int, default=1, metavar='S',
+                        help='random seed (default: 1)')
+    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
+                        help='how many batches to wait before logging training status')
+    parser.add_argument('--warmup', type=int, default=2, metavar='N',
+                        help='warm up rate for feature extractor')
+    parser.add_argument('--model-name', type=str, default="swinmodel",
+                        help='Name of the model for checkpointing')
+    args = parser.parse_args()
+
+    device = torch.device("cuda")
+    model = Model().to(device)
+    ckpt_combined_phase1_ft = "ridgeformer_checkpoints/combined_models_check/phase1_ft_hkpoly.pt"
+    ckpt_combined_phase2    = "ridgeformer_checkpoints/phase2_scratch.pt"
+
+    model.load_pretrained_models(ckpt_combined_phase1_ft, ckpt_combined_phase2)
+    model.freeze_backbone()
+    checkpoint_save_path = "ridgeformer_checkpoints/"
+    use_cuda = not args.no_cuda and torch.cuda.is_available()
+
+    if not os.path.exists("experiment_logs/"+args.model_name):
+        os.mkdir("experiment_logs/"+args.model_name)
+
+    log_writer = SummaryWriter("experiment_logs/"+args.model_name+"/",comment = str(args.batch_size)+str(args.lr_fusion))
+
+    torch.manual_seed(args.seed)
+
+    print("loading Normal RGB images -----------------------------")
+    train_dataset    = Combined_original(args.manifest_list,split="train")
+    val_dataset      = hktest(split="test")
+    
+    balanced_sampler = BalancedSampler(train_dataset, batch_size = args.batch_size, images_per_class = 2)
+    batch_sampler    = BatchSampler(balanced_sampler, batch_size = args.batch_size, drop_last = True)
+    
+    train_kwargs     = {'batch_sampler': batch_sampler}
+    test_kwargs      = {'batch_size':    args.test_batch_size}
+    
+    if use_cuda:
+        cuda_kwargs = {
+                       'num_workers': 1,
+                       'pin_memory': True
+                       }
+        train_kwargs.update(cuda_kwargs)
+        test_kwargs.update(cuda_kwargs)
+
+    
+    train_loader = torch.utils.data.DataLoader(train_dataset, **train_kwargs)
+    test_loader = torch.utils.data.DataLoader(val_dataset, **test_kwargs)
+    
+    print("Number of Trainable Parameters: - ", count_parameters(model))
+
+    loss_func           = DualMSLoss_FineGrained()
+    optimizer_fusion    = optim.AdamW(
+        [
+            {"params": model.output_logit_mlp.parameters(), "lr":args.lr_fusion},
+            {"params": model.fusion.parameters(),           "lr":args.lr_fusion},
+            {"params": model.sep_token,                     "lr":args.lr_fusion},
+            {"params": model.encoder_layer.parameters(),    "lr":args.lr_fusion},
+
+         ],
+        weight_decay=0.000001,
+        lr=args.lr_fusion)
+
+    scheduler = MultiStepLR(optimizer_fusion, milestones = [3,6,9,14], gamma=0.5)
+
+    cl2cl_lst,cb2cl_lst,eer_cl2cl_lst,eer_cb2cl_lst,cbcltf102_lst,cbcltf103_lst,cbcltf104_lst,clcltf102_lst,clcltf103_lst,clcltf104_lst = list(),list(),list(),list(),list(),list(),list(),list(),list(),list()
+    stepping = 1
+    for epoch in range(1, args.epochs + 1):
+        print(f"running epoch------ {epoch}")            
+        avg_step_loss,stepping = train(args, model, device, train_loader, test_loader, [optimizer_fusion], epoch, loss_func, [args.model_name,str(args.batch_size),str(args.lr_fusion)],stepping,log_writer, checkpoint_save_path)
+        
+        print(f"Learning Rate for {epoch} for linear = {scheduler.get_last_lr()}")
+        print(f"Learning Rate for {epoch} for swin = {scheduler.get_last_lr()}")
+        
+        log_writer.add_scalar('Liner_LR/epoch',scheduler.get_last_lr()[0],epoch)
+        log_writer.add_scalar('Swin_LR/epoch',scheduler.get_last_lr()[0],epoch)
+
+        scheduler.step()
+        
+        cl2clk1,cl2cbk1,eer_cb2cl,eer_cl2cl,cbcltf102,cbcltf103,cbcltf104,clcltf102,clcltf103,clcltf104 = hkpoly_test_fn(model, device, test_loader, epoch, [args.model_name,str(args.batch_size),str(args.lr_fusion)])
+        cl2cl_lst.append(cl2clk1)
+        cb2cl_lst.append(cl2cbk1)
+        eer_cl2cl_lst.append(eer_cl2cl)
+        eer_cb2cl_lst.append(eer_cb2cl)
+        cbcltf102_lst.append(cbcltf102)
+        cbcltf103_lst.append(cbcltf103)
+        cbcltf104_lst.append(cbcltf104)
+        clcltf102_lst.append(clcltf102)
+        clcltf103_lst.append(clcltf103)
+        clcltf104_lst.append(clcltf104)
+
+        log_writer.add_scalars('recall@1/epoch',{'CL2CL':cl2clk1,'CB2CL':cl2cbk1},epoch)
+        log_writer.add_scalars('EER/epoch',{'CL2CL':eer_cl2cl,'CB2CL':eer_cb2cl},epoch)
+        log_writer.add_scalars('TARFAR10^-2/epoch',{'CL2CL':clcltf102,'CB2CL':cbcltf102},epoch)
+        log_writer.add_scalars('TARFAR10^-4/epoch',{'CL2CL':clcltf104,'CB2CL':cbcltf104},epoch)
+        log_writer.add_scalar('AvgLoss/epoch',avg_step_loss,epoch)
+
+        torch.save(model.state_dict(), checkpoint_save_path + "combinedtrained_hkpolytest_" + args.model_name + "_" + str(args.lr_fusion) + "_" + str(args.batch_size) + str(epoch) + "_" + str(cl2clk1)+ "_" + str(cl2cbk1) + ".pt")
+    log_writer.close()
+
+    print(f"Maximum recall@1 for CL2CL: {max(cl2cl_lst)} at epoch {cl2cl_lst.index(max(cl2cl_lst))+1}")
+    print(f"Maximum recall@1 for CB2CL: {max(cb2cl_lst)} at epoch {cb2cl_lst.index(max(cb2cl_lst))+1}")
+    print(f"Minimum EER for CL2CL: {min(eer_cl2cl_lst)} at epoch {eer_cl2cl_lst.index(min(eer_cl2cl_lst))+1}")
+    print(f"Minimum EER for CB2CL: {min(eer_cb2cl_lst)} at epoch {eer_cb2cl_lst.index(min(eer_cb2cl_lst))+1}")
+    print(f"Maximum TAR@FAR=10^-2 for CB2CL: {max(cbcltf102_lst)} at epoch {cbcltf102_lst.index(max(cbcltf102_lst))+1}")
+    print(f"Maximum TAR@FAR=10^-3 for CB2CL: {max(cbcltf103_lst)} at epoch {cbcltf103_lst.index(max(cbcltf103_lst))+1}")
+    print(f"Maximum TAR@FAR=10^-4 for CB2CL: {max(cbcltf104_lst)} at epoch {cbcltf104_lst.index(max(cbcltf104_lst))+1}")
+    print(f"Maximum TAR@FAR=10^-2 for CL2CL: {max(clcltf102_lst)} at epoch {clcltf102_lst.index(max(clcltf102_lst))+1}")
+    print(f"Maximum TAR@FAR=10^-3 for CL2CL: {max(clcltf103_lst)} at epoch {clcltf103_lst.index(max(clcltf103_lst))+1}")
+    print(f"Maximum TAR@FAR=10^-4 for CL2CL: {max(clcltf104_lst)} at epoch {clcltf104_lst.index(max(clcltf104_lst))+1}")
+
+if __name__ == '__main__':
+    main()

utils.py

@@ -0,0 +1,117 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+class RetMetric(object):
+    def __init__(self, sim_mat, labels):
+        self.gallery_labels, self.query_labels = labels
+        self.sim_mat = sim_mat
+        self.is_equal_query = False
+
+    def recall_k(self, k=1):
+        m = len(self.sim_mat)
+
+        match_counter = 0
+
+        for i in range(m):
+            pos_sim = self.sim_mat[i][self.gallery_labels == self.query_labels[i]]
+            neg_sim = self.sim_mat[i][self.gallery_labels != self.query_labels[i]]
+            
+            thresh = np.sort(pos_sim)[-2] if self.is_equal_query and len(pos_sim) > 1 else np.max(pos_sim)
+
+            if np.sum(neg_sim > thresh) < k:
+                match_counter += 1
+        return float(match_counter) / m
+
+class Prev_RetMetric(object):
+    def __init__(self, feats, labels, cl2cl=True):
+
+        if len(feats) == 2 and type(feats) == list:
+            """
+            feats = [gallery_feats, query_feats]
+            labels = [gallery_labels, query_labels]
+            """
+            self.is_equal_query = False
+
+            self.gallery_feats, self.query_feats = feats
+            self.gallery_labels, self.query_labels = labels
+
+        else:
+            self.is_equal_query = True
+            self.gallery_feats = self.query_feats = feats
+            self.gallery_labels = self.query_labels = labels
+
+        self.sim_mat = np.matmul(self.query_feats, np.transpose(self.gallery_feats))
+        if cl2cl:
+            self.sim_mat = self.sim_mat * (1 - np.identity(self.sim_mat.shape[0]))
+        
+    def recall_k(self, k=1):
+        m = len(self.sim_mat)
+
+        match_counter = 0
+
+        for i in range(m):
+            pos_sim = self.sim_mat[i][self.gallery_labels == self.query_labels[i]]
+            neg_sim = self.sim_mat[i][self.gallery_labels != self.query_labels[i]]
+
+            thresh = np.sort(pos_sim)[-2] if self.is_equal_query else np.max(pos_sim)
+
+            if np.sum(neg_sim > thresh) < k:
+                match_counter += 1
+        return float(match_counter) / m
+    
+def compute_recall_at_k(similarity_matrix, p_labels, g_labels, k):
+    num_probes = p_labels.size(0)
+    recall_at_k = 0.0
+    for i in range(num_probes):
+        probe_label = p_labels[i]
+        sim_scores = similarity_matrix[i]
+        sorted_indices = torch.argsort(sim_scores, descending=True)
+        top_k_indices = sorted_indices[:k]
+        correct_in_top_k = any(g_labels[idx] == probe_label for idx in top_k_indices)
+        recall_at_k += correct_in_top_k
+    recall_at_k /= num_probes
+    return recall_at_k
+
+def count_parameters(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+def l2_norm(input):
+    input_size = input.size()
+    buffer = torch.pow(input, 2)
+    normp = torch.sum(buffer, 1).add_(1e-12)
+    norm = torch.sqrt(normp)
+    _output = torch.div(input, norm.view(-1, 1).expand_as(input))
+    output = _output.view(input_size)
+    
+    return output
+
+def compute_sharded_cosine_similarity(tensor1, tensor2, shard_size):
+    B, T, D = tensor1.shape
+    average_sim_matrix = torch.zeros((B, B), device=tensor1.device)
+
+    for start_idx1 in tqdm(range(0, B, shard_size)):
+        end_idx1 = min(start_idx1 + shard_size, B)
+
+        for start_idx2 in range(0, B, shard_size):
+            end_idx2 = min(start_idx2 + shard_size, B)
+
+            # Get the shard
+            shard_tensor1 = tensor1[start_idx1:end_idx1]
+            shard_tensor2 = tensor2[start_idx2:end_idx2]
+
+            # Reshape and expand
+            shard_tensor1_expanded = shard_tensor1.unsqueeze(1).unsqueeze(3)
+            shard_tensor2_expanded = shard_tensor2.unsqueeze(0).unsqueeze(2)
+
+            # Compute cosine similarity for the shard
+            shard_cos_sim = F.cosine_similarity(shard_tensor1_expanded, shard_tensor2_expanded, dim=-1)
+
+            # Sum up the cosine similarities
+            average_sim_matrix[start_idx1:end_idx1, start_idx2:end_idx2] += torch.sum(shard_cos_sim, dim=[2, 3])
+
+    # Normalize by the total number of elements (T*T)
+    average_sim_matrix /= (T * T)
+
+    return average_sim_matrix