anet_clip/backup/demo.py

import numpy as np

# Example similarity matrix with shape [10, 200]
similarity_matrix = np.random.rand(10, 200)

# Example range of indices for each step (stored in center and width arrays)
center = np.random.randint(0, 100, size=(10,))
width = np.random.randint(10, 20, size=(10,))

# Calculate the start and end indices for each step
start_indices = np.clip(center - width // 2, 0, similarity_matrix.shape[1])
end_indices = np.clip(center + width // 2, 0, similarity_matrix.shape[1])

# Generate column indices for each range
col_indices = np.arange(similarity_matrix.shape[1])

# Get topk values and corresponding indices
topk = 5
topk_values = []
topk_indices = []

for start, end in zip(start_indices, end_indices):
    # Slice the similarity matrix within the specified range
    range_values = similarity_matrix[:, start:end]

    # Find the indices of the topk values within the range
    sorted_indices = np.argsort(range_values, axis=1)[:, -topk:]
    sorted_indices += start  # Adjust indices to the absolute position

    # Flatten and concatenate the indices
    row_indices = np.arange(len(sorted_indices))[:, np.newaxis]
    indices_flat = np.ravel_multi_index((row_indices.flatten(), sorted_indices.flatten()), similarity_matrix.shape)
    
    # Append topk values and indices
    topk_values.append(np.take(similarity_matrix, indices_flat))
    topk_indices.append(np.column_stack((row_indices.repeat(topk, axis=1).flatten(), sorted_indices.flatten())))

# Convert lists to arrays
topk_values = np.array(topk_values)
topk_indices = np.array(topk_indices)

print("Topk values within the specified range:", topk_values)
print("Topk indices within the specified range:", topk_indices)