最近在做CBIR的时候碰到了这样一个问题:
给定一张图片,在数据库中检索和输入图片相似的图片。比如给定n张图片作为输入,输出数据库中各自top k检索结果。然后需要将这些图片分别再根据特征相似度得分进行重排。例:四张图片的top 3检索结果的index为
1: 23,34,32
2: 93,23,21
3: 52,12,13
4: 37,62,10
对应的局部特征相似度为:
1:0.3, 0.8, 0.1
2:0.4, 0.2, 0.8
3:0.1, 0.9, 0.7
4:0.6, 0.4, 0.8
需要分别将四张输入图片各自检索出的图片index根据得分进行重排,即输出应为:
1: 34,23,32
2: 21,93,23
3: 12,13,52
4: 10,37,62
传统方法,用for循环,数据量一大就很费时。想用numba加速,又需要避免跳出numpy模块。来看代码:
import numpy as np
idxs = np.arange(12).reshape(4,3)
scores = np.random.random((4,3))
print(labeks)
print(scores)
# idxs
[[ 0 1 2]
[ 3 4 5]
[ 6 7 8]
[ 9 10 11]]
# scores
[[0.13305936, 0.61524316, 0.06763024],
[0.77922301, 0.09916946, 0.76566095],
[0.47272886, 0.34700999, 0.21310556],
[0.66733795, 0.89457897, 0.95406421]]
基本思路是先将两个二维数组拼接成一个三维数组,然后按照score对数组进行排序,最后返回排序后的label,三行代码搞定:
def rerank_label_by_score(idx, score):
idx_score=np.stack((idx,score), axis=2)
sortIdx = np.argsort(-idx_score[:,:,1]) #因为从大到小排,所以取负
array_sorted = idx_score[np.arange(len(idx_score))[:,np.newaxis], sortIdx]
return array_sorted[:,:,0].astype(int)
label_reranked = rerank_label_by_score(labels, scores)
# label_reranked
[[ 1, 0, 2],
[ 3, 5, 4],
[ 6, 7, 8],
[11, 10, 9]]