#!/usr/bin/env catnip
# Classification d'image avec un ResNet18 pré-entraîné (torchvision, ImageNet)
# Charge le modèle + ses poids, lit une image locale, applique le preprocessing
# ImageNet officiel, puis renvoie le top-5 des classes prédites.
#
# DEPS: torch, torchvision

torch = import('torch')
tv_models = import('torchvision.models')
tv_io = import('torchvision.io')
import('pathlib', 'Path')

# Résolution relative au script (fonctionne depuis n'importe quel CWD)
script_dir = Path(META.file).parent

print("⇒ ResNet18 — classification ImageNet")

# Device : GPU si disponible, sinon CPU
device = 'cpu'
if torch.cuda.is_available() { device = 'cuda' }
print(f"  device : {device}")

# Modèle pré-entraîné ImageNet (1000 classes) + son preprocessing officiel
weights = tv_models.ResNet18_Weights.DEFAULT
model = tv_models.resnet18(weights=weights).to(device)
model.eval()
labels = weights.meta['categories']
preprocess = weights.transforms()

# Chargement de l'image (uint8 CHW) puis normalisation [0,1]
image_path = script_dir / "data/catnip_test.jpg"
raw_u8 = tv_io.read_image(str(image_path))
channels = int(raw_u8.shape[0])

# Adapter les canaux vers RGB (gris -> 3 canaux, RGBA -> RGB)
if channels == 1 { raw_u8 = raw_u8.repeat(tuple(3, 1, 1)) }
if channels > 3 { raw_u8 = raw_u8[:3] }

raw_image = torch.div(raw_u8.to(torch.float32), 255.0)
x = preprocess(raw_image).unsqueeze(0).to(device)

print(f"  image  : {image_path.name} {raw_image.shape}")
print(f"  entrée : {x.shape}")

# Inférence sans gradient
torch.set_grad_enabled(False)
probs = torch.softmax(model(x), dim=1)
topk = torch.topk(probs, k=5, dim=1)

print()
print("⇒ Top-5")
for i in range(5) {
    cls_idx = int(topk.indices[0][i].item())
    score = round(float(topk.values[0][i].item()), 4)
    print(f"  {i + 1}. {labels[cls_idx]:<24} (idx {cls_idx}) = {score}")
}

best_idx = int(topk.indices[0][0].item())
best_score = round(float(topk.values[0][0].item()), 4)
print()
print(f"⇒ Prédiction : {labels[best_idx]} (score {best_score})")