[深度学习-实战项⽬]以图搜图Resnet+LSH-特征编码图像检索相似度计算参考代码来源于
以图搜图
1. 写在最前⾯
⼊职新公司以后⼀直在搞项⽬，没什么时间写博客。
最近⼀个项⽬是以图搜图项⽬，主要⽤到的技术就是⽬标检测（yolo）+图像检索（ResNet+LSH）。
⽬标检测就不⽤多说了，成熟和现成的代码⼀抓⼀⼤把，主要问题就是在优化提升精度和性能上的摸索。
图像检索的技术也挺多，但是⽹上的资源相对较少，所以记录⼀下这段时间⽤到的⼀个代码。
这个以图搜图和⼈脸识别技术其实很像，可以说是⼀样。⽆⾮就是提取特征，然后进⾏相似度计算。所以相关的技术有ReID，Arcface，以及我在调研的时候有看到⼀个素描草图的图像匹配的研究。
2. 源码解析
（1）跑通代码-即测试⼀下⾃⼰的图⽚ demo.py
图片下载站源码我只⽤到⾥⾯的编码和检索部分。运⾏demo.py。直接跑通这部分，然后缺什么库函数去pip install就⾏了。
ieval_feature import AntiFraudFeatureDataset
ieval_index import EvaluteMap
if __name__ =='__main__':
hash_size =0
input_dim =2048
num_hashtables =1
img_dir ='ImageRetrieval/data'#存放所有图像库的图⽚
test_img_dir ='./images'# 待检索的图像
network ='./weights/gl18-tl-resnet50-gem-w-83fdc30.pth'# 模型权重
#下⾯这⼏个好像没有⽤，不管他
out_similar_dir ='./output/similar'
out_similar_file_dir ='./output/similar_file'
all_csv_file ='./output/aaa.csv'
feature_dict, lsh = AntiFraudFeatureDataset(img_dir, network).constructfeature(hash_size, input_dim, num_hashtables)
test_feature_dict = AntiFraudFeatureDataset(test_img_dir, network).test_feature()
EvaluteMap(out_similar_dir, out_similar_file_dir, all_csv_file).retrieval_images(test_feature_dict, lsh,3)
（2）特征编码
代码⾸先对img_dir中的所有图⽚进⾏特征提取：feature_dict, lsh = AntiFraudFeatureDataset(img_dir, network).constructfeature(hash_size,
input_dim, num_hashtables)
返回的feature_dict就是图⽚特征。（可以直接⽤余弦相似度进⾏相似计算）
但是这⾥还通过LSH对每张图⽚特征图进⾏0，1编号，所在这⾥后⾯⽤来图⽚检索的不是feature_dict，⽽是lsh，（应该是加速后⾯图⽚检索时候的速度）
进到特征编码那块代码retrieval_feature.py，⾥⾯主要对图⽚进⾏编码的函数对象是AntiFraudFeatureDataset
⾸先前⾯⼀⼤段到net.eval()，都是加载⽹络模型，可以看到模型选择有很多参数，这些参数对应⽹络的结构设置，（后⾯如果⽤⾃⼰的数据对⾃⼰的特征编码模型进⾏训练的话，要根据使⽤的不同模型参数进⾏修改）
这个函数ImageProcess是遍历⽬录底下的全部图⽚，并将他们的路径保存在数组中。
然后再这个函数extract_vectors中提取图像特征。（在这个⽬录底下ImageRetrieval-LSH/cirtorch/networks/imageretrievalnet.py ）主要也不需要怎么做修改，除⾮说你要修改⼀下图⽚的dataloader（这⾥是通过将所有图⽚路径保存下来做的dataset，因为每张图⽚的尺⼨可以不⼀样，Resnet⽹络的最后通过⼀个全连接层输出1 * 2048特征图。）
所以这⾥出来的vecs是N张图⽚的特征编码，每个特征编码是1 * 2048。
def constructfeature(self, hash_size, input_dim, num_hashtables):
multiscale ='[1]'
print(">> Loading network:\n>>>> '{}'".format(selfwork))
state = torch.load(selfwork)
net_params ={}
net_params['architecture']= state['meta']['architecture']
net_params['pooling']= state['meta']['pooling']
net_params['local_whitening']= state['meta'].get('local_whitening',False)
net_params['regional']= state['meta'].get('regional',False)
net_params['whitening']= state['meta'].get('whitening',False)
net_params['mean']= state['meta']['mean']
net_params['std']= state['meta']['std']
net_params['pretrained']=False
# network initialization
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
print(">>>> loaded network: ")
a_repr())
# setting up the multi-scale parameters
ms =list(eval(multiscale))
print(">>>> Evaluating scales: {}".format(ms))
# moving network to gpu and eval mode
if torch.cuda.is_available():
net.cuda()
net.eval()
# set up the transform 数据预处理
normalize = transforms.Normalize(
a['mean'],
a['std']
)
transform = transforms.Compose([
transforms.ToTensor(),
normalize
]
)
# extract database and query vectors 对图⽚进⾏编码提取数据库图⽚特征
print('>> ')
images = ImageProcess(self.img_dir).process()
vecs, img_paths = extract_vectors(net, images,1024, transform, ms=ms)
feature_dict =dict(zip(img_paths,list(vecs.detach().cpu().numpy().T)))
# index
lsh = LSHash(hash_size=int(hash_size), input_dim=int(input_dim), num_hashtables=int(num_hashtables))
for img_path, vec in feature_dict.items():
lsh.index(vec.flatten(), extra_data=img_path)
# ## 保存索引模型
# with open(self.feature_path, "wb") as f:
#    pickle.dump(feature_dict, f)
# with open(self.index_path, "wb") as f:
#    pickle.dump(lsh, f)
print("extract feature is done")
return feature_dict, lsh
（3）图像检索
这⾥图像检索这块我没怎么改动，因为只是测试⼀下⾃⼰训练后的模型的效果⽐较⽅便查看⽤的。所以我只是修改了输出的数量。
这⾥如果要输出多个Top，要⾃⼰多加⼏个，（也可以⾃⼰写个循环，我⽐较懒，没有写）然后后⾯我还显⽰出了得分情况，（因为后⾯要进⾏模型的对⽐）
def find_similar_img_gyz(self, feature_dict, lsh, num_results):
for q_path, q_vec in feature_dict.items():
try:
response = lsh.query(q_vec.flatten(), distance_func="cosine")# , num_results=int(num_results) # print(response[0][1])
# print(np.rint(100 * (1 - response[0][1])))
query_img_path0 = response[0][0][1]
query_img_path1 = response[1][0][1]
query_img_path2 = response[2][0][1]
query_img_path3 = response[3][0][1]
query_img_path4 = response[4][0][1]
score_img_path0 = response[0][1]
score_img_path1 = response[1][1]
score_img_path2 = response[2][1]
score_img_path3 = response[3][1]
score_img_path4 = response[4][1]
# score0 = response[0][1]
# score0 = np.rint(100 * (1 - score0))
print('**********************************************')
print('input img: {}'.format(q_path))
print('query0 img: {}'.format(query_img_path0),
' score:{}'.format(np.rint(100*(1- score_img_path0))))
print('query1 img: {}'.format(query_img_path1),
' score:{}'.format(np.rint(100*(1- score_img_path1))))
print('query2 img: {}'.format(query_img_path2),
' score:{}'.format(np.rint(100*(1- score_img_path2))))
print('query3 img: {}'.format(query_img_path3),
' score:{}'.format(np.rint(100*(1- score_img_path3))))
print('query4 img: {}'.format(query_img_path4),
' score:{}'.format(np.rint(100*(1- score_img_path4))))
except:
continue
3. 训练⾃⼰的数据集
（1）训练参数配置
# network architecture and initialization options
parser.add_argument('--arch','-a', metavar='ARCH', default='resnet50', choices=model_names, help='model architecture: '+
' | '.join(model_names)+
' (default: resnet101)')
parser.add_argument('--pool','-p', metavar='POOL', default='gem', choices=pool_names,
help='pooling options: '+
' | '.join(pool_names)+
' (default: gem)')
parser.add_argument('--local-whitening','-lw', dest='local_whitening', action='store_true',
help='train model with learnable local whitening (linear layer) before the pooling')
parser.add_argument('--regional','-r', dest='regional', action='store_true',
help='train model with regional pooling using fixed grid')
parser.add_argument('--whitening','-w', dest='whitening', action='store_true',
help='train model with learnable whitening (linear layer) after the pooling')
parser.add_argument('--not-pretrained', dest='pretrained', action='store_false',
help='initialize model with random weights (default: pretrained on imagenet)')
parser.add_argument('--loss','-l', metavar='LOSS', default='contrastive',
choices=loss_names,
help='training loss options: '+
' | '.join(loss_names)+
' (default: contrastive)')
parser.add_argument('--loss-margin','-lm', metavar='LM', default=0.7,type=float,
help='loss margin: (default: 0.7)')
（2）训练数据的准备
这⾥训练数据⽤的是retrieval-SfM-120k但是因为数据集38个GB，⽹速不⾏，在外⽹上下不下来，所以⽓急败坏的我直接看他的标签⽂
件retrieval-SfM-120k.pkl。
这个⽂件就是⼀个字典格式⽂件，⼤概分了⼏层如下，因为我没有准备验证集和测试集，所以训练时候测试和验证那部分我直接删去了（主要因为测试集的格式和训练集不⼀样，我懒得再去解析另⼀个数据集的格式）。
{ train : {
cids : [ ], cluster : [ ], qidxs : [ ], pidxs : [ ]
},
val : {…}
}
① cids：主要⽤来存放所有图⽚的路径，所以不管你图⽚存放在哪，只要有图⽚路径即可。数组长度就是总的图⽚数量。
② cluster：这个是存放该图⽚的类别，数组的长度和cids⼀样，类别⼀⼀对应cids的图⽚（retrieval-SfM-120k 是有713个建筑物所以是713类，依据⾃⼰的数据集⽽定，我的数据集每对图⽚都是⼀个类，所以有⼏千个类别 ）。
③ qidxs&pidxs：这个qidx是存放查询的query图⽚，对应位置的pidx是存放和他匹配的positive图⽚，这样就形成了⼀对正样例。然后这两个数组存放的是前⾯cids的索引index，对应的是cids[qidxs[1]] -> cids[pidxs[1]]。
所以⾃⼰准备⼀个⾃⼰数据集的pkl⽂件，就可以训练了。
更新⼀下制作标签的代码（在这⼀节最后，⾃⼰写的，所以有的粗糙。）
把你准备训练的数据放在⼀个⽬录下，例如：dirs = 'cirtorch/ImageRetrieval_dataset/train'如下图。train下每个⽂件夹都是相似图⽚的集合。（我⾃⼰每个相似图只有两张，所以每个qp⽂件夹下只有两张，这份代码应该可以⽀持多张相似图，⼏个⽉前写的了，有点健忘 。）
然后得到pkl⽂件夹就可以修改训练代码了。

版权声明：本站内容均来自互联网，仅供演示用，请勿用于商业和其他非法用途。如果侵犯了您的权益请与我们联系QQ:729038198，我们将在24小时内删除。