|
import argparse
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
|
|
from backbones import get_model
|
|
|
|
|
|
@torch.no_grad()
|
|
def inference(weight, name, img):
|
|
if img is None:
|
|
img = np.random.randint(0, 255, size=(112, 112, 3), dtype=np.uint8)
|
|
else:
|
|
img = cv2.imread(img)
|
|
img = cv2.resize(img, (112, 112))
|
|
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
img = np.transpose(img, (2, 0, 1))
|
|
img = torch.from_numpy(img).unsqueeze(0).float()
|
|
img.div_(255).sub_(0.5).div_(0.5)
|
|
net = get_model(name, fp16=False)
|
|
net.load_state_dict(torch.load(weight))
|
|
net.eval()
|
|
feat = net(img).numpy()
|
|
print(feat)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description='PyTorch ArcFace Training')
|
|
parser.add_argument('--network', type=str, default='r50', help='backbone network')
|
|
parser.add_argument('--weight', type=str, default='')
|
|
parser.add_argument('--img', type=str, default=None)
|
|
args = parser.parse_args()
|
|
inference(args.weight, args.network, args.img)
|
|
|
