"""Identify a tree from a photo.

    pip install timm torch pillow huggingface_hub
    python example.py photo.jpg
"""
import json
import sys

import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

REPO_ID = "Maximaxou2/french-trees-vit-small"

model = timm.create_model(f"hf-hub:{REPO_ID}", pretrained=True).eval()
transform = timm.data.create_transform(**timm.data.resolve_data_config({}, model=model))
labels = json.load(open(hf_hub_download(REPO_ID, "labels.json"), encoding="utf8"))
# Calibrated temperature: makes the percentages honest (fitted on the validation split).
temperature = json.load(open(hf_hub_download(REPO_ID, "config.json"))).get("temperature", 1.0)

image = Image.open(sys.argv[1]).convert("RGB")
with torch.no_grad():
    probs = (model(transform(image).unsqueeze(0)) / temperature).softmax(-1)[0]

for p, i in zip(*probs.topk(3)):
    label = labels[i]
    print(f"{p:6.1%}  {label['latin']:28} {label['fr']}")