2021全国人工智能大赛(NAIC)视觉编码赛道初赛一阶段baseline分享

Baseline思路

1、由于一阶段只做reid任务，所以直接计算特征之间的余弦距离或欧式距离就可以得到0.890的指标。

2、大赛鼓励先进行特征压缩编码再进行reid【但是一阶段肯定原始特征做reid是最香的了emmm】

代码

说明：

1、直接跑的话，cpu要足够大。
2、推荐使用faiss库，cpu或gpu版本，可在清华镜像下载安装。

from collections import OrderedDict
import numpy as np
import torch
import torch.nn as nn
import os
import json
from collections import OrderedDict

query_path='2021人工智能/test_A/query_feature_A'
gallery_path='2021人工智能/test_A/gallery_feature_A'
save_path='2021人工智能/result.json'

query_list = os.listdir(query_path)
gallery_list = os.listdir(gallery_path)

num_query = len(query_list)
num_gallery = len(gallery_list)

x = []
y = []
for i in range(num_query):
    with open(os.path.join(query_path, query_list[i]), 'rb') as r:
        img_q = np.frombuffer(r.read(), dtype=np.float32)
        x.append(img_q)
x = np.asarray(x)
# print(x.shape) # [20000, 2048]

for j in range(num_gallery):
    with open(os.path.join(gallery_path, gallery_list[j]), 'rb') as r:
        img_g = np.frombuffer(r.read(), dtype=np.float32)
        y.append(img_g)
y = np.asarray(y)
# print(y.shape) #[428794, 2048]

x=torch.tensor(x)
y=torch.tensor(y)

m, n = x.size(0), y.size(0)

## 这里用的是欧式距离，可自己换为余弦距离。另外输入数据并没norm。
dist_m = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(m, n) + \
       torch.pow(y, 2).sum(dim=1, keepdim=True).expand(n, m).t()
dist_m.addmm_(1, -2, x, y.t())  ## [m ,n]

dist_m = dist_m.data.numpy()

del x  ## 防止cpu爆掉
del y

indices = np.argsort(dist_m, axis=1)

del dist_m

gallery_array = []
results = OrderedDict()
results_save = OrderedDict()

for k in range(num_query):
    gallery_array.append(gallery_list)  
gallery_array = np.asarray(gallery_array)  ## 行为m，每一行顺序都和distmat保持一致

for v in range(num_query):
    results[query_list[v].split('.')]=gallery_array[v][indices[v]]  ### 排好序的结果
    results_save[query_list[v].split('.')] = gallery_array[v][:100] ### 保存前100个
    
with open(save_path,'w',encoding='utf-8') as fg:
    save_submit_path=json.dumps(results_save, indent=4)
    fg.write(save_submit_path)

总结

1、以上代码结果：0.88-0.89之间，如果还想提点，建议思路：re-ranking（做reid方向的大佬们知道我在讲什么…）