| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152 |
- import torch
- import torch.nn as nn
- import torch.nn.functional as F
- from scipy.optimize import linear_sum_assignment
- from .loss_utils import box_cxcywh_to_xyxy, generalized_box_iou
- class HungarianMatcher(nn.Module):
- def __init__(self, cost_class=2.0, cost_bbox=5.0, cost_giou=2.0, alpha=0.25, gamma=2.0):
- super().__init__()
- self.cost_class = cost_class
- self.cost_bbox = cost_bbox
- self.cost_giou = cost_giou
- self.alpha = alpha
- self.gamma = gamma
- assert self.cost_class != 0 or self.cost_bbox != 0 or self.cost_giou != 0, "all costs cant be 0"
- @torch.no_grad()
- def forward(self, outputs, targets):
- bs, num_queries = outputs["pred_logits"].shape[:2]
- # We flatten to compute the cost matrices in a batch
- out_prob = F.sigmoid(outputs["pred_logits"].flatten(0, 1))
- out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4]
- # Also concat the target labels and boxes
- tgt_ids = torch.cat([v["labels"] for v in targets])
- tgt_bbox = torch.cat([v["boxes"] for v in targets])
- # Compute the classification cost
- out_prob = out_prob[:, tgt_ids]
- neg_cost_class = (1 - self.alpha) * (out_prob**self.gamma) * (-(1 - out_prob + 1e-8).log())
- pos_cost_class = self.alpha * ((1 - out_prob)**self.gamma) * (-(out_prob + 1e-8).log())
- cost_class = pos_cost_class - neg_cost_class
- # Compute the L1 cost between boxes
- cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1)
- # Compute the giou cost betwen boxes
- cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox))
-
- # Final cost matrix
- C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou
- C = C.view(bs, num_queries, -1).cpu()
- sizes = [len(v["boxes"]) for v in targets]
- indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]
- return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
|
