class Reid_feature():
def __init__(self):
args = get_parser().parse_args()
cfg = setup_cfg(args)
self.demo = FeatureExtractionDemo(cfg, parallel=args.parallel)
def __call__(self, img_list):
t1 = time.time()
feat = self.demo.run_on_batch(img_list)
# print('reid time:', time.time() - t1, len(img_list))
return feat
# args.input = glob.glob(os.path.expanduser(args.input[0]))
# assert args.input, "The input path(s) was not found"
# for path in tqdm.tqdm(os.listdir(args.input)):
t1 = time.time()
img = cv2.imread(args.input)
im = img[:, :, ::-1] # reid 前处理
im = cv2.resize(im, (128, 256), interpolation=cv2.INTER_CUBIC)
img_list.append(torch.as_tensor(im.astype("float32").transpose(2, 0, 1))[None])
img_list.append(torch.as_tensor(im.astype("float32").transpose(2, 0, 1))[None])
img_list.append(torch.as_tensor(im.astype("float32").transpose(2, 0, 1))[None])
img_list.append(torch.as_tensor(im.astype("float32").transpose(2, 0, 1))[None])
img_list.append(torch.as_tensor(im.astype("float32").transpose(2, 0, 1))[None])
img_list.append(torch.as_tensor(im.astype("float32").transpose(2, 0, 1))[None])
img_list.append(torch.as_tensor(im.astype("float32").transpose(2, 0, 1))[None])
# img_list.append(img)
# feat = demo.run_on_image(img)
# print(img_list[0].shape)
feat = demo.run_on_batch(img_list)
print('pytorch time:', time.time() - t1)
pytorch_output = feat.numpy()
# print(pytorch_output.shape)
# sim = np.dot(pytorch_output, onnx_output.transpose())
# print('====sim:', sim)
# np.testing.assert_allclose(pytorch_output, onnx_output, rtol=1e-4, atol=1e-6)
# print("恭喜你 ^^ ，onnx 和 pytorch 结果一致 ， Exported model has been executed decimal=5 and the result looks good!")
# np.save(os.path.join(args.output, path.replace('.jpg', '.npy').split('/')[-1]), feat)