File size: 1,180 Bytes
e362006 | 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 | import numpy as np
# ImageNet class labels (top-5 for demo)
IMAGENET_CLASSES = {
0: 'tench', 1: 'goldfish', 2: 'great white shark', 3: 'tiger shark',
207: 'golden retriever', 248: 'Eskimo dog', 281: 'tabby cat',
282: 'tiger cat', 283: 'Persian cat', 284: 'Siamese cat',
386: 'African elephant', 388: 'giant panda', 402: 'acoustic guitar',
404: 'airliner', 417: 'balloon', 430: 'basketball',
504: 'coffee mug', 530: 'digital clock', 549: 'dumbbell',
582: 'grille', 634: 'carton', 673: 'mouse', 700: 'paper towel',
764: 'skyscraper', 817: 'sports car', 850: 'teddy bear',
954: 'banana', 967: 'espresso', 972: 'cliff', 988: 'daisy'
}
def postprocess(output, top_k=5):
"""Get top-k predictions from model output."""
probs = softmax(output[0])
top_indices = np.argsort(probs)[::-1][:top_k]
results = []
for idx in top_indices:
label = IMAGENET_CLASSES.get(int(idx), f'class_{idx}')
results.append({
'class_id': int(idx),
'label': label,
'confidence': float(probs[idx])
})
return results
def softmax(x):
e_x = np.exp(x - np.max(x))
return e_x / e_x.sum()
|