#!/usr/bin/env python
|
# -*- coding: utf-8 -*-
|
# @Time : 2020/10/26 11:48
|
# @Author : Scheaven
|
# @File : inference_net.py
|
# @description:
|
|
import torch
|
import sys
|
sys.path.append('.')
|
|
from data.data_utils import read_image
|
from predictor import ReID_Model
|
from config import get_cfg
|
from data.transforms.build import build_transforms
|
from engine.defaults import default_argument_parser, default_setup
|
import time
|
|
def setup(args):
|
"""
|
Create configs and perform basic setups.
|
"""
|
cfg = get_cfg()
|
cfg.merge_from_file(args.config_file)
|
cfg.merge_from_list(args.opts)
|
cfg.freeze()
|
default_setup(cfg, args)
|
return cfg
|
|
if __name__ == '__main__':
|
args = default_argument_parser().parse_args()
|
cfg = setup(args)
|
cfg.defrost()
|
cfg.MODEL.BACKBONE.PRETRAIN = False
|
|
model = ReID_Model(cfg)
|
|
test_transforms = build_transforms(cfg, is_train=False)
|
# print (args.img_a1)
|
img_a1 = read_image(args.img_a1)
|
img_a2 = read_image(args.img_a2)
|
img_b1 = read_image(args.img_b1)
|
img_b2 = read_image(args.img_b2)
|
|
img_a1 = test_transforms(img_a1)
|
img_a2 = test_transforms(img_a2)
|
img_b1 = test_transforms(img_b1)
|
img_b2 = test_transforms(img_b2)
|
|
out = torch.zeros((2, *img_a1.size()), dtype=img_a1.dtype)
|
out[0] += img_a1
|
out[1] += img_a2
|
t1 = time.time()
|
qurey_feat = model.run_on_image(out)
|
t2 = time.time()
|
print("t2-t1:", t2-t1)
|
|
similarity1 = torch.cosine_similarity(qurey_feat[0], qurey_feat[1], dim=0)
|
t3 = time.time()
|
print("t2-t1::", t3-t2, similarity1)
|
