Views
No views yet
1from urllib.request import urlopen
2
3import einops
4import numpy as np
5import onnxruntime as ort
6from PIL import Image
7
8def softmax(x):
9 y = np.exp(x - np.max(x))
10 return y / y.sum(axis=0)
11
12IMG_URL = 'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'
13IN1K_CLASSES_URL = 'https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt'
14
15session = ort.InferenceSession('edgenext_small.usi_in1k.ort')
16# session = ort.InferenceSession('edgenext_small.usi_in1k.onnx')
17
18labels = urlopen(IN1K_CLASSES_URL).read().decode().splitlines()
19img = np.array(
20 Image.open(urlopen(IMG_URL))
21 .resize(session._sess.inputs_meta[0].shape[2:])
22)
23
24# e.g. in1k norm stats
25mean = .485, .456, .406
26sd = .229, .224, .225
27img = (img / 255. - mean) / sd
28
29# to clearly illustrate format ort expects
30img = einops.rearrange(img, 'h w c -> 1 c h w').astype(np.float32)
31out = session.run(None, {session.get_inputs()[0].name: img})
32
33out = softmax(out[0][0])
34topk = np.argsort(out)[::-1][:5]
35for i in topk:
36 print(f'{out[i]:.2f}', labels[i])