# infer.py
from PIL import Image
from hf_model import PretrainedAgeEstimator

def predict_age(image_path: str):
    est = PretrainedAgeEstimator()
    img = Image.open(image_path)
    age, top = est.predict(img, topk=3)
    return age, top

if __name__ == "__main__":
    import sys, json
    path = sys.argv[1]
    age, top = predict_age(path)
    print(json.dumps({
        "estimated_age": round(age, 1),
        "top_classes": [(lbl, round(prob, 4)) for lbl, prob in top]
    }, indent=2))