| 9 | |
| 10 | @torch.no_grad() |
| 11 | def 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 | |
| 29 | if __name__ == "__main__": |