File size: 1,899 Bytes
987ed1b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | import torch
import pdb
def euclidean_distance(src, tar, reduction='mean'):
# B, (N), T, D
diff = src - tar
dist = torch.norm(diff, dim=-1)
if reduction == 'mean':
return dist.mean(dim=-1)
elif reduction == 'none':
return dist
def coverage_distance(src, tar, penalty=1000.0):
B, S, _ = src.shape
_, T, _ = tar.shape
# Compute all pairwise Euclidean distances
dist_matrix = euclidean_distance(tar.unsqueeze(2), src.unsqueeze(1), 'none')
# Prepare to track the minimum distance for each target and the indices of matches
min_distances = torch.full((B, T), float('inf'), device=src.device)
matched_indices = torch.full((B, T), -1, dtype=torch.long, device=src.device)
last_matched_indices = torch.full((B,), -1, dtype=torch.long, device=src.device)
for t in range(T):
for s in range(S):
valid_mask = (s > last_matched_indices[:, None]).squeeze()
valid_distances = torch.where(valid_mask, dist_matrix[:, t, s], float('inf'))
min_values, min_idxs = torch.min(valid_distances.unsqueeze(-1), dim=1)
min_mask = (min_values < min_distances[:, t])
min_distances[:, t] = torch.where(min_mask, min_values, min_distances[:, t])
matched_indices[:, t] = torch.where(min_mask, min_idxs + s, matched_indices[:, t])
# Update last matched indices
last_matched_indices = torch.where(min_mask, s * torch.ones_like(last_matched_indices), last_matched_indices)
# Apply penalties where no match was found
no_match = (matched_indices[:, t] == -1)
min_distances[:, t] = torch.where(no_match, torch.tensor(penalty, device=src.device), min_distances[:, t])
# Calculate the final distances
final_distances = min_distances.sum(dim=1)
# print('Matched Indices:', matched_indices)
return final_distances
|