#!/usr/bin/env catnip
# Inférence batchée avec ResNet18 (torchvision, ImageNet)
# Construit un batch à partir de plusieurs variantes d'une même image
# (broadcasting Catnip), exécute une seule passe vectorisée, puis renvoie
# la classe prédite pour chaque variante.
#
# DEPS: torch, torchvision

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

script_dir = Path(META.file).parent

print("⇒ ResNet18 — inférence batchée")

device = 'cpu'
if torch.cuda.is_available() { device = 'cuda' }
print(f"  device : {device}")

weights = tv_models.ResNet18_Weights.DEFAULT
model = tv_models.resnet18(weights=weights).to(device)
model.eval()
labels = weights.meta['categories']
preprocess = weights.transforms()

# Image source normalisée [0,1]
raw_u8 = tv_io.read_image(str(script_dir / "data/catnip_test.jpg"))
if int(raw_u8.shape[0]) > 3 { raw_u8 = raw_u8[:3] }
raw_image = torch.div(raw_u8.to(torch.float32), 255.0)

# Chaque variante est une transformation de l'image d'origine
struct Variant { name: str; tensor; }

variants = list(
    Variant("original", raw_image),
    Variant("flip_h", torch.flip(raw_image, dims=list(2))),
    Variant("flip_v", torch.flip(raw_image, dims=list(1))),
    Variant("bright", torch.clamp(raw_image * 1.3, 0.0, 1.0)),
)

# Broadcasting : preprocessing de chaque variante, puis stack en un seul batch
batch = torch.stack(variants.[(v) => { preprocess(v.tensor) }]).to(device)

print(f"  batch  : {batch.shape}")

torch.set_grad_enabled(False)
probs = torch.softmax(model(batch), dim=1)
preds = torch.argmax(probs, dim=1)

print()
print("⇒ Résultats")
for (i, v) in enumerate(variants) {
    idx = int(preds[i].item())
    score = round(float(probs[i][idx].item()), 4)
    print(f"  {v.name:<10} -> {labels[idx]:<24} ({score})")
}