# encoding: utf-8 """ @author: liaoxingyu @contact: sherlockliao01@gmail.com """ import argparse import json import os import sys import time import cv2 import numpy as np import torch from torch.backends import cudnn from modeling import Baseline cudnn.benchmark = True class Reid(object): def __init__(self): # self.cfg = self.prepare('config/softmax_triplet.yml') # self.num_classes = 413 # self.model = Baseline('resnet50_ibn', 100, 1) # state_dict = torch.load('/export/home/lxy/reid_baseline/logs/2019.8.12/bj/ibn_lighting/models/model_119.pth') # self.model.load_params_wo_fc(state_dict['model']) # self.model.cuda() # self.model.eval() self.model = torch.jit.load("reid_model.pt") # self.model.eval() # self.model.cuda() # example = torch.rand(1, 3, 256, 128) # example = example.cuda() # traced_script_module = torch.jit.trace(self.model, example) # traced_script_module.save("reid_model.pt") def demo(self, img_path): img = cv2.imread(img_path) img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB) img = cv2.resize(img, (128, 256)) img = img / 255.0 img = (img - [0.485, 0.456, 0.406]) / [0.229,0.224,0.225] img = img.transpose((2,0,1)).astype(np.float32) img = img[np.newaxis,:,:,:] data = torch.from_numpy(img).cuda().float() output = self.model(data) feat = output.cpu().data.numpy() return feat def prepare_gt(self,json_file): feat = [] label = [] with open(json_file,'r') as f: total = json.load(f) for index in total: label.append(index) feat.append(np.array(total[index])) time_label = [int(i[0:10]) for i in label] return np.array(feat),np.array(label),np.array(time_label) def compute_topk(self,k,feat,feats,label): #num_gallery = feats.shape[0] #new_feat = np.tile(feat,[num_gallery,1]) norm_feat = np.sqrt(np.sum(np.square(feat),axis = -1)) norm_feats = np.sqrt(np.sum(np.square(feats),axis = -1)) matrix = np.sum(np.multiply(feat,feats),axis=-1) dist = matrix / np.multiply(norm_feat,norm_feats) #print('feat:',feat.shape) #print('feats:',feats.shape) #print('label:',label.shape) #print('dist:',dist.shape) index = np.argsort(-dist) #print('index:',index.shape) result = [] for i in range(min(feats.shape[0],k)): print(dist[index[i]]) result.append(label[index[i]]) return result if __name__ == '__main__': img_path = '/export/home/lxy/reid_demo/imgs' reid = Reid() img1 = ['1-1.png', '1-2.png', '1-3.png', '1-4.png', '1-5.png'] img2 = ['2-1.png', '2-2.png'] for i in range(len(img1)): for j in range(len(img2)): out1 = reid.demo(os.path.join(img_path, img1[i])) out2 = reid.demo(os.path.join(img_path, img2[j])) innerProduct = np.dot(out1, out2.T) cosineSimilarity = innerProduct / (np.linalg.norm(out1, ord=2) * np.linalg.norm(out2, ord=2)) print('img {} and img {} cosine similarity is {:.4f}'.format(img1[i], img2[j], cosineSimilarity[0][0]))