Spaces:
Build error
Build error
File size: 1,032 Bytes
83d8d3c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
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)
|