mirror of https://github.com/JDAI-CV/fast-reid.git
66 lines
1.7 KiB
Python
66 lines
1.7 KiB
Python
import random
|
|
import numpy as np
|
|
import cv2
|
|
import fs
|
|
import argparse
|
|
import io
|
|
import sys
|
|
import torch
|
|
import time
|
|
import os
|
|
import torchvision.transforms as T
|
|
|
|
sys.path.append('../../..')
|
|
sys.path.append('../')
|
|
from fastreid.config import get_cfg
|
|
from fastreid.modeling.meta_arch import build_model
|
|
from fastreid.utils.file_io import PathManager
|
|
from fastreid.utils.checkpoint import Checkpointer
|
|
from fastreid.utils.logger import setup_logger
|
|
from fastreid.data import build_reid_train_loader, build_reid_test_loader
|
|
from fastreid.evaluation.rank import eval_market1501
|
|
|
|
from build.pybind_interface.ReID import ReID
|
|
|
|
|
|
FEATURE_DIM = 2048
|
|
GPU_ID = 0
|
|
|
|
def map(wrapper):
|
|
model = wrapper
|
|
cfg = get_cfg()
|
|
test_loader, num_query = build_reid_test_loader(cfg, "Market1501", T.Compose([]))
|
|
|
|
feats = []
|
|
pids = []
|
|
camids = []
|
|
|
|
for batch in test_loader:
|
|
for image_path in batch["img_paths"]:
|
|
t = torch.Tensor(np.array([model.infer(cv2.imread(image_path))]))
|
|
t.to(torch.device(GPU_ID))
|
|
feats.append(t)
|
|
pids.extend(batch["targets"].numpy())
|
|
camids.extend(batch["camids"].numpy())
|
|
|
|
feats = torch.cat(feats, dim=0)
|
|
q_feat = feats[:num_query]
|
|
g_feat = feats[num_query:]
|
|
q_pids = np.asarray(pids[:num_query])
|
|
g_pids = np.asarray(pids[num_query:])
|
|
q_camids = np.asarray(camids[:num_query])
|
|
g_camids = np.asarray(camids[num_query:])
|
|
|
|
|
|
distmat = 1 - torch.mm(q_feat, g_feat.t())
|
|
distmat = distmat.numpy()
|
|
all_cmc, all_AP, all_INP = eval_market1501(distmat, q_pids, g_pids, q_camids, g_camids, 5)
|
|
mAP = np.mean(all_AP)
|
|
print("mAP {}, rank-1 {}".format(mAP, all_cmc[0]))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
infer = ReID(GPU_ID)
|
|
infer.build("../build/sbs_R50-ibn.engine")
|
|
map(infer)
|