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()