| import numpy as np |
|
|
| |
| 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() |
|
|
