基于视频的行人重识别结果可视化

在网上搜集了很多帖子，想要白嫖可视化代码。但是发现并没有人分享，是视频重识别太小众了还单帧重识别太受宠了！！！！
为此在郑哲东的基础上改出了基于视频的行人重识别结果可视化。
郑哲东大佬的代码先放在这里：https://github.com/layumi/Person_reID_baseline_pytorch
单帧可视化的代码就是里面的demo.py
言归正传，开始向视频方向延申。

数据集读取

我们这里用的是Mars数据集，就不再赘述了，想要了解详情的去看余慈reid大佬的帖子：基于视频的行人再识别（2）：数据集管理data_manager

标准的loader流程

#######################################################################
# dataleader
def data_loader():
    dataset = data_manager.init_dataset(name=args.dataset)
    
    transform_test = T.Compose([
        T.Resize((args.height, args.width)),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    pin_memory = True if use_gpu else False
    
    queryloader = DataLoader(
        VideoDataset(dataset.query, seq_len=args.seq_len, sample='dense', transform=transform_test),
        batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
        pin_memory=pin_memory, drop_last=False,
    )

    galleryloader = DataLoader(
        VideoDataset(dataset.gallery, seq_len=args.seq_len, sample='dense', transform=transform_test),
        batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
        pin_memory=pin_memory, drop_last=False,
    )
    return dataset, trainloader, queryloader, galleryloader
#######################################################################

test结果保存

苦于测试耗时太长，因此将测试的结果保存，用于后续Rank可视化。保存cam是为了对齐，然而在实际操作中发现是多余的，可以省略，当然label还是要存的。

torch.save(qf, 'query_feature.pt')
torch.save(gf, 'gallery_feature.pt')
io.savemat('save_result.mat', {'query_cam':q_camids,
                               'query_label':q_pids,
                               'gallery_cam':g_camids,
                               'gallery_label':g_pids,
                              })

结果读取

接下来当然是读取了。

image_datasets = data_loader()[0]
# print(image_datasets)

result = scipy.io.loadmat('save_result.mat')
query_feature = torch.load('query_feature.pt')
gallery_feature = torch.load('gallery_feature.pt')

query_cam = result['query_cam'][0]
query_label = result['query_label'][0]
gallery_cam = result['gallery_cam'][0]
gallery_label = result['gallery_label'][0]

query_feature = query_feature.cuda()
gallery_feature = gallery_feature.cuda()

gallery排序

因为都是按照label做的排序，所以这里还是沿用的郑大佬的方法，没有做修改，具体流程就是在排除掉相同ID和相同cam的情况下，按照CMC曲线做降序排序，返回的index就是Rank排序的list。

#######################################################################
# sort the images
def sort_img(qf, ql, qc, gf, gl, gc):
    query = qf.view(-1, 1)
    # print(query.shape)
    score = torch.mm(gf, query)
    score = score.squeeze(1).cpu()
    score = score.numpy()
    # predict index
    index = np.argsort(score)  # from small to large
    index = index[::-1]
    query_index = np.argwhere(gl == ql)
    # same camera
    camera_index = np.argwhere(gc == qc)
    
    junk_index1 = np.argwhere(gl == -1)
    junk_index2 = np.intersect1d(query_index, camera_index)
    junk_index = np.append(junk_index2, junk_index1)

    mask = np.in1d(index, junk_index, invert=True)
    index = index[mask]
    return index
#######################################################################

然后丢一个id进行排序就可以了，说白了就是在gallery里找与query_indel相似度最高的排序。

i = opts.query_index
index = sort_img(query_feature[i], query_label[i], query_cam[i], gallery_feature, gallery_label, gallery_cam)

结果可视化

这里做6帧的Rank-4可视化，当然想做的其他的可以自行修改，为了美观嘛，做成两排的，每排3帧

########################################################################
# Visualize the rank result
img_paths, pid, camid = image_datasets.query[i]
query_label = query_label[i]
try:  # Visualize Ranking Result
    # Graphical User Interface is needed
    fig = plt.figure(figsize=(16, 4))
    for i in range(2):
        for j in range(3):
            ax = plt.subplot(2, 15, (1 if i == 0 else 16) + j)
            ax.axis('off')
            imshow(img_paths[(0 if i == 0 else 2) + j], 'query')
    
    for i in range(4):	# 要做其他Rank排序，将range(4)改成想要的list就行
        for j in range(2):
            for k in range(3):
                ax = plt.subplot(2, 15, (1 if j == 0 else 16) + i * 3 + k + 3)
                ax.axis('off')
                img_path, pid, camid = image_datasets.gallery[index[i]]
                label = gallery_label[index[i]]
                imshow(img_path[(0 if j == 0 else 2) + k])
                if label == query_label:
                    ax.set_title('%d' % (i + 1), color='green')
                else:
                    ax.set_title('%d' % (i + 1), color='red')