atss

import torch

def atss_sampling(anchors, gt_boxes, num_samples):
    """
    anchors: 预设的锚框，形状为(N, 4)
    gt_boxes: 真实的目标框，形状为(M, 4)
    num_samples: 需要采样的样本数量
    """
    N = anchors.shape[0]
    M = gt_boxes.shape[0]
    
    # 计算每个锚框与所有真实框的中心距离
    anchor_centers = (anchors[:, :2] + anchors[:, 2:]) / 2
    gt_centers = (gt_boxes[:, :2] + gt_boxes[:, 2:]) / 2
    distances = torch.cdist(anchor_centers, gt_centers)
    
    # 计算每个锚框的重要性分数
    scores = torch.sum(distances, dim=1)
    
    # 根据分数进行采样
    sampled_indices = torch.topk(scores, num_samples).indices
    
    return sampled_indices

# 示例使用
anchors = torch.tensor([[10, 10, 20, 20], [30, 30, 40, 40]])
gt_boxes = torch.tensor([[15, 15, 25, 25], [35, 35, 45, 45]])
num_samples = 1

sampled_indices = atss_sampling(anchors, gt_boxes, num_samples)
print("Sampled Indices:", sampled_indices)