RabbitRUI/ruispace

1import argparse2 3import cv24import numpy as np5import torch6 7from backbones import get_model8 9 10@torch.no_grad()11def inference(weight, name, img):12    if img is None:13        img = np.random.randint(0, 255, size=(112, 112, 3), dtype=np.uint8)14    else:15        img = cv2.imread(img)16        img = cv2.resize(img, (112, 112))17 18    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)19    img = np.transpose(img, (2, 0, 1))20    img = torch.from_numpy(img).unsqueeze(0).float()21    img.div_(255).sub_(0.5).div_(0.5)22    net = get_model(name, fp16=False)23    net.load_state_dict(torch.load(weight))24    net.eval()25    feat = net(img).numpy()26    print(feat)27 28 29if __name__ == "__main__":30    parser = argparse.ArgumentParser(description='PyTorch ArcFace Training')31    parser.add_argument('--network', type=str, default='r50', help='backbone network')32    parser.add_argument('--weight', type=str, default='')33    parser.add_argument('--img', type=str, default=None)34    args = parser.parse_args()35    inference(args.weight, args.network, args.img)36