models/matcher.py

"""
Instance Sequence Matching
Modified from DETR (https://github.com/facebookresearch/detr)
"""
import torch
from scipy.optimize import linear_sum_assignment
from torch import nn
import torch.nn.functional as F

from util.box_ops import box_cxcywh_to_xyxy, generalized_box_iou, multi_iou
from util.misc import nested_tensor_from_tensor_list

INF = 100000000

def dice_coef(inputs, targets):
    inputs = inputs.sigmoid()
    inputs = inputs.flatten(1).unsqueeze(1) # [N, 1, THW]
    targets = targets.flatten(1).unsqueeze(0) # [1, M, THW]
    numerator = 2 * (inputs * targets).sum(2)
    denominator = inputs.sum(-1) + targets.sum(-1)

    # NOTE coef doesn't be subtracted to 1 as it is not necessary for computing costs
    coef = (numerator + 1) / (denominator + 1)
    return coef

def sigmoid_focal_coef(inputs, targets, alpha: float = 0.25, gamma: float = 2):
    N, M = len(inputs), len(targets)
    inputs = inputs.flatten(1).unsqueeze(1).expand(-1, M, -1) # [N, M, THW]
    targets = targets.flatten(1).unsqueeze(0).expand(N, -1, -1) # [N, M, THW]

    prob = inputs.sigmoid()
    ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
    p_t = prob * targets + (1 - prob) * (1 - targets)
    coef = ce_loss * ((1 - p_t) ** gamma)

    if alpha >= 0:
        alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
        coef = alpha_t * coef

    return coef.mean(2) # [N, M]


