|
|
@ -71,7 +71,7 @@ class HungarianMatcher(nn.Module):
|
|
|
|
bs, nq, nc = pred_scores.shape
|
|
|
|
bs, nq, nc = pred_scores.shape
|
|
|
|
|
|
|
|
|
|
|
|
if sum(gt_groups) == 0:
|
|
|
|
if sum(gt_groups) == 0:
|
|
|
|
return [(torch.tensor([], dtype=torch.int32), torch.tensor([], dtype=torch.int32)) for _ in range(bs)]
|
|
|
|
return [(torch.tensor([], dtype=torch.long), torch.tensor([], dtype=torch.long)) for _ in range(bs)]
|
|
|
|
|
|
|
|
|
|
|
|
# We flatten to compute the cost matrices in a batch
|
|
|
|
# We flatten to compute the cost matrices in a batch
|
|
|
|
# [batch_size * num_queries, num_classes]
|
|
|
|
# [batch_size * num_queries, num_classes]
|
|
|
@ -107,7 +107,7 @@ class HungarianMatcher(nn.Module):
|
|
|
|
indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(gt_groups, -1))]
|
|
|
|
indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(gt_groups, -1))]
|
|
|
|
gt_groups = torch.as_tensor([0, *gt_groups[:-1]]).cumsum_(0)
|
|
|
|
gt_groups = torch.as_tensor([0, *gt_groups[:-1]]).cumsum_(0)
|
|
|
|
# (idx for queries, idx for gt)
|
|
|
|
# (idx for queries, idx for gt)
|
|
|
|
return [(torch.tensor(i, dtype=torch.int32), torch.tensor(j, dtype=torch.int32) + gt_groups[k])
|
|
|
|
return [(torch.tensor(i, dtype=torch.long), torch.tensor(j, dtype=torch.long) + gt_groups[k])
|
|
|
|
for k, (i, j) in enumerate(indices)]
|
|
|
|
for k, (i, j) in enumerate(indices)]
|
|
|
|
|
|
|
|
|
|
|
|
def _cost_mask(self, bs, num_gts, masks=None, gt_mask=None):
|
|
|
|
def _cost_mask(self, bs, num_gts, masks=None, gt_mask=None):
|
|
|
|