Upload folder using huggingface_hub

ASDA/engine/engine.py

@@ -0,0 +1,167 @@
+import time
+import matplotlib as mpl
+mpl.use('Agg')
+
+import numpy as np
+import torch
+import torch.nn.parallel
+import torch.optim
+from torch.autograd import Variable
+from torch.cuda.amp import autocast as autocast
+
+from model.model import *
+from dataset.data_loader import *
+from utils.losses import *
+from utils.parsing_metrics import *
+from utils.utils import *
+from utils.utils import dice_loss, sigmoid_focal_loss
+
+use_cuda = torch.cuda.is_available()
+print("use_cuda, ", use_cuda)
+
+
+def train_epoch(rank, args, train_loader, model, optimizer, epoch, scaler, logger):
+    print('train at epoch %d'%epoch)
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+    dice_losses = AverageMeter()
+    sigmoid_focal_losses = AverageMeter()
+    cos_losses = AverageMeter()
+    model.train()
+    end = time.time()
+
+    for batch_idx, (imgs, word_id, word_mask, bbox, seg_map) in enumerate(train_loader):
+        imgs = imgs.cuda(rank, non_blocking=True)
+        word_id = word_id.cuda(rank, non_blocking=True)
+        word_mask = word_mask.cuda(rank, non_blocking=True)
+        seg_map = seg_map.cuda(rank, non_blocking=True)
+        image = Variable(imgs)
+        word_id = Variable(word_id)
+        word_mask = Variable(word_mask)
+        seg_map = Variable(seg_map)
+
+        with autocast():
+            mask_out = model(image, word_id, word_mask)
+            loss = 0.
+            
+            mask_out_np = mask_out.data.cpu().numpy() # [bs, 1, 208, 208]
+            seg_map_np = seg_map.cpu().numpy() # [bs, 1, 208, 208]
+            seg_iou = cal_seg_iou_loss(seg_map_np, mask_out_np, args.seg_thresh)
+         
+            dice_loss_ = dice_loss(mask_out, seg_map)
+            sigmoid_focal_loss_ = sigmoid_focal_loss(mask_out, seg_map)
+
+            loss += dice_loss_ + sigmoid_focal_loss_
+
+        optimizer.zero_grad()
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+
+        losses.update(loss.item(), imgs.size(0))
+        dice_losses.update(dice_loss_.item(), imgs.size(0))
+        sigmoid_focal_losses.update(sigmoid_focal_loss_.item(), imgs.size(0))
+        cos_losses.update(seg_iou.mean().item(), imgs.size(0))
+        
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if rank == 0 and batch_idx % args.print_freq == 0:
+            print_str = 'Epoch: [{0}][{1}/{2}]\t' \
+                'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
+                'Loss {loss.val:.4f} ({loss.avg:.4f})\t' \
+                'dice_losses {dice_losses.val:.4f} ({dice_losses.avg:.4f})\t' \
+                'sigmoid_focal_losses {sigmoid_focal_losses.val:.4f} ({sigmoid_focal_losses.avg:.4f})\t' \
+                'IoU {cos_loss.val:.4f} ({cos_loss.avg:.4f})\t' \
+                .format(epoch, batch_idx, len(train_loader), batch_time=batch_time, loss=losses, dice_losses=dice_losses, sigmoid_focal_losses=sigmoid_focal_losses, cos_loss=cos_losses)
+            print(print_str)
+            logger.info(print_str)
+
+    return losses.avg
+
+def validate_epoch(args, val_loader, model, logger, mode='val'):
+    print('begin test')
+    batch_time = AverageMeter()
+    miou = AverageMeter()
+    miou_seg = AverageMeter()
+
+    prec=dict()
+    thresholds = np.arange(0.5, 1, 0.05)
+
+    for thresh in thresholds:
+        prec[thresh]= AverageMeter()
+
+    model.eval()
+    end = time.time()
+    idx = 0
+
+    t_all = []
+
+    for batch_idx, (imgs, word_id, word_mask, bbox, seg_map, ratio, dw, dh, im_id, phrase, draw_img) in enumerate(val_loader):
+        
+        imgs = imgs.cuda(0)
+        word_id = word_id.cuda(0)
+        word_mask = word_mask.cuda(0)
+        seg_map = seg_map.cuda(0)
+        image = Variable(imgs)
+        word_id = Variable(word_id)
+        word_mask = Variable(word_mask)
+        seg_map = Variable(seg_map)
+
+        t1 = time.time()
+        with torch.no_grad():
+            mask_out = model(image, word_id, word_mask)
+            mask_out = mask_out.sigmoid()
+
+        t2 = time.time()
+        t_all.append(t2-t1)
+
+        ## test: convert pred, gt box to original scale with meta-info
+        ih = seg_map.shape[-2]
+        iw = seg_map.shape[-1]
+        nh = int(ih * ratio)
+        nw = int(iw * ratio)
+        top, bottom = int(dh[0]), nh + int(dh[0])
+        left, right = int(dw[0]), nw + int(dw[0])
+        ratio = float(ratio)
+        new_shape = (iw, ih)
+        
+        ## revert image for visualization
+        seg_map_np = seg_map[0,:,:,:].data.cpu().numpy().transpose(1,2,0)
+        seg_map_np = cv2.resize(seg_map_np, new_shape, interpolation=cv2.INTER_CUBIC)
+        img_np = imgs[0,:,top:bottom,left:right].data.cpu().numpy().transpose(1,2,0)
+        img_np = cv2.resize(img_np, new_shape, interpolation=cv2.INTER_CUBIC)
+
+        img_np = Variable(torch.from_numpy(img_np.transpose(2,0,1)).cuda().unsqueeze(0))
+        
+        # seg
+        mask_out = mask_out[0].data.cpu().numpy().transpose(1,2,0)
+        mask_out = cv2.resize(mask_out, (args.size, args.size))
+        mask_out_np = mask_out[top:bottom, left:right]
+        mask_out_np = cv2.resize(mask_out_np, new_shape)
+        seg_iou, seg_prec = cal_seg_iou(seg_map[0].cpu().numpy(), mask_out_np, args.seg_thresh)
+        miou_seg.update(seg_iou, imgs.size(0))
+        for thresh in thresholds:
+            prec[thresh].update(seg_prec[thresh], imgs.size(0))
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+        if batch_idx % 1000 == 0:
+            print_str = '[{0}/{1}]\t' \
+                'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
+                'seg_iu {seg.val:.4f} ({seg.avg:.4f})\t' \
+                .format( \
+                    batch_idx, len(val_loader), batch_time=batch_time, seg=miou_seg)
+            print(print_str)
+            logger.info(print_str)  
+        idx = idx + 1
+        
+    print(miou_seg.avg)
+    for thresh in thresholds:
+            print("prec@%f: %f"%(thresh,float(prec[thresh].avg)))
+            logger.info("prec@%f:%f"%(thresh,float(prec[thresh].avg)))
+    logger.info("%f,%f"%(float(miou.avg), miou_seg.avg))
+    return miou_seg.avg, prec
+

ASDA/engine/engine_gref_sbert.py