class HungarianMatcher(nn.Module):
    """This class computes an assignment between the targets and the predictions of the network

    For efficiency reasons, the targets don't include the no_object. Because of this, in general,
    there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions,
    while the others are un-matched (and thus treated as non-objects).
    """

    def __init__(self, cost_class: float = 1, cost_bbox: float = 1, cost_giou: float = 1,
                       cost_mask: float = 1, cost_dice: float = 1, num_classes: int = 1):
        """Creates the matcher

        Params:
            cost_class: This is the relative weight of the classification error in the matching cost
            cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost
            cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost
            cost_mask: This is the relative weight of the sigmoid focal loss of the mask in the matching cost
            cost_dice: This is the relative weight of the dice loss of the mask in the matching cost
        """
        super().__init__()
        self.cost_class = cost_class
        self.cost_bbox = cost_bbox
        self.cost_giou = cost_giou
        self.cost_mask = cost_mask
        self.cost_dice = cost_dice
        self.num_classes = num_classes
        assert cost_class != 0 or cost_bbox != 0 or cost_giou != 0 \
            or cost_mask != 0 or cost_dice != 0, "all costs cant be 0"
        self.mask_out_stride = 4

    @torch.no_grad()
    def forward(self, outputs, targets):
        """ Performs the matching
        Params:
            outputs: This is a dict that contains at least these entries:
                 "pred_logits": Tensor of dim [batch_size, num_queries_per_frame, num_frames, num_classes] with the classification logits
                 "pred_boxes": Tensor of dim [batch_size, num_queries_per_frame, num_frames, 4] with the predicted box coordinates
                 "pred_masks": Tensor of dim [batch_size, num_queries_per_frame, num_frames, h, w], h,w in 4x size
            targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:
                 NOTE: Since every frame has one object at most
                 "labels": Tensor of dim [num_frames] (where num_target_boxes is the number of ground-truth
                           objects in the target) containing the class labels
                 "boxes": Tensor of dim [num_frames, 4] containing the target box coordinates
                 "masks": Tensor of dim [num_frames, h, w], h,w in origin size 
        Returns:
            A list of size batch_size, containing tuples of (index_i, index_j) where:
                - index_i is the indices of the selected predictions (in order)
                - index_j is the indices of the corresponding selected targets (in order)
            For each batch element, it holds:
                len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
        """
        src_logits = outputs["pred_logits"] 
        src_boxes = outputs["pred_boxes"]   
        src_masks = outputs["pred_masks"]   

        bs, nf, nq, h, w = src_masks.shape 

        # handle mask padding issue
        target_masks, valid = nested_tensor_from_tensor_list([t["masks"] for t in targets], 
                                                             size_divisibility=32,
                                                             split=False).decompose()
        target_masks = target_masks.to(src_masks) # [B, T, H, W]

        # downsample ground truth masks with ratio mask_out_stride
        start = int(self.mask_out_stride // 2)
        im_h, im_w = target_masks.shape[-2:]
        
        target_masks = target_masks[:, :, start::self.mask_out_stride, start::self.mask_out_stride] 
        assert target_masks.size(2) * self.mask_out_stride == im_h
        assert target_masks.size(3) * self.mask_out_stride == im_w

        indices = []
        for i in range(bs): 
            out_prob = src_logits[i].sigmoid() 
            out_bbox = src_boxes[i]            
            out_mask = src_masks[i]            

            tgt_ids = targets[i]["labels"]     
            tgt_bbox = targets[i]["boxes"]     
            tgt_mask = target_masks[i]         
            tgt_valid = targets[i]["valid"]    

            # class cost
            # we average the cost on valid frames
            cost_class = []
            for t in range(nf):
                if tgt_valid[t] == 0:
                    continue

                out_prob_split = out_prob[t]    
                tgt_ids_split = tgt_ids[t].unsqueeze(0)     

                # Compute the classification cost.
                alpha = 0.25
                gamma = 2.0
                neg_cost_class = (1 - alpha) * (out_prob_split ** gamma) * (-(1 - out_prob_split + 1e-8).log())
                pos_cost_class = alpha * ((1 - out_prob_split) ** gamma) * (-(out_prob_split + 1e-8).log())
                if self.num_classes == 1:  # binary referred
                    cost_class_split = pos_cost_class[:, [0]] - neg_cost_class[:, [0]]
                else:
                    cost_class_split = pos_cost_class[:, tgt_ids_split] - neg_cost_class[:, tgt_ids_split] 

                cost_class.append(cost_class_split)
            cost_class = torch.stack(cost_class, dim=0).mean(0)  # [q, 1]

            # box cost
            # we average the cost on every frame
            cost_bbox, cost_giou = [], []
            for t in range(nf):
                out_bbox_split = out_bbox[t]    
                tgt_bbox_split = tgt_bbox[t].unsqueeze(0)  

                # Compute the L1 cost between boxes
                cost_bbox_split = torch.cdist(out_bbox_split, tgt_bbox_split, p=1)  

                # Compute the giou cost betwen boxes
                cost_giou_split = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox_split),
                                                box_cxcywh_to_xyxy(tgt_bbox_split))
                
                cost_bbox.append(cost_bbox_split)
                cost_giou.append(cost_giou_split)
            cost_bbox = torch.stack(cost_bbox, dim=0).mean(0)
            cost_giou = torch.stack(cost_giou, dim=0).mean(0)

            # mask cost
            # Compute the focal loss between masks
            cost_mask = sigmoid_focal_coef(out_mask.transpose(0, 1), tgt_mask.unsqueeze(0))

            # Compute the dice loss betwen masks
            cost_dice = -dice_coef(out_mask.transpose(0, 1), tgt_mask.unsqueeze(0))

            # Final cost matrix
            C = self.cost_class * cost_class + self.cost_bbox * cost_bbox + self.cost_giou * cost_giou + \
                self.cost_mask * cost_mask + self.cost_dice * cost_dice  # [q, 1]

            # Only has one tgt, MinCost Matcher
            _, src_ind = torch.min(C, dim=0)
            tgt_ind = torch.arange(1).to(src_ind)  
            indices.append((src_ind.long(), tgt_ind.long()))
            
        # list[tuple], length is batch_size
        return indices
                

def build_matcher(args):
    if args.binary:
        num_classes = 1
    else:
        if args.dataset_file == 'ytvos':
            num_classes = 65 
        elif args.dataset_file == 'davis':
            num_classes = 78
        elif args.dataset_file == 'a2d' or args.dataset_file == 'jhmdb':
            num_classes = 1
        else: 
            num_classes = 91  # for coco
    return HungarianMatcher(cost_class=args.set_cost_class, 
                            cost_bbox=args.set_cost_bbox, 
                            cost_giou=args.set_cost_giou,
                            cost_mask=args.set_cost_mask,
                            cost_dice=args.set_cost_dice,
                            num_classes=num_classes)