| """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']}") | |