@@ -0,0 +1,347 @@
+import time
+import matplotlib as mpl
+mpl.use('Agg')
+
+import numpy as np
+import torch
+import torch.nn.parallel
+import torch.optim
+from torch.autograd import Variable
+from torch.cuda.amp import autocast as autocast
+
+from model.model_sbert_gref import *
+from dataset.data_loader import *
+from utils.losses import *
+from utils.parsing_metrics import *
+from utils.utils import *
+from utils.utils import dice_loss, sigmoid_focal_loss
+
+use_cuda = torch.cuda.is_available()
+print("use_cuda, ", use_cuda)
+
+
+def return_mask(emb_distance, verb_mask=None, rows_to_filter=None, cols_to_filter=None):
+    B_, B_ = emb_distance.shape
+    positive_mask = torch.zeros_like(emb_distance)
+    positive_mask.fill_diagonal_(1)  # Set diagonal elements to 1 for all cases
+    
+    if B_ < len(verb_mask):
+        # If B_ equals to 2*K (double the number of verb phrase)
+        for i in range(B_ // 2):
+            positive_mask[2 * i, 2 * i + 1] = 1
+            positive_mask[2 * i + 1, 2 * i] = 1
+    else:
+        # Process the case where we have a mix of sentences with and without verbs
+        i = 0
+        while i < B_:
+            if verb_mask[i] == 1:
+                positive_mask[i, i + 1] = 1
+                positive_mask[i + 1, i] = 1
+                i += 2
+            else:
+                i += 1  
+    negative_mask = torch.ones_like(emb_distance) - positive_mask
+    negative_mask = negative_mask.clone() 
+    
+    if rows_to_filter is not None and cols_to_filter is not None :
+        for row, col in zip(rows_to_filter, cols_to_filter):
+            negative_mask[row * 2, col * 2] = 0
+            negative_mask[row * 2, col * 2 + 1] = 0
+            negative_mask[row * 2 + 1, col * 2] = 0
+            negative_mask[row * 2 + 1, col * 2 + 1] = 0
+
+    return positive_mask, negative_mask
+
+
+def UniAngularLogitContrastLoss(total_fq, verb_mask, rows_to_filter, cols_to_filter, alpha=0.5, verbonly=True, m=0.5, tau=0.05, args=None):        
+    _, C, H, W = total_fq.shape
+
+    # Calculate embeddings
+    if verbonly :
+        B = total_fq[verb_mask].shape[0]
+        emb = torch.mean(total_fq[verb_mask], dim=(-1, -2)).reshape(B, C)
+        assert emb.shape[0] % 2 == 0, f"Embedding count {emb.shape[0]} is not divisible by 2."
+    else :
+        emb = torch.mean(total_fq, dim=-1)
+
+    B_ = emb.shape[0]
+    emb_i = emb.unsqueeze(1).repeat(1, B_, 1)  # (B_, B_, C)
+    emb_j = emb.unsqueeze(0).repeat(B_, 1, 1)  # (B_, B_, C)
+
+    sim = nn.CosineSimilarity(dim=-1, eps=1e-6)
+    sim_matrix = sim(emb_i, emb_j).reshape(B_, B_)  # (B_, B_)
+    sim_matrix = torch.clamp(sim_matrix, min=-0.9999, max=0.9999)
+
+    margin_in_radians = m / 57.2958  # Convert degrees to radians
+    theta_matrix = (torch.pi / 2) - torch.acos(sim_matrix)
+    # print("sim_matrix : ", sim_matrix)
+    # print("theta_matrix : ", theta_matrix)
+        
+    positive_mask, negative_mask = return_mask(sim_matrix, verb_mask, rows_to_filter, cols_to_filter)
+    # print("positive_mask : ", positive_mask)
+    # print("negative_mask : ", negative_mask)
+    # print("positive_mask requires_grad:", positive_mask.requires_grad, 
+    #     "device:", positive_mask.device, "dtype:", positive_mask.dtype)
+    # print("negative_mask requires_grad:", negative_mask.requires_grad, 
+    #     "device:", negative_mask.device, "dtype:", negative_mask.dtype)
+
+
+    theta_with_margin = theta_matrix.clone()
+    theta_with_margin[positive_mask.bool()] -= margin_in_radians  
+    logits = theta_with_margin / tau  # Scale with temperature
+    
+    # Compute exp logits for softmax
+    exp_logits = torch.exp(logits) 
+    pos_exp_logits = exp_logits * positive_mask
+    pos_exp_logits = pos_exp_logits.sum(dim=-1)
+    neg_exp_logits = exp_logits * negative_mask
+    neg_exp_logits = neg_exp_logits.sum(dim=-1)
+    
+    total_exp_logits = pos_exp_logits + neg_exp_logits
+
+    positive_loss = -torch.log(pos_exp_logits/ total_exp_logits)
+    angular_loss = positive_loss.mean()
+    # print("angular_loss : ", angular_loss)
+
+    return angular_loss, B_
+
+
+def train_epoch(rank, args, train_loader, model, optimizer, epoch, scaler, logger):
+    print('train at epoch %d'%epoch)
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+    dice_losses = AverageMeter()
+    sigmoid_focal_losses = AverageMeter()
+    cos_losses = AverageMeter()
+    model.train()
+    end = time.time()
+    
+    # argument for verb-centric radial contrastive loss
+    mlw = args.metric_loss_weight
+    metric_mode = args.metric_mode
+    filter_thres = args.filter_thres
+    metric_learning = args.metric_learning
+    
+    for batch_idx, (imgs, word_id, word_mask, bbox, seg_map, params) in enumerate(train_loader):
+        B = imgs.size(0) # Original Batch size
+        
+        hp_word_id = params['hp_word_id']
+        hp_word_mask = params['hp_word_mask']
+        hp_bert_embs = params['hardpos_emb'].cuda(non_blocking=True).squeeze(1)
+        pos_type = np.array(params['pos_type'])
+
+        pos_mask = torch.tensor(np.where(pos_type == 'hardpos', 1, 0))
+        
+        # print(hp_bert_embs.shape)
+        # print(imgs.shape, word_id.shape, word_mask.shape, seg_map.shape)
+
+        # hardpos flag outside the model
+        verb_masks = [] 
+        cl_masks = []
+        images = []  
+        targets = []
+        sentences_ = []
+        sentences_masked_ = []
+                
+        for idx in range(len(imgs)) :
+            sentences_.append(word_id[idx])
+            sentences_masked_.append(word_mask[idx])
+            images.append(imgs[idx])
+            targets.append(seg_map[idx])
+
+            # If verb exists, process it
+            if pos_mask[idx] :
+                verb_masks.extend([1, 1])  # Both original sentence and verb are marked
+                cl_masks.extend([1, 0])    # Only original sentence get marked
+                sentences_.append(hp_word_id[idx])
+                sentences_masked_.append(hp_word_mask[idx])
+                images.append(imgs[idx])
+                targets.append(seg_map[idx])
+            else:
+                verb_masks.append(0)
+                cl_masks.append(1)
+
+        imgs, seg_map, word_id, word_mask, verb_masks, cl_masks = \
+                                                        torch.stack(images).cuda(rank, non_blocking=True),\
+                                                        torch.stack(targets).cuda(rank, non_blocking=True),\
+                                                        torch.stack(sentences_).cuda(rank, non_blocking=True),\
+                                                        torch.stack(sentences_masked_).cuda(rank, non_blocking=True),\
+                                                        torch.tensor(verb_masks, dtype=torch.bool).cuda(rank, non_blocking=True),\
+                                                        torch.tensor(cl_masks, dtype=torch.bool).cuda(rank, non_blocking=True)
+
+        image = Variable(imgs)
+        word_id = Variable(word_id)
+        word_mask = Variable(word_mask)
+        seg_map = Variable(seg_map)
+        verb_masks = Variable(verb_masks)
+        cl_masks = Variable(cl_masks)        
+        
+        if hp_bert_embs.numel() > 0 :
+            mask = ~torch.all(hp_bert_embs == 0, dim=1) 
+            hp_bert_embs = hp_bert_embs[mask]
+            # print(hp_bert_embs.shape, hp_bert_embs.requires_grad, hp_bert_embs.device)
+            norms = torch.norm(hp_bert_embs, dim=-1, keepdim=True)
+            normed_embs = hp_bert_embs / norms
+            cosime_sim = torch.mm(normed_embs, normed_embs.T)
+            rows_to_filter, cols_to_filter = torch.where(cosime_sim > filter_thres) 
+
+            # print(normed_embs, normed_embs.requires_grad, normed_embs.device)
+            # print(cosime_sim, cosime_sim.requires_grad, cosime_sim.device)
+            # print("rows_to_filter : ", rows_to_filter, rows_to_filter.requires_grad)
+            # print("cols_to_filter : ", cols_to_filter, cols_to_filter.requires_grad)
+
+            
+            
+        with autocast():
+            mask_out_all, metric_tensors = model(image, word_id, word_mask)
+            loss = 0.            
+            
+            # get mask and seg_map for calculating existing loss function (iou loss, dice loss, sigmoid focal loss)
+            mask_out = mask_out_all[cl_masks]
+            seg_map_cl = seg_map[cl_masks]
+            
+            mask_out_np = mask_out.data.cpu().numpy() # [bs, 1, 208, 208]
+            seg_map_np = seg_map_cl.cpu().numpy() # [bs, 1, 208, 208]
+            seg_iou = cal_seg_iou_loss(seg_map_np, mask_out_np, args.seg_thresh)
+         
+            dice_loss_ = dice_loss(mask_out, seg_map_cl)
+            sigmoid_focal_loss_ = sigmoid_focal_loss(mask_out, seg_map_cl)
+
+            dice_weight, focal_weight = 1.0, 1.0
+            loss = (dice_weight * dice_loss_) + (focal_weight * sigmoid_focal_loss_)
+            
+            # get angular contrastive loss, which involves original & verb pharase pairs (only for pairs where hardpos verb phrase exists)
+            if metric_learning and sum(pos_mask) > 1 :
+                metric_weight = mlw
+                # NS means number of orig-verb pair where verb phrase exists.
+                metric_loss, NS = UniAngularLogitContrastLoss(metric_tensors, verb_masks, rows_to_filter, cols_to_filter, m=args.margin_value, tau=args.temperature, verbonly=True, args=args)
+                loss += metric_weight * metric_loss
+
+        optimizer.zero_grad()
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+
+        losses.update(loss.item(), B)
+        dice_losses.update(dice_loss_.item(), B)
+        sigmoid_focal_losses.update(sigmoid_focal_loss_.item(), B)
+        cos_losses.update(seg_iou.mean().item(), B)
+        
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if rank == 0 and batch_idx % args.print_freq == 0:
+            print_str = 'Epoch: [{0}][{1}/{2}]\t' \
+                'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
+                'Loss {loss.val:.4f} ({loss.avg:.4f})\t' \
+                'dice_losses {dice_losses.val:.4f} ({dice_losses.avg:.4f})\t' \
+                'sigmoid_focal_losses {sigmoid_focal_losses.val:.4f} ({sigmoid_focal_losses.avg:.4f})\t' \
+                'IoU {cos_loss.val:.4f} ({cos_loss.avg:.4f})\t' \
+                .format(epoch, batch_idx, len(train_loader), batch_time=batch_time, loss=losses, dice_losses=dice_losses, sigmoid_focal_losses=sigmoid_focal_losses, cos_loss=cos_losses)
+            print(print_str)
+            logger.info(print_str)
+
+    return losses.avg
+
+
+
+def validate_epoch(args, val_loader, model, logger, mode='val'):
+    print('begin test')
+    batch_time = AverageMeter()
+    miou = AverageMeter()
+    miou_seg = AverageMeter()
+
+    prec=dict()
+    thresholds = np.arange(0.5, 1, 0.05)
+
+    for thresh in thresholds:
+        prec[thresh]= AverageMeter()
+
+    model.eval()
+    end = time.time()
+    idx = 0
+
+    t_all = []
+    total_intersection = 0.0
+    total_union = 0.0
+    
+    for batch_idx, (imgs, word_id, word_mask, bbox, seg_map, ratio, dw, dh, im_id, phrase, draw_img) in enumerate(val_loader):
+        
+        imgs = imgs.cuda(0)
+        word_id = word_id.cuda(0)
+        word_mask = word_mask.cuda(0)
+        seg_map = seg_map.cuda(0)
+        image = Variable(imgs)
+        word_id = Variable(word_id)
+        word_mask = Variable(word_mask)
+        seg_map = Variable(seg_map)
+
+        t1 = time.time()
+        with torch.no_grad():
+            mask_out, _ = model(image, word_id, word_mask)
+            mask_out = mask_out.sigmoid()
+
+        t2 = time.time()
+        t_all.append(t2-t1)
+
+        ## test: convert pred, gt box to original scale with meta-info
+        ih = seg_map.shape[-2]
+        iw = seg_map.shape[-1]
+        nh = int(ih * ratio)
+        nw = int(iw * ratio)
+        top, bottom = int(dh[0]), nh + int(dh[0])
+        left, right = int(dw[0]), nw + int(dw[0])
+        ratio = float(ratio)
+        new_shape = (iw, ih)
+        
+        ## revert image for visualization
+        seg_map_np = seg_map[0,:,:,:].data.cpu().numpy().transpose(1,2,0)
+        seg_map_np = cv2.resize(seg_map_np, new_shape, interpolation=cv2.INTER_CUBIC)
+        img_np = imgs[0,:,top:bottom,left:right].data.cpu().numpy().transpose(1,2,0)
+        img_np = cv2.resize(img_np, new_shape, interpolation=cv2.INTER_CUBIC)
+
+        img_np = Variable(torch.from_numpy(img_np.transpose(2,0,1)).cuda().unsqueeze(0))
+        
+        # seg
+        mask_out = mask_out[0].data.cpu().numpy().transpose(1,2,0)
+        mask_out = cv2.resize(mask_out, (args.size, args.size))
+        mask_out_np = mask_out[top:bottom, left:right]
+        mask_out_np = cv2.resize(mask_out_np, new_shape)
+        # seg_iou, seg_prec = cal_seg_iou(seg_map[0].cpu().numpy(), mask_out_np, args.seg_thresh)
+        seg_iou, seg_prec, inter_sum, union_sum = cal_seg_iou2(seg_map_np, mask_out_np, args.seg_thresh)
+
+        miou_seg.update(seg_iou, imgs.size(0))
+        total_intersection += inter_sum
+        total_union += union_sum
+
+        for thresh in thresholds:
+            prec[thresh].update(seg_prec[thresh], imgs.size(0))
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+        if batch_idx % 1000 == 0:
+            print_str = '[{0}/{1}]\t' \
+                'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
+                'seg_iu {seg.val:.4f} ({seg.avg:.4f})\t' \
+                .format( \
+                    batch_idx, len(val_loader), batch_time=batch_time, seg=miou_seg)
+            print(print_str)
+            logger.info(print_str)  
+        idx = idx + 1
+    overall_iou = (total_intersection + 1e-10) / (total_union + 1e-10)
+
+    print("Mean IoU:", miou_seg.avg)
+    print("Overall IoU:", overall_iou)
+    logger.info("Mean IoU: %.4f" % miou_seg.avg)
+    logger.info("Overall IoU: %.4f" % overall_iou)
+ 
+    for thresh in thresholds:
+            print("prec@%f: %f"%(thresh,float(prec[thresh].avg)))
+            logger.info("prec@%f:%f"%(thresh,float(prec[thresh].avg)))
+    # logger.info("%f,%f"%(float(miou.avg), miou_seg.avg))
+    
+    return miou_seg.avg, prec
+
+

ASDA/engine/engine_gref_sbert_oiou.py

@@ -0,0 +1,340 @@
+import time
+import matplotlib as mpl
+mpl.use('Agg')
+
+import numpy as np
+import torch
+import torch.nn.parallel
+import torch.optim
+from torch.autograd import Variable
+from torch.cuda.amp import autocast as autocast
+
+from model.model_sbert_gref import *
+from dataset.data_loader import *
+from utils.losses import *
+from utils.parsing_metrics import *
+from utils.utils import *
+from utils.utils import dice_loss, sigmoid_focal_loss
+
+use_cuda = torch.cuda.is_available()
+print("use_cuda, ", use_cuda)
+
+
+def return_mask(emb_distance, verb_mask=None, rows_to_filter=None, cols_to_filter=None):
+    B_, B_ = emb_distance.shape
+    positive_mask = torch.zeros_like(emb_distance)
+    positive_mask.fill_diagonal_(1)  # Set diagonal elements to 1 for all cases
+    
+    if B_ < len(verb_mask):
+        # If B_ equals to 2*K (double the number of verb phrase)
+        for i in range(B_ // 2):
+            positive_mask[2 * i, 2 * i + 1] = 1
+            positive_mask[2 * i + 1, 2 * i] = 1
+    else:
+        # Process the case where we have a mix of sentences with and without verbs
+        i = 0
+        while i < B_:
+            if verb_mask[i] == 1:
+                positive_mask[i, i + 1] = 1
+                positive_mask[i + 1, i] = 1
+                i += 2
+            else:
+                i += 1  
+    negative_mask = torch.ones_like(emb_distance) - positive_mask
+    negative_mask = negative_mask.clone() 
+    
+    if rows_to_filter is not None and cols_to_filter is not None :
+        for row, col in zip(rows_to_filter, cols_to_filter):
+            negative_mask[row * 2, col * 2] = 0
+            negative_mask[row * 2, col * 2 + 1] = 0
+            negative_mask[row * 2 + 1, col * 2] = 0
+            negative_mask[row * 2 + 1, col * 2 + 1] = 0
+
+    return positive_mask, negative_mask
+
+
+def UniAngularLogitContrastLoss(total_fq, verb_mask, rows_to_filter, cols_to_filter, alpha=0.5, verbonly=True, m=0.5, tau=0.05, args=None):        
+    _, C, H, W = total_fq.shape
+
+    # Calculate embeddings
+    if verbonly :
+        B = total_fq[verb_mask].shape[0]
+        emb = torch.mean(total_fq[verb_mask], dim=(-1, -2)).reshape(B, C)
+        assert emb.shape[0] % 2 == 0, f"Embedding count {emb.shape[0]} is not divisible by 2."
+    else :
+        emb = torch.mean(total_fq, dim=-1)
+
+    B_ = emb.shape[0]
+    emb_i = emb.unsqueeze(1).repeat(1, B_, 1)  # (B_, B_, C)
+    emb_j = emb.unsqueeze(0).repeat(B_, 1, 1)  # (B_, B_, C)
+
+    sim = nn.CosineSimilarity(dim=-1, eps=1e-6)
+    sim_matrix = sim(emb_i, emb_j).reshape(B_, B_)  # (B_, B_)
+    sim_matrix = torch.clamp(sim_matrix, min=-0.9999, max=0.9999)
+
+    margin_in_radians = m / 57.2958  # Convert degrees to radians
+    theta_matrix = (torch.pi / 2) - torch.acos(sim_matrix)
+    # print("sim_matrix : ", sim_matrix)
+    # print("theta_matrix : ", theta_matrix)
+        
+    positive_mask, negative_mask = return_mask(sim_matrix, verb_mask, rows_to_filter, cols_to_filter)
+    # print("positive_mask : ", positive_mask)
+    # print("negative_mask : ", negative_mask)
+    # print("positive_mask requires_grad:", positive_mask.requires_grad, 
+    #     "device:", positive_mask.device, "dtype:", positive_mask.dtype)
+    # print("negative_mask requires_grad:", negative_mask.requires_grad, 
+    #     "device:", negative_mask.device, "dtype:", negative_mask.dtype)
+
+
+    theta_with_margin = theta_matrix.clone()
+    theta_with_margin[positive_mask.bool()] -= margin_in_radians  
+    logits = theta_with_margin / tau  # Scale with temperature
+    
+    # Compute exp logits for softmax
+    exp_logits = torch.exp(logits) 
+    pos_exp_logits = exp_logits * positive_mask
+    pos_exp_logits = pos_exp_logits.sum(dim=-1)
+    neg_exp_logits = exp_logits * negative_mask
+    neg_exp_logits = neg_exp_logits.sum(dim=-1)
+    
+    total_exp_logits = pos_exp_logits + neg_exp_logits
+
+    positive_loss = -torch.log(pos_exp_logits/ total_exp_logits)
+    angular_loss = positive_loss.mean()
+    # print("angular_loss : ", angular_loss)
+
+    return angular_loss, B_
+
+
+def train_epoch(rank, args, train_loader, model, optimizer, epoch, scaler, logger):
+    print('train at epoch %d'%epoch)
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+    dice_losses = AverageMeter()
+    sigmoid_focal_losses = AverageMeter()
+    cos_losses = AverageMeter()
+    model.train()
+    end = time.time()
+    
+    # argument for verb-centric radial contrastive loss
+    mlw = args.metric_loss_weight
+    metric_mode = args.metric_mode
+    filter_thres = args.filter_thres
+    metric_learning = args.metric_learning
+    
+    for batch_idx, (imgs, word_id, word_mask, bbox, seg_map, params) in enumerate(train_loader):
+        B = imgs.size(0) # Original Batch size
+        
+        hp_word_id = params['hp_word_id']
+        hp_word_mask = params['hp_word_mask']
+        hp_bert_embs = params['hardpos_emb'].cuda(non_blocking=True).squeeze(1)
+        pos_type = np.array(params['pos_type'])
+
+        pos_mask = torch.tensor(np.where(pos_type == 'hardpos', 1, 0))
+        
+        # print(hp_bert_embs.shape)
+        # print(imgs.shape, word_id.shape, word_mask.shape, seg_map.shape)
+
+        # hardpos flag outside the model
+        verb_masks = [] 
+        cl_masks = []
+        images = []  
+        targets = []
+        sentences_ = []
+        sentences_masked_ = []
+                
+        for idx in range(len(imgs)) :
+            sentences_.append(word_id[idx])
+            sentences_masked_.append(word_mask[idx])
+            images.append(imgs[idx])
+            targets.append(seg_map[idx])
+
+            # If verb exists, process it
+            if pos_mask[idx] :
+                verb_masks.extend([1, 1])  # Both original sentence and verb are marked
+                cl_masks.extend([1, 0])    # Only original sentence get marked
+                sentences_.append(hp_word_id[idx])
+                sentences_masked_.append(hp_word_mask[idx])
+                images.append(imgs[idx])
+                targets.append(seg_map[idx])
+            else:
+                verb_masks.append(0)
+                cl_masks.append(1)
+
+        imgs, seg_map, word_id, word_mask, verb_masks, cl_masks = \
+                                                        torch.stack(images).cuda(rank, non_blocking=True),\
+                                                        torch.stack(targets).cuda(rank, non_blocking=True),\
+                                                        torch.stack(sentences_).cuda(rank, non_blocking=True),\
+                                                        torch.stack(sentences_masked_).cuda(rank, non_blocking=True),\
+                                                        torch.tensor(verb_masks, dtype=torch.bool).cuda(rank, non_blocking=True),\
+                                                        torch.tensor(cl_masks, dtype=torch.bool).cuda(rank, non_blocking=True)
+
+        image = Variable(imgs)
+        word_id = Variable(word_id)
+        word_mask = Variable(word_mask)
+        seg_map = Variable(seg_map)
+        verb_masks = Variable(verb_masks)
+        cl_masks = Variable(cl_masks)        
+        
+        if hp_bert_embs.numel() > 0 :
+            mask = ~torch.all(hp_bert_embs == 0, dim=1) 
+            hp_bert_embs = hp_bert_embs[mask]
+            # print(hp_bert_embs.shape, hp_bert_embs.requires_grad, hp_bert_embs.device)
+            norms = torch.norm(hp_bert_embs, dim=-1, keepdim=True)
+            normed_embs = hp_bert_embs / norms
+            cosime_sim = torch.mm(normed_embs, normed_embs.T)
+            rows_to_filter, cols_to_filter = torch.where(cosime_sim > filter_thres) 
+
+            # print(normed_embs, normed_embs.requires_grad, normed_embs.device)
+            # print(cosime_sim, cosime_sim.requires_grad, cosime_sim.device)
+            # print("rows_to_filter : ", rows_to_filter, rows_to_filter.requires_grad)
+            # print("cols_to_filter : ", cols_to_filter, cols_to_filter.requires_grad)
+            
+        with autocast():
+            mask_out_all, metric_tensors = model(image, word_id, word_mask)
+            loss = 0.            
+            
+            # get mask and seg_map for calculating existing loss function (iou loss, dice loss, sigmoid focal loss)
+            mask_out = mask_out_all[cl_masks]
+            seg_map_cl = seg_map[cl_masks]
+            
+            mask_out_np = mask_out.data.cpu().numpy() # [bs, 1, 208, 208]
+            seg_map_np = seg_map_cl.cpu().numpy() # [bs, 1, 208, 208]
+            seg_iou = cal_seg_iou_loss(seg_map_np, mask_out_np, args.seg_thresh)
+         
+            dice_loss_ = dice_loss(mask_out, seg_map_cl)
+            sigmoid_focal_loss_ = sigmoid_focal_loss(mask_out, seg_map_cl)
+
+            dice_weight, focal_weight = 1.0, 1.0
+            loss = (dice_weight * dice_loss_) + (focal_weight * sigmoid_focal_loss_)
+            
+            # get angular contrastive loss, which involves original & verb pharase pairs (only for pairs where hardpos verb phrase exists)
+            if metric_learning and sum(pos_mask) > 1 :
+                metric_weight = mlw
+                # NS means number of orig-verb pair where verb phrase exists.
+                metric_loss, NS = UniAngularLogitContrastLoss(metric_tensors, verb_masks, rows_to_filter, cols_to_filter, m=args.margin_value, tau=args.temperature, verbonly=True, args=args)
+                loss += metric_weight * metric_loss
+
+        optimizer.zero_grad()
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+
+        losses.update(loss.item(), B)
+        dice_losses.update(dice_loss_.item(), B)
+        sigmoid_focal_losses.update(sigmoid_focal_loss_.item(), B)
+        cos_losses.update(seg_iou.mean().item(), B)
+        
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if rank == 0 and batch_idx % args.print_freq == 0:
+            print_str = 'Epoch: [{0}][{1}/{2}]\t' \
+                'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
+                'Loss {loss.val:.4f} ({loss.avg:.4f})\t' \
+                'dice_losses {dice_losses.val:.4f} ({dice_losses.avg:.4f})\t' \
+                'sigmoid_focal_losses {sigmoid_focal_losses.val:.4f} ({sigmoid_focal_losses.avg:.4f})\t' \
+                'IoU {cos_loss.val:.4f} ({cos_loss.avg:.4f})\t' \
+                .format(epoch, batch_idx, len(train_loader), batch_time=batch_time, loss=losses, dice_losses=dice_losses, sigmoid_focal_losses=sigmoid_focal_losses, cos_loss=cos_losses)
+            print(print_str)
+            logger.info(print_str)
+
+    return losses.avg
+
+def validate_epoch(args, val_loader, model, logger, mode='val'):
+    print('begin test')
+    batch_time = AverageMeter()
+    miou = AverageMeter()
+    miou_seg = AverageMeter()
+
+    prec=dict()
+    thresholds = np.arange(0.5, 1, 0.05)
+
+    for thresh in thresholds:
+        prec[thresh]= AverageMeter()
+
+    model.eval()
+    end = time.time()
+    idx = 0
+
+    t_all = []
+    total_intersection = 0.0
+    total_union = 0.0
+    
+    for batch_idx, (imgs, word_id, word_mask, bbox, seg_map, ratio, dw, dh, im_id, phrase, draw_img) in enumerate(val_loader):
+        
+        imgs = imgs.cuda(0)
+        word_id = word_id.cuda(0)
+        word_mask = word_mask.cuda(0)
+        seg_map = seg_map.cuda(0)
+        image = Variable(imgs)
+        word_id = Variable(word_id)
+        word_mask = Variable(word_mask)
+        seg_map = Variable(seg_map)
+
+        t1 = time.time()
+        with torch.no_grad():
+            mask_out, _ = model(image, word_id, word_mask)
+            mask_out = mask_out.sigmoid()
+
+        t2 = time.time()
+        t_all.append(t2-t1)
+
+        ## test: convert pred, gt box to original scale with meta-info
+        ih = seg_map.shape[-2]
+        iw = seg_map.shape[-1]
+        nh = int(ih * ratio)
+        nw = int(iw * ratio)
+        top, bottom = int(dh[0]), nh + int(dh[0])
+        left, right = int(dw[0]), nw + int(dw[0])
+        ratio = float(ratio)
+        new_shape = (iw, ih)
+        
+        ## revert image for visualization
+        seg_map_np = seg_map[0,:,:,:].data.cpu().numpy().transpose(1,2,0)
+        seg_map_np = cv2.resize(seg_map_np, new_shape, interpolation=cv2.INTER_CUBIC)
+        img_np = imgs[0,:,top:bottom,left:right].data.cpu().numpy().transpose(1,2,0)
+        img_np = cv2.resize(img_np, new_shape, interpolation=cv2.INTER_CUBIC)
+
+        img_np = Variable(torch.from_numpy(img_np.transpose(2,0,1)).cuda().unsqueeze(0))
+        
+        # seg
+        mask_out = mask_out[0].data.cpu().numpy().transpose(1,2,0)
+        mask_out = cv2.resize(mask_out, (args.size, args.size))
+        mask_out_np = mask_out[top:bottom, left:right]
+        mask_out_np = cv2.resize(mask_out_np, new_shape)
+        # seg_iou, seg_prec = cal_seg_iou(seg_map[0].cpu().numpy(), mask_out_np, args.seg_thresh)
+        seg_iou, seg_prec, inter_sum, union_sum = cal_seg_iou2(seg_map_np, mask_out_np, args.seg_thresh)
+
+        miou_seg.update(seg_iou, imgs.size(0))
+        total_intersection += inter_sum
+        total_union += union_sum
+
+        for thresh in thresholds:
+            prec[thresh].update(seg_prec[thresh], imgs.size(0))
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+        if batch_idx % 1000 == 0:
+            print_str = '[{0}/{1}]\t' \
+                'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
+                'seg_iu {seg.val:.4f} ({seg.avg:.4f})\t' \
+                .format( \
+                    batch_idx, len(val_loader), batch_time=batch_time, seg=miou_seg)
+            print(print_str)
+            logger.info(print_str)  
+        idx = idx + 1
+    overall_iou = (total_intersection + 1e-10) / (total_union + 1e-10)
+
+    print("Mean IoU:", miou_seg.avg)
+    print("Overall IoU:", overall_iou)
+    logger.info("Mean IoU: %.4f" % miou_seg.avg)
+    logger.info("Overall IoU: %.4f" % overall_iou)
+ 
+    for thresh in thresholds:
+            print("prec@%f: %f"%(thresh,float(prec[thresh].avg)))
+            logger.info("prec@%f:%f"%(thresh,float(prec[thresh].avg)))
+    # logger.info("%f,%f"%(float(miou.avg), miou_seg.avg))
+    return miou_seg.avg, overall_iou, prec

ASDA/engine/engine_oiou.py

@@ -0,0 +1,179 @@
+import time
+import matplotlib as mpl
+mpl.use('Agg')
+
+import numpy as np
+import torch
+import torch.nn.parallel
+import torch.optim
+from torch.autograd import Variable
+from torch.cuda.amp import autocast as autocast
+
+from model.model import *
+from dataset.data_loader import *
+from utils.losses import *
+from utils.parsing_metrics import *
+from utils.utils import *
+from utils.utils import dice_loss, sigmoid_focal_loss
+
+use_cuda = torch.cuda.is_available()
+print("use_cuda, ", use_cuda)
+
+
+def train_epoch(rank, args, train_loader, model, optimizer, epoch, scaler, logger):
+    print('train at epoch %d'%epoch)
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+    dice_losses = AverageMeter()
+    sigmoid_focal_losses = AverageMeter()
+    cos_losses = AverageMeter()
+    model.train()
+    end = time.time()
+
+    for batch_idx, (imgs, word_id, word_mask, bbox, seg_map) in enumerate(train_loader):
+        imgs = imgs.cuda(rank, non_blocking=True)
+        word_id = word_id.cuda(rank, non_blocking=True)
+        word_mask = word_mask.cuda(rank, non_blocking=True)
+        seg_map = seg_map.cuda(rank, non_blocking=True)
+        image = Variable(imgs)
+        word_id = Variable(word_id)
+        word_mask = Variable(word_mask)
+        seg_map = Variable(seg_map)
+
+        with autocast():
+            mask_out = model(image, word_id, word_mask)
+            loss = 0.
+            
+            mask_out_np = mask_out.data.cpu().numpy() # [bs, 1, 208, 208]
+            seg_map_np = seg_map.cpu().numpy() # [bs, 1, 208, 208]
+            seg_iou = cal_seg_iou_loss(seg_map_np, mask_out_np, args.seg_thresh)
+         
+            dice_loss_ = dice_loss(mask_out, seg_map)
+            sigmoid_focal_loss_ = sigmoid_focal_loss(mask_out, seg_map)
+
+            loss += dice_loss_ + sigmoid_focal_loss_
+
+        optimizer.zero_grad()
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+
+        losses.update(loss.item(), imgs.size(0))
+        dice_losses.update(dice_loss_.item(), imgs.size(0))
+        sigmoid_focal_losses.update(sigmoid_focal_loss_.item(), imgs.size(0))
+        cos_losses.update(seg_iou.mean().item(), imgs.size(0))
+        
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if rank == 0 and batch_idx % args.print_freq == 0:
+            print_str = 'Epoch: [{0}][{1}/{2}]\t' \
+                'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
+                'Loss {loss.val:.4f} ({loss.avg:.4f})\t' \
+                'dice_losses {dice_losses.val:.4f} ({dice_losses.avg:.4f})\t' \
+                'sigmoid_focal_losses {sigmoid_focal_losses.val:.4f} ({sigmoid_focal_losses.avg:.4f})\t' \
+                'IoU {cos_loss.val:.4f} ({cos_loss.avg:.4f})\t' \
+                .format(epoch, batch_idx, len(train_loader), batch_time=batch_time, loss=losses, dice_losses=dice_losses, sigmoid_focal_losses=sigmoid_focal_losses, cos_loss=cos_losses)
+            print(print_str)
+            logger.info(print_str)
+
+    return losses.avg
+
+def validate_epoch(args, val_loader, model, logger, mode='val'):
+    print('begin test')
+    batch_time = AverageMeter()
+    miou = AverageMeter()
+    miou_seg = AverageMeter()
+
+    prec=dict()
+    thresholds = np.arange(0.5, 1, 0.05)
+
+    for thresh in thresholds:
+        prec[thresh]= AverageMeter()
+
+    model.eval()
+    end = time.time()
+    idx = 0
+
+    t_all = []
+    total_intersection = 0.0
+    total_union = 0.0
+    
+    for batch_idx, (imgs, word_id, word_mask, bbox, seg_map, ratio, dw, dh, im_id, phrase, draw_img) in enumerate(val_loader):
+        
+        imgs = imgs.cuda(0)
+        word_id = word_id.cuda(0)
+        word_mask = word_mask.cuda(0)
+        seg_map = seg_map.cuda(0)
+        image = Variable(imgs)
+        word_id = Variable(word_id)
+        word_mask = Variable(word_mask)
+        seg_map = Variable(seg_map)
+
+        t1 = time.time()
+        with torch.no_grad():
+            mask_out = model(image, word_id, word_mask)
+            mask_out = mask_out.sigmoid()
+
+        t2 = time.time()
+        t_all.append(t2-t1)
+
+        ## test: convert pred, gt box to original scale with meta-info
+        ih = seg_map.shape[-2]
+        iw = seg_map.shape[-1]
+        nh = int(ih * ratio)
+        nw = int(iw * ratio)
+        top, bottom = int(dh[0]), nh + int(dh[0])
+        left, right = int(dw[0]), nw + int(dw[0])
+        ratio = float(ratio)
+        new_shape = (iw, ih)
+        
+        ## revert image for visualization
+        seg_map_np = seg_map[0,:,:,:].data.cpu().numpy().transpose(1,2,0)
+        seg_map_np = cv2.resize(seg_map_np, new_shape, interpolation=cv2.INTER_CUBIC)
+        img_np = imgs[0,:,top:bottom,left:right].data.cpu().numpy().transpose(1,2,0)
+        img_np = cv2.resize(img_np, new_shape, interpolation=cv2.INTER_CUBIC)
+
+        img_np = Variable(torch.from_numpy(img_np.transpose(2,0,1)).cuda().unsqueeze(0))
+        
+        # seg
+        mask_out = mask_out[0].data.cpu().numpy().transpose(1,2,0)
+        mask_out = cv2.resize(mask_out, (args.size, args.size))
+        mask_out_np = mask_out[top:bottom, left:right]
+        mask_out_np = cv2.resize(mask_out_np, new_shape)
+        # seg_iou, seg_prec = cal_seg_iou(seg_map[0].cpu().numpy(), mask_out_np, args.seg_thresh)
+        seg_iou, seg_prec, inter_sum, union_sum = cal_seg_iou2(seg_map_np, mask_out_np, args.seg_thresh)
+
+        miou_seg.update(seg_iou, imgs.size(0))
+        total_intersection += inter_sum
+        total_union += union_sum
+
+        for thresh in thresholds:
+            prec[thresh].update(seg_prec[thresh], imgs.size(0))
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+        if batch_idx % 1000 == 0:
+            print_str = '[{0}/{1}]\t' \
+                'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
+                'seg_iu {seg.val:.4f} ({seg.avg:.4f})\t' \
+                .format( \
+                    batch_idx, len(val_loader), batch_time=batch_time, seg=miou_seg)
+            print(print_str)
+            logger.info(print_str)  
+        idx = idx + 1
+    overall_iou = (total_intersection + 1e-10) / (total_union + 1e-10)
+
+    print("Mean IoU:", miou_seg.avg)
+    print("Overall IoU:", overall_iou)
+    logger.info("Mean IoU: %.4f" % miou_seg.avg)
+    logger.info("Overall IoU: %.4f" % overall_iou)
+ 
+    for thresh in thresholds:
+            print("prec@%f: %f"%(thresh,float(prec[thresh].avg)))
+            logger.info("prec@%f:%f"%(thresh,float(prec[thresh].avg)))
+    # logger.info("%f,%f"%(float(miou.avg), miou_seg.avg))
+    return miou_seg.avg, overall_iou, prec
+

ASDA/engine/engine_rcc_sbert.py

@@ -0,0 +1,258 @@
+import time
+import matplotlib as mpl
+mpl.use('Agg')
+
+import numpy as np
+import torch
+import torch.nn.parallel
+import torch.optim
+from torch.autograd import Variable
+from torch.cuda.amp import autocast as autocast
+
+from model.model_sbert_gref import *
+from dataset.data_loader import *
+from utils.losses import *
+from utils.parsing_metrics import *
+from utils.utils import *
+from utils.utils import dice_loss, sigmoid_focal_loss
+
+use_cuda = torch.cuda.is_available()
+print("use_cuda, ", use_cuda)
+
+
+def return_mask(emb_distance, rows_to_filter=None, cols_to_filter=None):
+    B_, B_ = emb_distance.shape
+    positive_mask = torch.zeros_like(emb_distance)
+    positive_mask.fill_diagonal_(1)  # Set diagonal elements to 1 for all cases        
+    negative_mask = torch.ones_like(emb_distance) - positive_mask
+    negative_mask = negative_mask.clone() 
+    
+    if rows_to_filter is not None and cols_to_filter is not None :
+        for row, col in zip(rows_to_filter, cols_to_filter):
+            negative_mask[row , col] = 0
+    return positive_mask, negative_mask
+
+
+def UniAngularLogitContrastLoss(total_fq, rows_to_filter, cols_to_filter, alpha=0.5, verbonly=True, m=0.5, tau=0.05, args=None):    
+    _, C, H, W = total_fq.shape
+
+    B = total_fq.shape[0]
+    emb = torch.mean(total_fq, dim=(-1, -2)).reshape(B, C)
+
+    B_ = emb.shape[0]
+    emb_i = emb.unsqueeze(1).repeat(1, B_, 1)  # (B_, B_, C)
+    emb_j = emb.unsqueeze(0).repeat(B_, 1, 1)  # (B_, B_, C)
+
+    sim = nn.CosineSimilarity(dim=-1, eps=1e-6)
+    sim_matrix = sim(emb_i, emb_j).reshape(B_, B_)  # (B_, B_)
+    sim_matrix = torch.clamp(sim_matrix, min=-0.9999, max=0.9999)
+
+    margin_in_radians = m / 57.2958  # Convert degrees to radians
+    theta_matrix = (torch.pi / 2) - torch.acos(sim_matrix)
+    # print("sim_matrix : ", sim_matrix)
+    # print("theta_matrix : ", theta_matrix)
+    positive_mask, negative_mask = return_mask(sim_matrix, rows_to_filter, cols_to_filter)
+
+    theta_with_margin = theta_matrix.clone()
+    theta_with_margin[positive_mask.bool()] -= margin_in_radians  
+    logits = theta_with_margin / tau  # Scale with temperature
+    
+    # Compute exp logits for softmax
+    exp_logits = torch.exp(logits) 
+    pos_exp_logits = exp_logits * positive_mask
+    pos_exp_logits = pos_exp_logits.sum(dim=-1)
+    neg_exp_logits = exp_logits * negative_mask
+    neg_exp_logits = neg_exp_logits.sum(dim=-1)
+    
+    total_exp_logits = pos_exp_logits + neg_exp_logits
+
+
+    positive_loss = -torch.log(pos_exp_logits/ total_exp_logits)
+    angular_loss = positive_loss.mean()
+
+    return angular_loss
+
+def train_epoch(rank, args, train_loader, model, optimizer, epoch, scaler, logger):
+    print('train at epoch %d'%epoch)
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+    dice_losses = AverageMeter()
+    sigmoid_focal_losses = AverageMeter()
+    cos_losses = AverageMeter()
+    model.train()
+    end = time.time()
+    
+    # argument for verb-centric radial contrastive loss
+    mlw = args.metric_loss_weight
+    metric_mode = args.metric_mode
+    filter_thres = args.filter_thres
+    metric_learning = args.metric_learning
+    
+    for batch_idx, (imgs, word_id, word_mask, bbox, seg_map, params) in enumerate(train_loader):
+        B = imgs.size(0) # Original Batch size        
+        hp_bert_embs = params['hardpos_emb'].cuda(non_blocking=True).squeeze(1)
+        
+        imgs = imgs.cuda(rank, non_blocking=True)
+        word_id = word_id.cuda(rank, non_blocking=True)
+        word_mask = word_mask.cuda(rank, non_blocking=True)
+        seg_map = seg_map.cuda(rank, non_blocking=True)
+        image = Variable(imgs)
+        word_id = Variable(word_id)
+        word_mask = Variable(word_mask)
+        seg_map = Variable(seg_map)
+        
+        if hp_bert_embs.numel() > 0 :
+            # print(hp_bert_embs.shape, hp_bert_embs.requires_grad, hp_bert_embs.device)
+            norms = torch.norm(hp_bert_embs, dim=-1, keepdim=True)
+            normed_embs = hp_bert_embs / norms
+            cosime_sim = torch.mm(normed_embs, normed_embs.T)
+            rows_to_filter, cols_to_filter = torch.where(cosime_sim > filter_thres) 
+
+                        
+        with autocast():
+            mask_out, metric_tensors = model(image, word_id, word_mask)
+            loss = 0.       
+            
+            # get mask and seg_map for calculating existing loss function (iou loss, dice loss, sigmoid focal loss)
+            
+            mask_out_np = mask_out.data.cpu().numpy() # [bs, 1, 208, 208]
+            seg_map_np = seg_map.cpu().numpy() # [bs, 1, 208, 208]
+            seg_iou = cal_seg_iou_loss(seg_map_np, mask_out_np, args.seg_thresh)
+         
+            dice_loss_ = dice_loss(mask_out, seg_map)
+            sigmoid_focal_loss_ = sigmoid_focal_loss(mask_out, seg_map)
+
+            dice_weight, focal_weight = 1.0, 1.0
+            loss = (dice_weight * dice_loss_) + (focal_weight * sigmoid_focal_loss_)
+            
+            # get angular contrastive loss, which involves original & verb pharase pairs (only for pairs where hardpos verb phrase exists)
+            if metric_learning :
+                metric_weight = mlw
+                metric_loss = UniAngularLogitContrastLoss(metric_tensors, rows_to_filter, cols_to_filter, m=args.margin_value, tau=args.temperature, verbonly=True, args=args)
+
+                loss += metric_weight * metric_loss
+
+        optimizer.zero_grad()
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+
+        losses.update(loss.item(), B)
+        dice_losses.update(dice_loss_.item(), B)
+        sigmoid_focal_losses.update(sigmoid_focal_loss_.item(), B)
+        cos_losses.update(seg_iou.mean().item(), B)
+        
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if rank == 0 and batch_idx % args.print_freq == 0:
+            print_str = 'Epoch: [{0}][{1}/{2}]\t' \
+                'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
+                'Loss {loss.val:.4f} ({loss.avg:.4f})\t' \
+                'dice_losses {dice_losses.val:.4f} ({dice_losses.avg:.4f})\t' \
+                'sigmoid_focal_losses {sigmoid_focal_losses.val:.4f} ({sigmoid_focal_losses.avg:.4f})\t' \
+                'IoU {cos_loss.val:.4f} ({cos_loss.avg:.4f})\t' \
+                .format(epoch, batch_idx, len(train_loader), batch_time=batch_time, loss=losses, dice_losses=dice_losses, sigmoid_focal_losses=sigmoid_focal_losses, cos_loss=cos_losses)
+            print(print_str)
+            logger.info(print_str)
+
+    return losses.avg
+
+def validate_epoch(args, val_loader, model, logger, mode='val'):
+    print('begin test')
+    batch_time = AverageMeter()
+    miou = AverageMeter()
+    miou_seg = AverageMeter()
+
+    prec=dict()
+    thresholds = np.arange(0.5, 1, 0.05)
+
+    for thresh in thresholds:
+        prec[thresh]= AverageMeter()
+
+    model.eval()
+    end = time.time()
+    idx = 0
+
+    t_all = []
+    total_intersection = 0.0
+    total_union = 0.0
+    
+    for batch_idx, (imgs, word_id, word_mask, bbox, seg_map, ratio, dw, dh, im_id, phrase, draw_img) in enumerate(val_loader):
+        
+        imgs = imgs.cuda(0)
+        word_id = word_id.cuda(0)
+        word_mask = word_mask.cuda(0)
+        seg_map = seg_map.cuda(0)
+        image = Variable(imgs)
+        word_id = Variable(word_id)
+        word_mask = Variable(word_mask)
+        seg_map = Variable(seg_map)
+
+        t1 = time.time()
+        with torch.no_grad():
+            mask_out, _ = model(image, word_id, word_mask)
+            mask_out = mask_out.sigmoid()
+
+        t2 = time.time()
+        t_all.append(t2-t1)
+
+        ## test: convert pred, gt box to original scale with meta-info
+        ih = seg_map.shape[-2]
+        iw = seg_map.shape[-1]
+        nh = int(ih * ratio)
+        nw = int(iw * ratio)
+        top, bottom = int(dh[0]), nh + int(dh[0])
+        left, right = int(dw[0]), nw + int(dw[0])
+        ratio = float(ratio)
+        new_shape = (iw, ih)
+        
+        ## revert image for visualization
+        seg_map_np = seg_map[0,:,:,:].data.cpu().numpy().transpose(1,2,0)
+        seg_map_np = cv2.resize(seg_map_np, new_shape, interpolation=cv2.INTER_CUBIC)
+        img_np = imgs[0,:,top:bottom,left:right].data.cpu().numpy().transpose(1,2,0)
+        img_np = cv2.resize(img_np, new_shape, interpolation=cv2.INTER_CUBIC)
+
+        img_np = Variable(torch.from_numpy(img_np.transpose(2,0,1)).cuda().unsqueeze(0))
+        
+        # seg
+        mask_out = mask_out[0].data.cpu().numpy().transpose(1,2,0)
+        mask_out = cv2.resize(mask_out, (args.size, args.size))
+        mask_out_np = mask_out[top:bottom, left:right]
+        mask_out_np = cv2.resize(mask_out_np, new_shape)
+        # seg_iou, seg_prec = cal_seg_iou(seg_map[0].cpu().numpy(), mask_out_np, args.seg_thresh)
+        seg_iou, seg_prec, inter_sum, union_sum = cal_seg_iou2(seg_map_np, mask_out_np, args.seg_thresh)
+
+        miou_seg.update(seg_iou, imgs.size(0))
+        total_intersection += inter_sum
+        total_union += union_sum
+
+        for thresh in thresholds:
+            prec[thresh].update(seg_prec[thresh], imgs.size(0))
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+        if batch_idx % 1000 == 0:
+            print_str = '[{0}/{1}]\t' \
+                'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
+                'seg_iu {seg.val:.4f} ({seg.avg:.4f})\t' \
+                .format( \
+                    batch_idx, len(val_loader), batch_time=batch_time, seg=miou_seg)
+            print(print_str)
+            logger.info(print_str)  
+        idx = idx + 1
+    overall_iou = (total_intersection + 1e-10) / (total_union + 1e-10)
+
+    print("Mean IoU:", miou_seg.avg)
+    print("Overall IoU:", overall_iou)
+    logger.info("Mean IoU: %.4f" % miou_seg.avg)
+    logger.info("Overall IoU: %.4f" % overall_iou)
+ 
+    for thresh in thresholds:
+            print("prec@%f: %f"%(thresh,float(prec[thresh].avg)))
+            logger.info("prec@%f:%f"%(thresh,float(prec[thresh].avg)))
+    # logger.info("%f,%f"%(float(miou.avg), miou_seg.avg))
+    return miou_seg.avg, overall_iou, prec
+

ASDA/engine/tmp.py

@@ -0,0 +1,292 @@
+import os
+import time
+from tqdm import tqdm
+import cv2
+import numpy as np
+import torch
+import pdb
+import torch.cuda.amp as amp
+import torch.distributed as dist
+import torch.nn.functional as F
+import wandb
+from loguru import logger
+from utils.dataset_verbonly import tokenize
+from utils.misc import (AverageMeter, ProgressMeter, concat_all_gather,
+                        trainMetricGPU)
+
+## todo : add oIoU metric
+def train(train_loader, model, optimizer, scheduler, scaler, epoch,  args):
+    # torch.autograd.set_detect_anomaly(True)
+    batch_time = AverageMeter('Batch', ':2.2f')
+    data_time = AverageMeter('Data', ':2.2f')
+    lr = AverageMeter('Lr', ':1.6f')
+    loss_meter = AverageMeter('Loss', ':2.4f')
+    iou_meter = AverageMeter('IoU', ':2.2f')
+    pr_meter = AverageMeter('Prec@50', ':2.2f')
+    progress = ProgressMeter(
+        len(train_loader),
+        [batch_time, data_time, lr, loss_meter, iou_meter, pr_meter],
+        prefix="Training: Epoch=[{}/{}] ".format(epoch, args.epochs))
+
+
+    model.train()
+    time.sleep(2)
+    end = time.time()
+
+    # size_list = [320, 352, 384, 416, 448, 480, 512]
+    # idx = np.random.choice(len(size_list))
+    # new_size = size_list[idx]
+
+    for i, (image, text, target, hardpos, params) in enumerate(train_loader):
+        data_time.update(time.time() - end)
+
+        # data
+        image = image.cuda(non_blocking=True)
+        text = text.cuda(non_blocking=True)
+        target = target.cuda(non_blocking=True).unsqueeze(1)
+        hardpos = hardpos.cuda(non_blocking=True)
+        hp_emb = params['hardpos_emb'].cuda(non_blocking=True)
+
+        with amp.autocast():
+            pred, target, loss = model(image, text, target, hardpos, hp_emb) # , fq, vis, word, state
+
+   
+        # backward
+        optimizer.zero_grad()
+        # scaler.scale(loss).backward()
+        scaler.scale(loss).backward()
+        # loss.backward()
+
+        # for name, param in model.named_parameters():
+        #     if param.grad is not None:
+        #         if torch.isinf(param.grad).any() or torch.isnan(param.grad).any():
+        #             print(f"Inf/NaN in gradients: {name}")
+        # for name, param in model.named_parameters():
+        #     if param.grad is not None:
+        #         grad_norm = param.grad.norm()
+        #         if torch.isnan(grad_norm):
+        #             print(f"NaN gradient detected in {name}")
+
+        if args.max_norm:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
+
+        # optimizer.step()
+        # scheduler.step()
+        scaler.step(optimizer)
+        scaler.update()
+
+        # metric
+        iou, pr5 = trainMetricGPU(pred, target, 0.35, 0.5)
+        dist.all_reduce(loss.detach())
+        dist.all_reduce(iou)
+        dist.all_reduce(pr5)
+        loss = loss / dist.get_world_size()
+        iou = iou / dist.get_world_size()
+        pr5 = pr5 / dist.get_world_size()
+
+        loss_meter.update(loss.item(), image.size(0))
+        iou_meter.update(iou.item(), image.size(0))
+        pr_meter.update(pr5.item(), image.size(0))
+        lr.update(scheduler.get_last_lr()[-1])
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        # if (i + 1) % args.print_freq == 0:
+        #     progress.display(i + 1)
+        #     if dist.get_rank() in [-1, 0]:
+        #         wandb.log(
+        #             {
+        #                 "time/batch": batch_time.val,
+        #                 "time/data": data_time.val,
+        #                 "training/lr": lr.val,
+        #                 "training/loss": loss_meter.val,
+        #                 "training/iou": iou_meter.val,
+        #                 "training/prec@50": pr_meter.val,
+        #             },
+        #             step=epoch * len(train_loader) + (i + 1))
+
+
+@torch.no_grad()
+def validate(val_loader, model, epoch, args):
+    iou_list = []
+    I_list = []
+    U_list = []
+    model.eval()
+    time.sleep(2)
+    for imgs, texts, masks, param in val_loader:
+        # data
+        imgs = imgs.cuda(non_blocking=True)
+        texts = texts.cuda(non_blocking=True)
+        # inference
+        preds = model(imgs, texts)
+        preds = torch.sigmoid(preds)
+        if preds.shape[-2:] != imgs.shape[-2:]:
+            preds = F.interpolate(preds,
+                                  size=imgs.shape[-2:],
+                                  mode='bicubic',
+                                  align_corners=True).squeeze(1)
+        # process one batch
+        # for pred, mask_dir, mat, ori_size in zip(preds, param['mask_dir'],
+        #                                          param['inverse'],
+        #                                          param['ori_size']):
+        #     h, w = np.array(ori_size)
+        #     mat = np.array(mat)
+        #     pred = pred.cpu().numpy()
+        #     pred = cv2.warpAffine(pred, mat, (w, h),
+        #                           flags=cv2.INTER_CUBIC,
+        #                           borderValue=0.)
+        #     pred = np.array(pred > 0.35)
+        #     mask = cv2.imread(mask_dir, flags=cv2.IMREAD_GRAYSCALE)
+        #     mask = mask / 255.
+        #     # iou
+        #     inter = np.logical_and(pred, mask)
+        #     union = np.logical_or(pred, mask)
+        #     iou = np.sum(inter) / (np.sum(union) + 1e-6)
+        #     iou_list.append(iou)
+        #     I_list.append(inter)
+        #     U_list.append(union)
+        for pred, mask in zip(preds, masks):
+            # h, w = np.array(ori_size)
+            # mat = np.array(mat)
+            pred = pred.cpu().numpy()
+            # pred = cv2.warpAffine(pred, mat, (w, h),
+            #                       flags=cv2.INTER_CUBIC,
+            #                       borderValue=0.)
+            pred = np.array(pred > 0.35)
+            # mask = cv2.imread(mask_dir, flags=cv2.IMREAD_GRAYSCALE)
+            # mask = mask / 255.
+            mask = mask.numpy()
+            # iou
+            inter = np.logical_and(pred, mask)
+            union = np.logical_or(pred, mask)
+            iou = np.sum(inter) / (np.sum(union) + 1e-6)
+            I_list.append(inter)
+            U_list.append(union)
+            iou_list.append(iou)
+
+    iou_list = np.stack(iou_list)
+    iou_list = torch.from_numpy(iou_list).to(imgs.device)
+    iou_list = concat_all_gather(iou_list)
+    
+    I_list = np.stack(I_list)
+    I_list = torch.from_numpy(I_list).to(imgs.device)
+    I_list = concat_all_gather(I_list)
+ 
+    U_list = np.stack(U_list)
+    U_list = torch.from_numpy(U_list).to(imgs.device)
+    U_list = concat_all_gather(U_list)
+
+    overall_I = I_list.sum().item()
+    overall_U = U_list.sum().item()
+    overall_IoU = overall_I / (overall_U + 1e-6)  # to avoid division by zero
+
+    
+    prec_list = []
+    for thres in torch.arange(0.5, 1.0, 0.1):
+        tmp = (iou_list > thres).float().mean()
+        prec_list.append(tmp)
+    iou = iou_list.mean()
+    prec = {}
+    temp = '  '
+    for i, thres in enumerate(range(5, 10)):
+        key = 'Pr@{}'.format(thres * 10)
+        value = prec_list[i].item()
+        prec[key] = value
+        temp += "{}: {:.2f}  ".format(key, 100. * value)
+    head = 'Evaluation: Epoch=[{}/{}]  IoU={:.2f}  OIoU={:.4f}'.format(
+        epoch, args.epochs, 100. * iou.item(), 100. * overall_IoU)
+    logger.info(head + temp)
+    # print(head)
+    
+    # return three results : mIoU, oIoU and prec results
+    return iou.item(), overall_IoU, prec
+
+
+@torch.no_grad()
+def inference(test_loader, model, args):
+    iou_list = []
+    I_list = []
+    U_list = []
+
+    tbar = tqdm(test_loader, desc='Inference:', ncols=100)
+    model.eval()
+    time.sleep(2)
+    for img, mask, param in tbar:
+        # data
+        # img = img.cuda(non_blocking=True)
+        # mask = cv2.imread(param['mask_dir'][0], flags=cv2.IMREAD_GRAYSCALE)
+        img = img.cuda(non_blocking=True)
+        mask = mask[0].cpu().numpy()
+        
+        # dump image & mask
+        if args.visualize:
+            seg_id = param['seg_id'][0].cpu().numpy()
+            img_name = '{}-img.jpg'.format(seg_id)
+            mask_name = '{}-mask.png'.format(seg_id)
+            cv2.imwrite(filename=os.path.join(args.vis_dir, img_name),
+                        img=param['ori_img'][0].cpu().numpy())
+            cv2.imwrite(filename=os.path.join(args.vis_dir, mask_name),
+                        img=mask)
+        # multiple sentences
+        for sent in param['sents']:
+            # mask = mask / 255.
+            text = tokenize(sent, args.word_len, True)
+            text = text.cuda(non_blocking=True)
+            # inference
+            pred = model(img, text)
+            pred = torch.sigmoid(pred)
+            if pred.shape[-2:] != img.shape[-2:]:
+                pred = F.interpolate(pred,
+                                     size=img.shape[-2:],
+                                     mode='bicubic',
+                                     align_corners=True).squeeze()
+            # process one sentence
+            # h, w = param['ori_size'].numpy()[0]
+            # mat = param['inverse'].numpy()[0]
+            pred = pred.cpu().numpy()
+            # pred = cv2.warpAffine(pred, mat, (w, h),
+            #                       flags=cv2.INTER_CUBIC,
+            #                       borderValue=0.)
+            pred = np.array(pred > 0.35)
+            # iou
+            inter = np.logical_and(pred, mask)
+            union = np.logical_or(pred, mask)
+            iou = np.sum(inter) / (np.sum(union) + 1e-6)
+            iou_list.append(iou)
+            I_list.append(inter)
+            U_list.append(union)
+            # dump prediction
+            if args.visualize:
+                pred = np.array(pred*255, dtype=np.uint8)
+                sent = "_".join(sent[0].split(" "))
+                pred_name = '{}-iou={:.2f}-{}.png'.format(seg_id, iou*100, sent)
+                cv2.imwrite(filename=os.path.join(args.vis_dir, pred_name),
+                            img=pred)
+    logger.info('=> Metric Calculation <=')
+    iou_list = np.stack(iou_list)
+    iou_list = torch.from_numpy(iou_list).to(img.device)
+
+    I_list = np.stack(I_list)
+    I_list = torch.from_numpy(I_list).to(img.device)
+    U_list = np.stack(U_list)
+    U_list = torch.from_numpy(U_list).to(img.device)
+    overall_I = I_list.sum().item()
+    overall_U = U_list.sum().item()
+    overall_IoU = overall_I / (overall_U + 1e-6)  # to avoid division by zero
+
+    prec_list = []
+    for thres in torch.arange(0.5, 1.0, 0.1):
+        tmp = (iou_list > thres).float().mean()
+        prec_list.append(tmp)
+    iou = iou_list.mean()
+    prec = {}
+    for i, thres in enumerate(range(5, 10)):
+        key = 'Pr@{}'.format(thres*10)
+        value = prec_list[i].item()
+        prec[key] = value
+    logger.info('IoU={:.2f}  OIoU={:.4f}'.format(100.*iou.item(), 100. * overall_IoU))
+    print('IoU={:.2f}  OIoU={:.4f}'.format(100.*iou.item(), 100. * overall_IoU))
+    for k, v in prec.items():
+        logger.info('{}: {:.2f}.'.format(k, 100.*v))
+
+    return iou.item(), overall_IoU, prec