"""Is this Simba? ConvNeXt-Tiny features (timm) + logistic-regression head."""
import sys
import joblib, timm, torch
from PIL import Image, ImageOps
from huggingface_hub import hf_hub_download

head = joblib.load(hf_hub_download("makactuary/simba-or-not", "head.joblib"))
backbone = timm.create_model("convnext_tiny.fb_in22k_ft_in1k", pretrained=True, num_classes=0).eval()
transform = timm.data.create_transform(**timm.data.resolve_data_config({}, model=backbone))

def is_simba(path):
    with torch.inference_mode():
        x = backbone(transform(ImageOps.exif_transpose(Image.open(path)).convert("RGB")).unsqueeze(0)).numpy()
    p = float(head.predict_proba(x)[0, 1])
    return ("Simba" if p >= 0.5 else "not Simba"), round(p, 3)

if __name__ == "__main__":
    print(is_simba(sys.argv[1]))